{"id":5648335,"date":"2023-07-02T06:12:52","date_gmt":"2023-07-02T10:12:52","guid":{"rendered":"https:\/\/lightning.ai\/pages\/?p=5648335"},"modified":"2024-02-01T09:50:34","modified_gmt":"2024-02-01T14:50:34","slug":"pytorch-memory-vit-llm","status":"publish","type":"post","link":"https:\/\/lightning.ai\/pages\/community\/tutorial\/pytorch-memory-vit-llm\/","title":{"rendered":"Optimizing Memory Usage for Training LLMs and Vision Transformers in PyTorch"},"content":{"rendered":"

Key takeaway<\/h3>
\nPeak memory consumption is a common bottleneck when training deep learning models such as vision transformers and LLMs. This article provides a series of techniques that can lower memory consumption in PyTorch by approximately 20x without sacrificing modeling performance and prediction accuracy.
\n<\/div>\n

Introduction<\/span><\/h2>\n

In this article, we will be exploring 10 easily-accessible techniques to reduce memory usage in PyTorch. These techniques are cumulative, meaning we can apply them on top of one another. <\/span><\/p>\n

We will begin working with a vision transformer from PyTorch’s <\/span>Torchvision<\/span><\/a><\/span> library to provide simple code examples that you can execute on your own machine without downloading and installing too many code and dataset dependencies. The self-contained baseline training script consists of ~100 lines of code (ignoring the whitespaces and code comments). All code examples are available <\/span>here on GitHub<\/span><\/a><\/span>.<\/span><\/p>\n

Here’s an outline of the sections and techniques we are going to cover (you can access these sections directly from the Table of Contents on the right):<\/span><\/p>\n

    \n
  1. Finetuning a Vision Transformer<\/span><\/li>\n
  2. Automatic Mixed-Precision Training<\/span><\/li>\n
  3. Lower-Precision Training<\/span><\/li>\n
  4. Training with Reduced Batch Size<\/span><\/li>\n
  5. Gradient Accumulation and Microbatches<\/span><\/li>\n
  6. Choosing Leaner Optimizers<\/span><\/li>\n
  7. Instantiating Models on the Target Device<\/span><\/li>\n
  8. Distributed Training and Tensor Sharding<\/span><\/li>\n
  9. Activation Checkpointing<\/li>\n
  10. Parameter Offloading <\/span><\/li>\n
  11. Putting It All Together: Training an LLM<\/span><\/li>\n<\/ol>\n

    While we are working with a vision transformer here (the ViT-L-16 model from the paper <\/span>An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale<\/span><\/a><\/span>), all the techniques used in this article transfer to other models as well: Convolutional networks, large language models (LLMs), and others.<\/span><\/p>\n

    Furthermore, after introducing one technique at a time using the abovementioned vision transformer example, we will apply these to train a BigBird-Roberta LLM on a text classification task. It wouldn’t be possible to train such a model on consumer hardware without these techniques.<\/span><\/p>\n

    PS: Note that there are many sections in this article. To not bloat this article further, I will keep each section purposefully short but provide links to more detailed articles on the individual topics.<\/span><\/p>\n

     <\/p>\n

    1) Finetuning a Vision Transformer<\/span><\/h2>\n

    To simplify the PyTorch code for the experiments, we will be introducing the <\/span>open-source Fabric library<\/span><\/a><\/span>, which allows us to apply various advanced PyTorch techniques (automatic mixed-precision training, multi-GPU training, tensor sharding, etc.) with a handful (instead of dozens) lines of code.<\/span><\/p>\n

    The difference between simple PyTorch code and the modified one to use Fabric is subtle and involves only minor modifications, as highlighted in the code below:<\/span><\/p>\n

    \"\"<\/p>\n

     <\/p>\n

    As mentioned above, these minor changes now provide a gateway to utilize advanced features in PyTorch, as we will see in a bit, without restructuring any more of the existing code.<\/span><\/p>\n

    To summarize the figure above, the main 3 steps for converting plain PyTorch code to PyTorch+Fabric are as follows:<\/span><\/p>\n

    \"\"<\/p>\n

     <\/p>\n

      \n
    1. \n

      Import Fabric and instantiate a Fabric object.<\/span><\/p>\n<\/li>\n

    2. \n

      Use Fabric to set up the model, the optimizer, and the data loader.<\/span><\/p>\n<\/li>\n

    3. \n

      Call <\/span>fabric.backward()<\/code><\/span> on the loss instead of the usual <\/span>loss.backward()<\/code><\/span><\/p>\n<\/li>\n<\/ol>\n

      The vision transformer is based on the <\/span>original ViT architecture<\/span><\/a><\/span>, and the code is available here for inspection. Note that we are finetuning the model for classification instead of training it from scratch to optimize predictive performance.<\/span><\/p>\n

      As a quick sanity check, the predictive performance and memory consumption using plain PyTorch and PyTorch with Fabric remains exactly the same (+\/- expected fluctuations due to randomness):<\/span><\/p>\n

      Plain PyTorch (<\/span>01_pytorch-vit.py<\/span><\/a><\/span>):<\/span><\/strong><\/span><\/p>\n

      Time elapsed 17.94 min<\/span>\r\nMemory used: 26.79 GB<\/span>\r\nTest accuracy 95.85%<\/span><\/pre>\n
      PyTorch with Fabric (<\/span>01-2_pytorch-fabric.py<\/span><\/a><\/span>)<\/span><\/strong><\/span><\/p>\n
      Time elapsed 17.88 min<\/span>\r\nMemory used: 26.84 GB<\/span>\r\nTest accuracy 96.06%<\/span><\/pre>\n
      As an optional exercise, you are welcome to experiment with the code and replace<\/span><\/p>\n
      model<\/span> =<\/span> vit_l_16<\/span>(weights<\/span>=<\/span>ViT_L_16_Weights<\/span>.IMAGENET1K_V1<\/span>)<\/span><\/pre>\n
      with<\/span><\/p>\n
      model<\/span> =<\/span> vit_l_16<\/span>(weights<\/span>=<\/span>None<\/span>)<\/span><\/pre>\n
      This will train the same vision transformer architecture from scratch instead of finetuning it. If you carry out this exercise, you’ll see that the prediction accuracy drops from >96% down to ~60%:<\/span><\/p>\n
      \"\"<\/p>\n
       <\/p>\n
      2) Automatic Mixed-Precision <\/span><\/h2>\n
      In the previous section, we modified our PyTorch code using Fabric. Why go through all this hassle? As we will see below, we can now try advanced techniques, like mixed-precision and distributed training, by only changing one line of code.<\/span><\/p>\n
      We will start with mixed-precision training, which has become the recent norm for training deep neural networks. <\/span><\/p>\n
      Applying Mixed-Precision Training<\/span><\/strong><\/span><\/p>\n
      We can apply mixed-precision training with only one small modification, changing <\/span><\/p>\n
      fabric<\/span> =<\/span> Fabric<\/span>(accelerator<\/span>=<\/span>\"cuda\"<\/span>, devices<\/span>=<\/span>1<\/span>)<\/span><\/pre>\n
      to the following:<\/span><\/p>\n
      fabric<\/span> =<\/span> Fabric<\/span>(accelerator<\/span>=<\/span>\"cuda\"<\/span>, devices<\/span>=<\/span>1<\/span>, precision<\/span>=<\/span>\"16-mixed\"<\/span>)<\/span><\/pre>\n
      As a result, our memory consumption is reduced from 26.84 GB to 18.21 GB without sacrificing prediction accuracy, as shown below.<\/span><\/p>\n
       <\/p>\n
      \"a\"
      Comparing <\/span>01-2_pytorch-fabric.py<\/span><\/a><\/span> and <\/span>02_mixed-precision.py<\/span><\/a><\/span><\/p><\/div>\n
      As a bonus, mixed-precision training doesn’t only reduce memory usage but also reduces the runtime 6-fold (from 17.88 min to 3.45 min), which is a nice, added benefit; however, the focus of this particular article is on memory consumption to not complicate it further.<\/span><\/p>\n
      What Is Mixed-Precision Training?<\/span><\/strong><\/span><\/p>\n
      Mixed precision training uses both 16-bit and 32-bit precision to ensure no loss in accuracy. The computation of gradients in the 16-bit representation is much faster than in the 32-bit format and saves a significant amount of memory. This strategy is beneficial, especially when we are memory or compute-constrained.<\/span><\/p>\n
      It’s called “mixed-“rather than “low-“precision training because we don’t transfer <\/span>all<\/span><\/em><\/span> parameters and operations to 16-bit floats. Instead, we switch between 32-bit and 16-bit operations during training, hence, the term “mixed” precision.<\/span><\/p>\n
      As illustrated in the figure below, mixed-precision training involves converting weights to lower-precision (FP16) for faster computation, calculating gradients, converting gradients back to higher-precision (FP32) for numerical stability, and updating the original weights with the scaled gradients.<\/span><\/p>\n
      This approach allows for efficient training while maintaining the accuracy and stability of the neural network. <\/span><\/p>\n
       <\/p>\n
      \"\"<\/p>\n
      For additional details, I recommend reading my more detailed standalone article <\/span>Accelerating Large Language Models with Mixed-Precision Techniques<\/span><\/a><\/span>, where I dive deeper into the underlying concepts.<\/span><\/p>\n
      3) Lower-Precision Training<\/span><\/h2>\n
      We can also take it a step further and try running with “full” lower 16-bit precision (instead of mixed-precision, which converts intermediate results to a 32-bit representation.)<\/span><\/p>\n
      We can enable lower-precision training by changing<\/span><\/p>\n
      fabric<\/span> =<\/span> Fabric<\/span>(accelerator<\/span>=<\/span>\"cuda\"<\/span>, precision<\/span>=<\/span>\"16-mixed\"<\/span>)<\/span><\/pre>\n
      to the following: <\/span><\/p>\n
      fabric<\/span> =<\/span> Fabric<\/span>(accelerator<\/span>=<\/span>\"cuda\"<\/span>, precision<\/span>=<\/span>\"16-true\"<\/span>)<\/span><\/pre>\n
      However, you may notice that when running this code, you’ll encounter NaN values in the loss:<\/span><\/p>\n
      Epoch: 0001\/0001 | Batch 0000\/0703 | Loss: 2.4105<\/span>\r\nEpoch: 0001\/0001 | Batch 0300\/0703 | Loss: nan<\/span>\r\nEpoch: 0001\/0001 | Batch 0600\/0703 | Loss: nan<\/span>\r\n...<\/span><\/pre>\n
      This is because regular 16-bit floats can only represent numbers between -65,504 and 65,504:<\/span><\/p>\n
      In<\/span> [1<\/span>]: import<\/span> torch<\/span><\/span>\r\n\u200b<\/span>\r\nIn<\/span> [2<\/span>]: torch<\/span>.finfo<\/span>(torch<\/span>.float16<\/span>)<\/span>\r\nOut<\/span>[2<\/span>]: finfo<\/span>(resolution<\/span>=<\/span>0.001<\/span>, min<\/span>=-<\/span>65504<\/span>, max<\/span>=<\/span>65504<\/span>, eps<\/span>=<\/span>0.000976562<\/span>, smallest_normal<\/span>=<\/span>6.10352e-05<\/span>, tiny<\/span>=<\/span>6.10352e-05<\/span>, dtype<\/span>=<\/span>float16<\/span>)<\/span><\/pre>\n
      So, to avoid the NaN issue, we can use the “<\/span>bf16-true<\/code><\/span>” setting.<\/span><\/p>\n
      fabric<\/span> =<\/span> Fabric<\/span>(accelerator<\/span>=<\/span>\"cuda\"<\/span>, precision<\/span>=<\/span>\"bf16-true\"<\/span>)<\/span><\/pre>\n
      As a result, we can reduce the memory consumption even further down to 13.82 GB (again, without sacrificing accuracy):<\/span><\/p>\n
      \"a\"
      Comparing <\/span>03_bfloat16.py<\/span><\/a><\/span> to the previous codes<\/span><\/p><\/div>\n
      What Is Bfloat16?<\/span><\/strong><\/span><\/p>\n
      The “bf16” in <\/span>\"bf16-mixed\"<\/code><\/span> stands for <\/span>Brain Floating Point<\/span><\/a><\/span> (bfloat16). Google developed this format for machine learning and deep learning applications, particularly in their Tensor Processing Units (TPUs). Bfloat16 extends the dynamic range compared to the conventional float16 format at the expense of decreased precision.<\/span><\/p>\n
      \"\"<\/p>\n
       <\/p>\n
      The extended dynamic range helps bfloat16 to represent very large and very small numbers, making it more suitable for deep learning applications where a wide range of values might be encountered. However, the lower precision may affect the accuracy of certain calculations or lead to rounding errors in some cases. But in most deep learning applications, this reduced precision has minimal impact on modeling performance.<\/span><\/p>\n
      While bfloat16 was originally developed for TPUs, this format is now supported by several NVIDIA GPUs as well, beginning with the A100 Tensor Core GPUs, which are part of the NVIDIA Ampere architecture.<\/span><\/p>\n
      You can check whether your GPU supports <\/span>bfloat16<\/code><\/span> via the following code:<\/span><\/p>\n
      >>><\/span> import<\/span> torch<\/span><\/span>\r\n>>><\/span> torch<\/span>.cuda<\/span>.is_bf16_supported<\/span>()<\/span>\r\nTrue<\/span><\/span><\/pre>\n
      4) Reducing the Batchsize<\/span><\/h2>\n
      Let’s tackle one of the big elephants in the room: why don’t we simply reduce the batch size? This is usually always an option to reduce memory consumption. However, it can sometimes result in worse predictive performance since it alters the training dynamics. (For more details, <\/span>see Lecture 9.5 in my Deep Learning Fundamentals course<\/span><\/a><\/span>.) <\/span><\/p>\n
      Either way, let’s reduce the batch size to see how that affects the results. It turns out we can lower the batch size to 16, which brings memory consumption down to 5.69 GB, without sacrificing performance:<\/span><\/p>\n
      \"a\"
      Comparing <\/span>04_lower-batchsize.py<\/span><\/a><\/span> to the previous codes.<\/span><\/p><\/div>\n
       <\/p>\n
      5) Using Gradient Accumulation to Create Microbatches<\/span><\/h2>\n
      Gradient accumulation is a way to virtually increase the batch size during training, which is very useful when the available GPU memory is insufficient to accommodate the desired batch size. Note that this only affects the runtime, not the modeling performance.<\/span><\/p>\n
      In gradient accumulation, gradients are computed for smaller batches and accumulated (usually summed or averaged) over multiple iterations instead of updating the model weights after every batch. Once the accumulated gradients reach the target \u201cvirtual\u201d batch size, the model weights are updated with the accumulated gradients.<\/span><\/p>\n
      To enable gradient accumulation, there are only two small modifications to the forward and backward pass required:<\/span><\/p>\n
      \"\"
      Code modification in <\/span>05_gradient-accum.py<\/span><\/a><\/span><\/p><\/div>\n
       <\/p>\n
      I covered gradient accumulation in more detail in my article <\/span>Finetuning LLMs on a Single GPU Using Gradient Accumulation<\/span><\/a><\/span>. <\/span><\/p>\n
      Using an effective batch size of 16 and 4 accumulation steps means we will use an actual batch size of 4 (since 16 \/ 4 = 4).<\/span><\/p>\n
      \"a\"
      Result of <\/span>05_gradient-accum.py<\/span><\/a><\/span><\/p><\/div>\n
      A disadvantage of this technique is that it increases the runtime from 3.96 min to 12.91 min.<\/span><\/p>\n
      Of course, we could even go smaller and e 16 accumulation steps. This would lead to a microbatch size of 1, reducing the memory size further (about 75%), but I’ll leave this as an optional exercise.<\/span><\/p>\n
      6) Using a Leaner Optimizer<\/span><\/h2>\n
      Did you know that the popular Adam optimizer comes with additional parameters? For instance, Adam has 2 additional optimizer parameters (a mean and a variance) for each model parameter. <\/span><\/p>\n
      So, by swapping Adam with a stateless optimizer like SGD, we can reduce the number of parameters by 2\/3, which can be quite significant when working with vision transformers and LLMs.<\/span><\/p>\n
      The downside of plain SGD is that it usually has worse convergence properties. So, let’s swap Adam with SGD and introduce a cosine decay learning rate scheduler to compensate for this and achieve better convergence.<\/span><\/p>\n
      In short, we will be swapping the previously used Adam optimizer:<\/span><\/p>\n
      optimizer<\/span> =<\/span> torch<\/span>.optim<\/span>.Adam<\/span>(model<\/span>.parameters<\/span>(), lr<\/span>=<\/span>5e-5<\/span>)<\/span><\/pre>\n
      with an SGD optimizer plus scheduler:<\/span><\/p>\n
      optimizer<\/span> =<\/span> torch<\/span>.optim<\/span>.SGD<\/span>(model<\/span>.parameters<\/span>(), lr<\/span>=<\/span>0.01<\/span>)<\/span>\r\n\u200b<\/span>\r\nnum_steps<\/span> =<\/span> NUM_EPOCHS<\/span> *<\/span> len<\/span>(train_loader<\/span>)<\/span>\r\nscheduler<\/span> =<\/span> torch<\/span>.optim<\/span>.lr_scheduler<\/span>.CosineAnnealingLR<\/span>(<\/span>\r\n \u00a0 \u00a0optimizer<\/span>, T_max<\/span>=<\/span>num_steps<\/span>)<\/span><\/pre>\n
      With this change, we are able to have the peak memory consumption while maintaining ~97% classification accuracy:<\/span><\/p>\n
      \"a\"
      Result of <\/span>06_sgd-with-scheduler.py<\/span><\/a><\/span><\/p><\/div>\n
      If you want to learn more, I discussed learning rate schedulers (including cosine decay with a 1-cycle schedule) in more detail in my <\/span>Unit 6.5 of my Deep Learning Fundamentals class<\/span><\/a><\/span>.<\/span><\/p>\n
      7) Creating the Model on the Target Device with Desired Precision<\/span><\/h2>\n
      When we instantiate a model in PyTorch, we usually create it on the CPU device first, and then we transfer it onto the target device and convert it to the desired precision:<\/span><\/p>\n
      model<\/span> =<\/span> vit_l_16<\/span>(weights<\/span>=<\/span>ViT_L_16_Weights<\/span>.IMAGENET1K_V1<\/span>)<\/span>\r\nmodel<\/span>.cuda<\/span>().float16<\/span>()<\/span><\/pre>\n
      This can be inefficient considering the intermediate model representation in full precision on the CPU. Instead, we can directly create the model in desired precision on the target device (e.g., GPU) using the <\/span>init_module<\/code><\/span> context in Fabric:<\/span><\/p>\n
      import<\/span> lightning<\/span> as<\/span> L<\/span><\/span>\r\n\u200b<\/span>\r\nfabric<\/span> =<\/span> Fabric<\/span>(accelerator<\/span>=<\/span>\"cuda\"<\/span>, devices<\/span>=<\/span>1<\/span>, precision<\/span>=<\/span>\"16-true\"<\/span>)<\/span>\r\n\u200b<\/span>\r\nwith<\/span> fabric<\/span>.init_module<\/span>():<\/span>\r\n \u00a0 \u00a0model<\/span> =<\/span> vit_l_16<\/span>(weights<\/span>=<\/span>ViT_L_16_Weights<\/span>.IMAGENET1K_V1<\/span>)<\/span><\/pre>\n
      In this specific case (model), the peak memory during the forward pass is larger than the model size in its full precision representation. So, we will benchmark the <\/span>fabric.init_module<\/code><\/span> approach just for the model loading itself.<\/span><\/p>\n