If you want to try your hand at fine-tuning an LLM (Large Language Model): one of the first things you’re going to need to know is “will it fit on my GPU”.

The pre-eminent guide to estimating (VRAM) memory requirements is Transformer Math 101. It bears mentioning, though, that its heuristics are written in the context of frameworks such as GPT-NeoX and Megatron-DeepSpeed. These are specialist frameworks for training at scale, and are not what a beginner would reach for.

Are Transformer Math’s memory estimates applicable to idiomatic PyTorch, without additional frameworks? We put it to the test.

In brief, we found:

• Memory costs measured on small models extrapolate well to larger models
• PyTorch AMP costs are similar to the frameworks discussed in Transformer Math, but (see next)
• DeepSpeed ZeRO’s fused optimiser saves 2 bytes/param on mixed-precision over PyTorch AMP
• Gradient accumulation can be ~free, but costs more if used in concert with AMP or DDP
• AMP usually costs ~2 bytes/param. Cost increases when gradient accumulation is enabled, or becomes ~free if used in concert with DDP
• DDP usually costs ~4 bytes/param, but becomes cheaper if used in concert with AMP
• DDP can be made 2.5 bytes/param cheaper by enabling gradient_as_bucket_view
• Only the smallest LLMs can be trained under DDP with realistic settings, even on 80GiB GPUs.
• Training large LLMs requires more advanced distributed training techniques, to shard memory costs over multiple GPUs
• Parameter-efficient finetuning is a more realistic option for the GPU-poor

Read on to see our reasoning and evidence.

to reveal more detail. We hide deep-dives in these collapsible sections to keep the article snappy.

Experimental setup

We built a simple training loop (mem_llm.py). It loads a model via HF (HuggingFace) transformers, and trains said model to complete random sequences. By running this training loop for a few steps, then taking a measurement: we can determine our peak memory usage.

We will evaluate the cost of three common techniques provided by an ML accelerator (explanations to follow later):

• Mixed-precision computation
• DDP (Distributed Data Parallel) parallelism
• Gradient accumulation

In our trainer: this functionality is built using PyTorch idioms. Our implementation is similar to how HF accelerate operates in its most basic configurations (i.e. when it is not configured to delegate to other accelerators such as DeepSpeed or Megatron).

What we expect to see

Before we dive into measurements: let’s estimate how much memory we’ll need in theory. We will use reasoning similar to Transformer Math.

For non-distributed training,
Total Training Memory = Model + Optimiser + Activations + Gradients

Our experimental setup eliminates optimiser costs, through means of a stateless optimiser (SGD without momentum).

These measures narrow down the possible sources of overhead, if detected.

Estimated baseline costs in single-precision

First up, let’s establish the terminology we’re going to use:

The memory costs for pure-float32 training are relatively straightforward:

For performance reasons, and to reduce the cost of activations: it is standard to train in mixed-precision.

Estimated baseline costs in mixed-precision

Mixed-precision costs are framework-specific. We will interpret the costs for the frameworks considered in Transformer Math, and also make an educated guess at the costs of PyTorch’s built-in mixed-precision implementation, AMP.

Transformer Math was validated against GPT-NeoX (circa April 2023), and was cross-checked against Megatron-DeepSpeed (against which its predictions held, for everything except activation checkpointing).
Its estimates are expected to be predictive to within ±10%, including on other accelerators such as Megatron-LM, PyTorch FSDP, and ColossalAI.
Thanks go to Quentin Anthony at EleutherAI for confirming the above.

GPT-NeoX implements mixed-precision via DeepSpeed ZeRO. In both DeepSpeed ZeRO and Megatron-DS: the model is held in half-precision, whilst the optimiser takes a float32 copy of the parameters.

PyTorch AMP (Automated Mixed Precision) implements mixed-precision moreorless the opposite way around. The model is held in float32, and AMP manages a half-precision copy.

Our understanding is as follows:

The above mixed-precision costs assume that other features such as gradient accumulation and distributed training are disabled. We will now consider the costs of those features in isolation.

Estimated cost of gradient accumulation

Gradient accumulation is a technique which computes the gradients for a minibatch, one microbatch at a time, summing each result (divided by the number of microbatches). This is typically done to save memory (at the expense of wall time), for a given minibatch size.
In distributed training, gradient accumulation can also be exploited to reduce the frequency with which synchronization of gradients is required, reducing the proportion of training time spent on communication overhead.

When running in mixed-precision: we expect gradient accumulation to cost:

Estimated cost of distributed training

There are many ways to distribute training. Among the simplest is DDP (Distributed Data Parallel), which replicates the model across multiple GPUs. We wish to measure its costs here, because its ease-of-use makes it a good fit for hobbyist or small-scale training.

DDP doesn’t reduce the size of the model or optimiser states. It enables us to achieve bigger minibatch sizes, or give smaller microbatches to each GPU, reducing the per-GPU cost of activations. The experiment we’re doing here has already minimized the cost of activations. Hence we don’t expect DDP to save us memory, but we are interested in measuring the overhead incurred by the extra communication it necessitates.

We expect PyTorch DDP to cost 4 bytes/param, as its distributed Reducer creates gradient buckets for each parameter. This cost can be reduced by enabling gradient_as_bucket_view.

The frameworks studied in Transformer Math can be distributed in more exotic ways than DDP, capable of sharding the model and optimiser states. DeepSpeed ZeRO allocates a ~477MiB bucket for allreduce operations, a fixed cost that should not scale with model size.

What we observed experimentally

We trained with batch-of-1, short sequences, and a stateless optimiser for 2 steps.
Where gradient accumulation was used: we ran each step for 3 microsteps.

For 7b models, we used naïve model parallelism to vertically-shard costs over two 48GiB GPUs. This type of parallelism is non-performant and therefore unrealistic, but serves as a good way to simulate having one larger GPU.

Note also that when Nvidia markets the capacity of a GPU, they use the measure ‘GB’ but they actually mean GiB.

We were not able to measure DDP costs 7b models, as we had no remaining GPUs over which to replicate the training.

Memory cost (MiB)

We can see that almost none of these training configurations fit on an Nvidia RTX 4090 (24GiB), the largest consumer Nvidia GPU at the time of writing.
It’s even a tight fit on a lot of server Nvidia GPUs (e.g. 40GiB and 48GiB). Bear in mind how these training configurations constitute a lower bound (we aren’t paying for optimiser state, and our activations are unrealistically small).

We see that enabling all of these training techniques is punitive, doubling our memory costs.

It’s clear why Nvidia’s top-end 80GiB GPUs are desirable, and why there is interest in training on AMD’s 128GiB GPUs. We also see that realistically, the training of large models must rely upon the sharding of memory costs across GPUs.

Let’s now try to understand these costs as a function of model size.

Memory cost (bytes/param)

Here we divide the memory usage by the number of model parameters.

Remarkably, some very regular patterns emerge, scaling moreorless linearly with model size.

Our (strawman) baseline of “everything in float32” is indeed 8 bytes/param. This baseline should be low-overhead; it will incur less fragmentation, as it does not need to allocate as many temporary buffers as the other configurations.

Our estimate that “mixed-precision should cost 10 bytes/param” looks close (it’s actually slightly cheaper), but the theory “DDP costs 4 bytes/param more” only holds true when gradient accumulation is enabled. Something to dig into.

Let’s visualise these datapoints as relative overheads compared to our float32 baseline.

Memory cost (bytes/param), relative to float32 baseline

Some results here match our expectations, but there are some surprises also.

Gradient accumulation on its own is almost free (as expected), but it has interesting interactions depending on which other features are enabled. Using it in concert with mixed-precision costs more than the sum of either feature’s cost in isolation.

DDP in isolation costs about 4 bytes/param, as expected. This can be reduced by enabling gradient_as_bucket_view. We conducted a limited exploration of its benefits:

DDP, AMP and GA exhibit interesting interactions with each other. Let’s determine “how much would it cost to enable X, given Y and/or Z are enabled”.

Memory cost (bytes/param) of enabling mixed-precision

Given a starting point of “DDP and GA are enabled/disabled”: how much does it cost to enable mixed-precision?

Memory cost (bytes/param) of enabling Gradient accumulation

Given a starting point of “DDP and AMP are enabled/disabled”: how much does it cost to enable gradient accumulation?

Both DDP and AMP have adverse interactions with gradient accumulation. Enabling either will increase the cost, but the cost of enabling both is greater than the sum of their individual overheads.

This is an overhead we encountered previously when using HF accelerate libraries, but now we see that it reproduces in pure PyTorch.

Memory cost (bytes/param) of enabling DDP

Given a starting point of “AMP and GA are enabled/disabled”: how much does it cost to enable DDP?

In most cases, it costs 4 bytes/param to enable DDP. There is one synergetic case, whereby if mixed-precision is used without gradient accumulation, then the cost of DDP decreases.

Conclusions

The memory costs we derived for PyTorch AMP aren’t a million miles away from what’s described in Transformer Math. Moreover, these heuristics predicted our empirical results reasonably well.

The training techniques studied here incurred greater overheads when used in combination than when used individually, perhaps due to memory fragmentation. Some combinations appeared to exhibit synergies, perhaps due to re-use of freed temporary memory.

For DDP specifically, we learned that gradient_as_bucket_view can be enabled, to save a huge amount of memory (~2.5 bytes/param).

With all techniques enabled (mixed-precision training, gradient accumulation, DDP): our lower bound is somewhere in the region of 15 bytes/param.

It’s clear that beginner-level techniques won’t cut it for fine-tuning LLMs; we need to shard the model parameters, gradients, optimiser states and data over multiple GPUs, trading DDP for more advanced techniques such as PP (pipeline parallelism), FSDP (fully-sharded data parallel), ZeRO (Zero Redundancy Optimizer), and/or TP (tensor parallelism).

As daunting as this looks: we should spare a thought for those who confront pretraining (training a model from scratch), which entails larger batch sizes, longer training runs, and more failed attempts. It can cost $250k to pretrain a 7b model, and 7b is on the small side for LLMs!

But there is an easier way! For the rest of us, parameter-efficient finetuning methods exist, which drastically lower the compute requirements, making fine-tuning possible on consumer GPUs. We hope to explore these techniques in a future blog post.

Acknowledgements

Thanks to collaborators at EleutherAI — including the authors of Transformer Math 101 — for answering questions during the writing of this article. Thanks also to Tim Dettmers for explaining optimiser costs.

Citing

@article{birch2023llmmem,
  title   = "LLM finetuning memory requirements",
  author  = "Birch, Alex",
  journal = "blog.scottlogic.com",
  year    = "2023",
  url     = "https://blog.scottlogic.com/2023/11/24/llm-mem.html"
}

References

This list is limited to our primary reference (Transformer Math 101) and arXiv references; articles and papers which provided BibTeX records. For a more complete list of our references: please refer to the hyperlinks in this article.

[1] Anthony et al. Transformer Math 101 blog.eleuther.ai 2023
[2] Li et al. PyTorch Distributed: Experiences on Accelerating Data Parallel Training VLDB 2020
[3] Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness arXiv preprint arXiv:2205.14135 (2022)
[4] Kingma et al. Adam: A Method for Stochastic Optimization ICLR 2015
[5] Loshchilov et al. Decoupled Weight Decay Regularization ICLR 2019
[6] Dettmers et al. 8-bit Optimizers via Block-wise Quantization ICLR 2022
[7] Huang et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism CVPR 2018
[8] Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models arXiv preprint arXiv:1910.02054 (2019)