Lowering PyTorch's Memory Consumption for Selective Differentiation

Authors: samarth  bhatia
felix  dangel
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-02 17:29:23

Abstract: Memory is a limiting resource for many deep learning tasks.
Beside the neural network weights, one main memory consumer is the computation graph built up by automatic differentiation (AD) for backpropagation.
We observe that PyTorch's current AD implementation neglects information about parameter differentiability when storing the computation graph.
This information is useful though to reduce memory whenever gradients are requested for a parameter subset, as is the case in many modern fine-tuning tasks.
Specifically, inputs to layers that act linearly in their parameters and inputs (dense, convolution, or normalization layers in evaluation mode) can be discarded whenever the parameters are marked as non-differentiable. 
We provide a drop-in, differentiability-agnostic implementation of such layers and demonstrate its ability to reduce memory without affecting run time on popular convolution- and attention-based architectures.

Main paper: 
 introduction & motivation: 
 The success of many deep learning applications is driven by scaling computational resources (Thompson et al., 2020) . One important resource is GPU memory, specifically on low-and mid-end GPUs which usually offer between 6 to 16 GiB. Therefore, down-scaling the computational demands of deep learning is an important objective to widen its accessibility to researchers and practitioners with fewer hardware resources. Two major memory consumers are the network weights, and the computation graph stored by the automatic differentiation (AD) engine. There exist various Figure 1 : PyTorch's AD is sometimes not agnostic to parameter differentiability. We consider a deep CNN made of size-preserving convolutions and measure the forward pass's peak memory when processing a mini-batch of size (256, 8, 256, 256) , requiring 512 MiB memory. Memory increases linearly in the number of layers when all parameters are marked differentiable and remains constant when all parameters are marked non-differentiable. Surprisingly, when only one layer's parameters are marked as differentiable the memory increases as if all subsequent parameters were marked differentiable. Our drop-in solution stores layer inputs depending on parameter differentiability and reduces memory compared to the current PyTorch implementation. approaches to reduce their memory burden; e.g., the parameters can be compressed with low-precision data types (quantization (Hubara et al., 2018; Li et al., 2017; Nagel et al., 2021) ) or sparsified (Hassibi & Stork, 1992; Frantar & Alistarh, 2022) , and the computation graph can be off-loaded to CPU (Ren et al., 2021) , compressed (Chen et al., 2021) , randomized (Oktay et al., 2021) , or partially recorded and re-computed (gradient checkpointing (Griewank & Walther, 2008; Chen et al., 2016) ).Here, we consider a special AD scenario we call selective differentiation where gradients are requested only for a subset of the computation graph leafs (the neural network inputs and parameters). This approach has gained a lot of practical relevance in the era of large foundation models. For instance, fine-tuning techniques for pre-trained models rely on training only a subset of layers (Zhao et al., 2024; Lee et al., 2023; Zhu et al., 2024; Pan et al., 2024) , or extend them with parameter-efficient adapters (Hu et al., 2022) which are then trained, keeping the original weights fixed. But selective differentiation is also common in other applications such as generating adversarial examples (Goodfellow et al., 2015) and neural style transfer (Gatys et al., 2016) which optimize the input to a network with frozen weights.We find that PyTorch's (Paszke et al., 2019) AD allows for additional, simple, optimizations to further reduce memory consumption in the context of selective differentiation:1. We observe that PyTorch's AD sometimes neglects the differentiability of layer parameters when storing the computation graph (Figure 1 ). This information is useful though as it allows discarding inputs to linear layers whose parameters are marked as non-differentiable.2. We provide a drop-in implementation of various layers that is agnostic to the parameter differentiability and demonstrate on various convolutional neural networks (CNNs) and attention-based large language models (LLMs) that it lowers the default implementation's memory footprint without increasing run time.This easy-to-use insight benefits many tasks with selective differentiation and enables them to scale further. We hope it will stimulate future research into AD optimizations. 
selective differentiation in pytorch: 
 PyTorch users can specify whether they want to compute gradients w.r.t. a tensor through its requires grad attribute, which is dominantly inherited by child tensors: If any input (parent) to an operation is marked differentiable, the output (child) will also be differentiable and require the computation graph to be stored for backpropagation. As we will see, PyTorch does not always check the differentiability of all parents, but rather stores the computation graph as if all parents were differentiable once it encounters a differentiable input. In the following example, we demonstrate this behavior, and how it misses out on possible optimizations for linear operations with non-differentiable parameters.Experimental procedure: We use PyTorch 2.2.1 and measure the forward pass's peak memory of different neural networks as a proxy for the computation graph size stored by the AD engine. We choose peak memory as it is a reasonable proxy for the computation graph size and the relevant metric for causing out-of-memory errors in practise. On GPU, we measure peak memory using torch.cuda.max memory allocated() , the maximum memory allocated by CUDA. To assess run time differences of our implementation, we compare a forward and backward pass with PyTorch. Each measurement is (Zagoruyko et al., 2019) .performed in a separate Python session to avoid memory (de)allocation leaks between consecutive runs.We start with a synthetic example to probe the internals of PyTorch's AD that are responsible for identifying and storing the computation graph during a forward pass (summarized in Figure 2 ). We consider a deep CNN consisting entirely of convolutions without bias. Each convolution preserves its input size (kernel size 3, unit padding and stride) and number of channels and we vary its depth as well as the parameter differentiability. As input, we choose a mini-batch of size (256, 8, 256, 256) . Storing this tensor, and each intermediate output generated by a layer, requires 512 MiB in single precision. Memory consumed by the convolution kernels of shape (8,8,3,3) is negligible compared to the hidden features. Figure 1 summarizes our findings.First, we investigate the computation graph size when marking all or no parameters as differentiable. When all parameters are differentiable, all layer inputs must be stored to compute gradients. Consequently, we observe a linear relation with a slope corresponding to the 512 MiB consumed by each intermediate input (Figure 1 , fully differentiable).When no parameter is differentiable, the memory consumption flattens after more than two layers (Figure 1 , fully nondifferentiable). This is because during a forward pass both input and output tensors of a layer are allocated in memory, in addition to the network's input. Hence, at most two tensors are in memory at a time for a single layer, while at most three are allocated for two or more layers.Next, we observe the memory consumption in the context of selective differentiation. We consider two scenarios: In the first, all parameters after and including the fourth layer are differentiable, hence all layer inputs after the third layer must be stored in memory. In the second, only the fourth layer is differentiable and therefore only that layer's input is necessary to compute gradients. However, we observe that both scenarios exhibit the same memory footprint (Figure 1 , layers 4+ differentiable and layer 4 differentiable)!We conclude that PyTorch stores a convolution layer's input with requires grad = True although the layer's parameters might be non-differentiable. This information can be useful to reduce the information stored for backpropagation, as we show with our implementation (Figure 1 , layer 4 differentiable (ours)).In Appendix B, we repeat similar experiments with fullyconnected, transpose convolution, and batch normalization layers. We find that PyTorch's (transpose) convolutions and batch normalization layers in evaluation mode are not agnostic to the differentiability of their weights. Interestingly, fully-connected layers are already optimized.Our implementation of memory-saving layers is straightforward and does not require low-level code. For each layer we create a new torch.autograd.Function AD primitive, and its associated torch.nn.Module layer. The primitive uses the same forward and backward routines as the original operation (from torch.nn.functional and torch.ops.aten ), but has additional logic for the information that is saved to the AD tape, which we describe below. Hence, our implementation shares the performance of PyTorch's. We also provide a converter that replaces supported layers of a net with our memory-saving equivalents.Consider a convolution layer Z = W * X + b with input X, output Z, kernel W , and bias b. Its input Jacobian depends on W , the weight Jacobian on X, and the bias Jacobian has no dependency (e.g. Dangel, 2023, Chapter 2.3). During the forward pass, we check the differentiability of W , X, and only store tensors required by the Jacobians that will be applied during backpropagation. The same dependency pattern holds for other layers that process their inputs linearly w.r.t. their weight and input, and add a bias term, such as (transpose) convolution (Chellapilla et al., 2006) and batch normalization (Ioffe & Szegedy, 2015) layers in evaluation mode. We implement them in exactly the same fashion.Real neural nets contain additional layers that are interleaved with linear layers: activations, dropout, pooling, etc. The information stored by such layers may overlap with the input to a linear layer, e.g. if an activation layer feeds into a convolution. Depending on the implementation details, our described optimizations might not apply in such cases because the linear layer's input tensor could still be stored by the preceding activation. We encountered this effect for ReLU layers in real-world CNNs (Section 3), but were able to overcome it using a customized ReLU implementation that saves a boolean mask of the output. As PyTorch currently only supports 1-byte booleans, this leads to a 4x reduction of the stored tensor's size, which could be further reduced to 32x (1-bit booleans are in the works). For dropout layers, we accept adding a slight computational overhead to avoid storing the output by saving the random state and re-generating the mask during backpropagation.This underlines an important challenge for saving memory in selective differentiation whenever multiple layers use the same tensor for backpropagation. Inputs to linear layers, even if their parameters are non-differentiable, might still be stored by neighboring layers. We believe it should often be possible to resolve this scenario-albeit through careful implementation-e.g. whenever the Jacobian can be implemented by either storing the layer input or output, as is the case for various activation functions, and pooling layers. 
real-world examples: 
 Here we measure the effect of our memory-saving layers on real-world CNNs. We consider four different scenarios:All: All network parameters are differentiable. This serves as reference to establish similar performance of our layers in the absence of selective differentiation.Input: Only the input to the neural net is differentiable. This situation resembles constructing adversarial examples (Goodfellow et al., 2015) or style-transfers (Gatys et al., 2016) by optimizing noisy inputs.Surgical: Only the first quarter of layers are differentiable. This situation is similar to surgical fine-tuning (Lee et al., 2023) , which splits a network into different blocks, each containing a subset of layers, to be trained one at a time.Norm: Only normalization layers are differentiable, resembling layer norm fine-tuning in LLMs (Zhao et al., 2024) .We report GPU results on an NVIDIA RTX A6000 with 48 GiB of VRAM. All CNNs are fed inputs of size (64, 3, 224, 224) and we measure according to the procedure described in Section 2. For object detection models, the batch size is 4 and 2 boxes are predicted per input image.For ResNet-101, we take a detailed look at our memorysaving layer's effects. It is a decently powerful, modern model and frequently used as a backbone for other architectures such as CLIP (Radford et al., 2021) and LAVA (Gurram et al., 2022) . It contains dense, convolution, batch normalization, ReLU, and max/average pooling layers. Table 1 summarizes our findings (model in training mode).PyTorch is often unaware of selective differentiation:In the upper part of Table 1 , we see that the default PyTorch implementation uses the same amount of memory for any scenario. In fact, marking only the input to a neural net as differentiable consumes as much memory as marking all parameters, although the former does not require storing the inputs to fully-connected and convolution layers. This confirms the findings on the toy model from Section 2.Layer interactions diminish memory savings: To gradually investigate the effect of our layers, we only swap out the convolutions; without observing any effects. At first, this seems counter-intuitive. However, a closer look at the architecture (Table 2 ) reveals that the convolutions are preceded by ReLU activations which unconditionally store their outputs and render our convolution layer's optimizations ineffective as described in Section 2.2.2.After swapping out ReLU with our mask-based custom implementation, we observe substantial memory savings. E.g., memory consumption for 'Input' dropped to roughly two thirds. Only when all parameters are differentiable, we see a slight increase in memory after swapping out ReLUs. This is to be expected as the mask stored by our custom implementation requires additional storage.The bottom half of Table 1 shows that run time for a fixed case remains equal up to measurement noise. This confirms that our memory-saving layers share the default implementation's performance. To further solidify our findings, we now evaluate our layers on other popular and commonly used CNNs, including ResNet-18 (He et al., 2016) , VGG-16 (Simonyan & Zisserman, 2015) , ResNeXt101-64x4d (Xie et al., 2017) , EfficientNetv2-L (Tan & Le, 2019; 2021) , MobileNetv3-L (Sandler et al., 2018; Howard et al., 2019) , FCN with a ResNet-101 backbone (Long et al., 2015) , DeepLabv3 with a ResNet-101 backbone (Chen et al., 2017) , Faster-RCNN with a ResNet-50 backbone (Ren et al., 2015) , and SSDLite with a MobileNetv3-L backbone (Sandler et al., 2018) . Table 3 summarizes the comparison.On this large repertoire of networks, we observe the same effects as on ResNet-101 from the previous section: Due to our customized ReLU implementation, memory is slightly higher in the absence of selective differentiation. Run time is unaffected by swapping in our layers, while the selective differentiation scenarios consistently show lower memory consumption (as low as two thirds), underlining the usefulness of our approach in this context.So far, we used all batch normalization layers in training mode. In Appendix A, we experiment with BatchNorm2d in evaluation mode. This allows discarding normalization layer inputs whenever the parameters are marked nondifferentiable and enables further memory savings, with reductions up to 6x and no runtime overhead (Table 5 ).Selective differentiation is of high practical relevance for large language models (LLMs), especially with the increasing popularity of vision language models (VLMs) and the release of models such as LLaVa, PaLI, PaLI-Gemma and GPT-4o (Liu et al., 2023; Chen et al., 2023b; a) . Another area where this is quite important is using a (light) modalityspecific encoder to encode (and project) the input of different modalities into the embedding space of the LLM, such as in scientific ML (Shen et al., 2024) .Here, we investigate the impact of our memory-saving layers on attention-based models. Popular libraries like Hugging-Face that provide access to such models often implement attention through linear layers, which we did not find to suffer from the behavior of convolutions. However, an additional challenge in these architectures is dropout, which may interact with inputs to linear layers similarly to other activations (Section 2.2.2). Thus, we find our tricks for ReLU and dropout helpful to avoid storing inputs to linear layers if their weights are marked non-differentiable.We consider the same cases as in Section 3 and all nets consume embeddings of size (64, 256, H), with hidden size H of a network (given in Table 4 ). We test on the vanilla Transformer (Vaswani et al., 2017) , T5 (Raffel et al., 2020) , LLaMa3 (Touvron et al., 2023) , Mistral (Jiang et al., 2023) and Phi3 (Abdin et al., 2024; Gunasekar et al., 2023) . The results are shown in Table 4 . With LLaMa (8B parameters), Mistral (7B parameters) and Phi3 (4B parameters) being large models, they are loaded in bfloat16 data type and the batch size is reduced to 8 for LLaMa and 16 for the others. We observe that our approach also enables memory savings on attention-based models and consistently achieves a smaller memory footprint. 
conclusion: 
 We have shown how to improve the memory consumption of PyTorch's automatic differentiation in the context of selective differentiation where gradients are only requested for a subset of variables-a common situation in modern fine-tuning tasks. Our approach is based on the insight that PyTorch sometimes ignores the differentiability of tensors; specifically in (transpose) convolutions and batch normalization layers in evaluation mode. To overcome this, we provide a drop-in implementation which takes into account the differentiability of all tensors when storing the computation graph. Empirically, we demonstrated the effectiveness of our approach to reduce memory in multiple selective differentiation cases without affecting run time, on both convolution-and attention-based architectures. Our method is easy to use, requiring only a single call to a converter function that replaces all supported layers with our equivalents.Next, we plan to study its impact on real-world applications, such as parameter-efficient LLM fine-tuning with low-rank adapters (Hu et al., 2022) .  
 
 Appendix: Our experiments with a deep CNN of size-preserving convolutions from Figure 1 and Section 2.1 revealed that Py-Torch's 2d convolution layer stores its input tensor whenever it is differentiable, irrespective of the weight's differentiability. Here, we perform analogous experiments, but with other layers we suspect to exhibit similar behaviour. Specifically, we consider layers whose forward pass is linear w.r.t. both the layer input and the layer weights. This includes fully-connected layers ( torch.nn.Linear ), convolution layers ( torch.nn.ConvNd ), transpose convolution layers ( torch.nn.ConvTransposeNd ), and batch normalization layers in evaluation mode (e.g. torch.nn.BatchNorm2d ).Layers are set up to preserve their input size and we feed mini-batches that require 512 MiB storage in single precision:• (512, 1024, 256) for torch.nn.Linear• (4096, 8, 4096) for torch.nn.Conv1d and torch.nn.ConvTranspose1d• (512, 8, 256, 256) for torch.nn.BatchNorm2d , torch.nn.Conv2d , and torch.nn.ConvTranspose2d• (64, 8, 64, 64, 64) for torch.nn.Conv3d and torch.nn.ConvTranspose3dCompared to the input and intermediate activations, the memory footprint of a layer's weight is negligible:• (256, 256) for torch.nn.Linear• (8, 8, 3) for torch.nn.Conv1d and torch.nn.ConvTranspose1d• (8) for torch.nn.BatchNorm2d• (8, 8, 3, 3) for torch.nn.Conv2d and torch.nn.ConvTranspose2d• (8, 8, 3, 3, 3) for torch.nn.Conv3d and torch.nn.ConvTranspose3dOur results are summarized in Figure 3 . We can see that PyTorch's convolutions, transpose convolutions, and batch normalization (in evaluation mode) are not agnostic to the differentiability of their weights. PyTorch's fully-connected layer, however, is already agnostic to the differentiability of its weights. nn.BatchNorm2d in evaluation mode are not agnostic to parameter differentiability. 


Decision: ## Strengths
* This paper describes a very clear idea, that does allow to obtain better memory consumption
* Experiments are convincing
## Weaknesses
* It is not clear how general the results of the paper can be. 
* The experiments are very preliminary and require more generalization and automation
* Some of the insight could be found by reading the code rather than experimenting

The general sentiment about this paper is that the idea is interesting, despite the results being rather preliminary. I recommend
acceptance as a poster.