Memory: a Hopper case study

How to fit it on 80 GiB.Well, technically, nvidia-smi will report only 79.65 GiB. But we'll use 80 GiB as our yardstick in our post because this is a roofline analysis and we need to have some headroom anyway to account for intermediates which aren't modeled here.

GPU memory is like oxygen: you don't really think about it until you run out, at which point you care about nothing else. If your model is OOMing, it's good to have an expectation about how much memory your model should be using--given a choice of parallelism, activation checkpointing and low precision--so you can tell if your implementation is buggy and using more memory than it should be. It's also helpful to know what levers you have for reducing memory usage: you will often be in a regime where your hero run just barely doesn't fit in memory, and you're trying to make it fit without sacrificing too much MFU. For example, the original DeepSeek-V3 pretrain was done on 2048 H800 GPUs (same amount of memory as H100 SXM)--as we'll see, memory is pretty tight in this setting!

Credit: This roofline analysis was heavily inspired by an analysis done by Daniel Haziza. However, any mistakes are my own!

What is memory used for in training?

While your exact memory usage depends on implementation details of your training framework, it's still possible to reason through what the major consumers of memory will be for a given configuration of choices. Let's first set the stage with a naive, global memory accounting of DeepSeek-V3 training assuming compute being done in BF16 (accounting for memory usage in FP8 gives rise to additional complications, which I'd like to address in a dedicated section) and no activation checkpointing.

Parameters. In the previous post, we did a tally of DSv3's parameters, sans MTP.For simplicity, I will continue not to model MTP for this post. Once we add pipeline parallelism, the MTP will be stored on PP rank 0, which has less transformer blocks and is not the memory bottleneck even with MTP. To do compute in BF16, we will want the parameters to be available in BF16. This gives us the following memory usage:

This bar chart needs some explanation. Notably, the x-axis is in logarithmic scale, and the light blue bars are a decomposition of the weights into three subcategories: expert weights, non-expert weights (in the transformer block), and non-transformer-block weights, labeled "emb + lm head" (the initial embedding, the LM head, and — riding along — the final RMSNorm.) Though it's not evident from log scale, these three subcategories cover all the weights in the model. Similarly, don't be deceived into thinking that the non-expert weights are "half" of the expert weights: there are over 40× more expert weights than non-expert weights!

Why logarithmic? Many of the interventions we do involve dividing a quantity by some factor. As a short preview, when we distribute parameters for expert parallelism, we'll divide the expert weights exactly by the EP factor. A log scale lets us express the change from applying EP64 by shortening an expert weights bar by six ticks (each tick is a factor of two), no matter what the original starting quantity was, and allows us to keep a consistent x-scale across all the diagrams in this essay.

Gradients. We need memory to store computed gradients before we update parameters with them. As a simplifying assumption, we will assume that the memory for gradients stays live for the entirety of the optimizer step; while you can avoid this in non-PP settings (since you can try to all-reduce the gradient as soon as possible and then update the weights ASAP so you can reuse the gradient memory for something else), with PP--or really any microbatching--you will need to keep the gradients separate as you accumulate microbatches into it before you do a parameter update.

Separately, in the DeepSeek-V3 paper, it was described that they used FP32 gradients for training stability reasons, so we preserve this choice here--thereby doubling the memory usage of gradients relative to parameters. (In log scale, this uniformly shifts all of these bars to the right one tick.)

Optimizer states. DeepSeek-V3 was trained with AdamW with BF16 first and second moment and FP32 master weights. This means we additionally need another eight bytes for every parameter--double the size of the gradients.

Note: treating the FP32 master weights as "optimizer state" is a DeepSpeed convention. With PyTorch FSDPWe're discussing the implementation here, not ZeRO-3 (which is how FSDP is sometimes interpreted). By default PyTorch FSDP is ZeRO-3 (e.g., FULLY_SHARD), but it also supports other sharding strategies., the master weights are just thought of as the parameters, and FSDP is responsible for casting the master weights to the compute precision and all-gathering them. When reshard_after_forward=False, the compute weights just get kept around, so both the master and compute weights need to be accounted for in a roofline. Because in this essay we'll be using ZeRO to describe the sharding we're doing, we'll stick to the DeepSpeed convention instead of using an explanation that is more natural for PyTorch FSDP's implementation.

Activations. Intermediate activations that are saved for backwards typically contribute to the high watermark memory usage (which usually occurs at the end of forwards / beginning of backwards). We tally up all of the saved for backwards activations, which here are assumed to mostly be saved in the compute precision BF16 (modulo a few small always-FP32 artifacts: router state, attention lse, RMSNorm rstd). This tally assumes a 4096 sequence length per microbatch.

Putting it together. Our current accounting requires more than 100x more memory than a single H100 has. We cannot even fit all of the activations we want to save for backwards.

There are two obvious things we can do here: ZeRO (at ZeRO-3, this would reduce optimizer state, gradient and parameter sizes by a factor of DP) and expert parallelism (reducing expert weights by a factor of EP). However, we cannot do full ZeRO-3, as it turns out a communication roofline shows we will either OOM or be communication bound (we leave this analysis to a future post). And with only ZeRO-1 (optimizer state only--including the FP32 master parameters) and EP, we still cannot fit, so we need to bring out the big guns.

Pipeline parallelism

The primary benefit of pipeline parallelism is that it lets us further reduce the per-device footprint of parameters, gradients and optimizer states by splitting stages of transformer blocks across the PP axis (but in a cheaper way than ZeRO-3, as we only send/recv activations/gradients across PP ranks, rather than the entire sharded weights). However, besides being complicated to implement, there are some subtleties on how exactly PP interacts with the memory usage of your model. The art of building a pipeline schedule is surprisingly deep, so to keep this section short we'll focus specifically on DualPipeV, an improved version of the DualPipe schedule that DeepSeek-V3 was originally trained on.

Let's briefly recap the important facts about DualPipeV. DualPipeV is closely related to a family of what we call "zero bubble" schedules, where the weight gradient computation, which is not critical path in backward computation, is split out from the rest of the gradient computation and used to fill in bubbles (dead time) that would occur when a PP rank runs out of work to do. The "V" in its name refers to the fact that each PP rank is assigned two chunks of the model, and the execution of a single microbatch goes "down" and then "up" in a V shape.

OK, so how does pipeline parallelism affect the memory usage of our model? First, we split parameters/gradients/optimizer state across PP ranks. Because we're looking at a V schedule, we assign two model chunks to each physical PP rank, with roughly four transformer blocks per chunk--click "fold onto the 8 ranks" to see the mapping:

You may have noticed that there is a slight imbalance arising from the special cases at the beginning and end of the model (embedding, dense blocks, LM head). It will turn out that this will be compensated for by the fact that the first PP rank doesn't need to save as many activations for backwards (because there are fewer transformer blocks assigned to it.)

Second, pipeline parallelism implies the use of microbatching, where we take our local batch and split it even further into microbatches, which are then processed in an interleaved manner (this microbatching, by the way, is why we need to maintain dedicated memory for gradients). This per-microbatch processing means that the amount of memory we need to save activations for backwards is not based on the local batch size (global batch size / DP), but instead based on the microbatch size and the pipeline schedule: we don't get to release saved activation memory until we've finished running the backwards of a microbatch. When we have a V shape, so that a single PP rank is responsible for multiple chunks, we furthermore must be careful to talk about the saved activation memory for a specific microbatch and chunk (because the PP rank is responsible for computing two disjoint chunks of transformer blocks on a single microbatch--the V!)

Below, we've reproduced the classic DualPipeV schedule diagram, but by selecting a specific PP rank there is a second visualization of the lifetime of saved activations for each (microbatch, chunk) span. (Microbatches are identified by number, and then a chunk is either light or dark.) If you aren't convinced the spans line up with the schedule, you can click a span to highlight the corresponding entries in the schedule.

It turns out that in the DualPipeV schedule, we only ever need to hold 17 (microbatch, chunk) activations (colloquially, we might call this 8.5 microbatches, because each microbatch is associated with two chunks). This is precisely "d + 1/2" (where d is the number of physical PP ranks) as seen in the cut-in-half row of the table here. The 1/2 term here means that applying DualPipeV pipelining slightly increases the amount of memory you need to save for backwards!

This suffices for the rest of our memory analysis. However, to give some broader context: DualPipeV is not "the schedule to rule them all." In general, zero bubble schedules form a Pareto curve, where you trade off between bubbles and activation memory usage. You would like to reduce bubbles (increase MFU) as much as possible, subject to the constraint of fitting the model in memory. In fact, the next section will be about another such Pareto curve!

Activation checkpointing

With pipeline parallelism, our parameters/gradients/optimizer states are comfortably fitting on a single H100, but we still cannot fit the activations we need to save for backwards. The classic technique for reducing the amount of memory we need for backwards is activation checkpointing: we trade MFU for memory, by simply dropping some quantities in the initial forward and recomputing them in backwards. This strategy is especially attractive as we can change how much recompute we do without changing bitwise numericsYou may have to sweat a bit to actually achieve this: your forward pass must be bitwise deterministic, and you need to accurately track and restore RNG state during recompute. Dropout is so 2022, but stochastic rounding is my new best friend. (this is not true for the low precision tricks in the next section.)

While there are a number of different ways to reason about activation checkpointing, in this post I am going to explain AC on DeepSeek-V3 using torch_remat's programming model. For every operation, we can decide whether to recompute it or not. If we recompute it, it's not necessary to save any of the tensors it needs for backwards--however, the inputs to the operation must be available somehow during recompute, either because: (1) it was independently saved for backwards by another operation, (2) the producer is being recomputed, or (3) we manually saved it (potentially increasing your memory usage!)

You can play around with this belowAI disclosure: the tooltips on the diagram are AI authored and I haven't hand edited them (yet).--every GEMM is also annotated with tall black/gray rectangles indicating how many FLOPs we are spending on forward/recompute (respectively). The chart is preset to the DeepSeek-V3 published recompute strategy, but you should also try setting full recompute and then picking which specific operations you want to save activations for backwards for.

Note that the FLOP counts here are per token, so the routed experts FLOPs are 8x'ed because each token gets processed by eight experts. You have some degrees of freedom in how you setup your recompute/save strategy. However, there are some primary themes:

• You want to avoid recomputing expensive operations, like the FFN GEMMs and the EP communication. On the other hand, non-GEMMs are often cheap (or free, if they are fused) to recompute, so we choose to recompute them even when it's only a very small (or no) memory improvement.
• While your general mental model should be that going from recompute to save will increase the memory usage of your model, it's not always true! In particular, whenever you have an arrow from save to recompute, you have to save inputs into the recompute region. So, for example, you never want to actually recompute the routed + shared addition because you don't want to recompute the a2a combine, and this means marking it as recompute increases the amount of memory you need (there's a little ⚠️ icon to tell you if you've forgotten!) You can also get into situations where toggling save/recompute doesn't change the end-to-end memory usage, because you manage to avoid saving something at the cost of saving something else (⇄ indicates this situation).

Low precision

The second lever we have for activation memory is low precision. Reducing the precision of our GEMMs has multiple benefits:

1. It reduces the amount of memory the GEMM needs to read (low precision inputs are smaller).
2. It can increase tokens per second, because GPUs have more TFLOPs for lower precision tensor cores.
3. It reduces the amount of memory you need to save for backwards (but there is a cost to this).

Because we're doing an analysis for Hopper, we will be primarily looking at Hopper-era quantization schemes. Hopper's MMA instruction supported FP8 without native support for arbitrary tensor/block scaling. Instead, this scaling was done outside the MMA. The original quantization recipes published by NVIDIA advocated for tensor-wide descaling (where scaling can be done entirely after the GEMM), but it turns out trying to apply a single scale to an entire tensor is not great for numerics. DeepSeek-V3 popularized a 128-element 1D block scaling scheme (implemented by incrementally scaling during the K reduction). Later, NVIDIA natively added support for scaling 32-element 1D blocks with Blackwell MXFP8.

I am somewhat reluctant to go super deep into the peculiarities of Hopper quantization, as some of the problems it has to solve don't transfer to Blackwell, but there is one universal feature of weight quantization that is worth talking through. When we are quantizing over 1D blocks, an awkwardness arises because a weight needs to be usable both untransposed and transposed:

• Y = X·W
• dX = dY·Wᵀ
• dW = Xᵀ·dY

Thus, one strategy is to have weights that are quantized into square blocks (128x128 on Hopper, 32x32 on Blackwell); the scale for the block can be repeated to form individual Nx1 or 1xN scales for the actual GEMM. On Blackwell, you can even all-gather the (block quantized) parameters in MXFP8, reducing the memory usage of the (working) parameters, because Blackwell GEMM supports in all combinations of transpose (T) and non-transpose (N) layouts.

However, miserably, on Hopper we must have both WGMMA arguments in physically the correct layout, which means we always have to physically transpose and maintain another copy of the parameter in FP8ᵀ. So if you decide to store your compute weights in FP8 rather than BF16, you don't actually end up saving any compute weight memory in the end (two FP8 = one BF16)--in fact, it's slightly worse because you also have to store the scales for FP8. A more common recipe is to maintain the compute weights in BF16 and quantize the weights on the fly right before they're used.

The strict layout requirements also mean that input activations saved for backwards also need to be transposed. Unlike weights, these aren't quantized into square blocks, to keep the scales for different tokens independent. DeepSeek-V3 will often just requantize the activations in backwards: in cases where it was desirable for the activation saved for backwards to be in FP8, it will dequantize and then requantize it in backwards so that it can be transposed.

The diagram below shows you a DeepSeek-V3 style low precision recipe based on their description of low precision compute in their paperThis diagram is somewhat extrapolated from the DeepSeek-V3 paper, as the paper doesn't spell out every detail. Concretely, we assume the MLA latents were saved as BF16 for recompute and various statistics were kept in FP32 (attention lse, RMSNorm rstd and router state). We also assume that DSv3 did in fact pack their 12-bit E5M6 format at exactly 1.5 bytes per element., and you can use it to explore the consequences of precision changes. For precision (ahem), we prefer not to use the fp8 in the diagram, instead writing e4m3 to indicate the exponent/mantissa size of the fp8 quantities.

Some things to notice from the diagram:

• The precision buttons toggle purely the compute precision (which you can see from how the FLOPs tall rectangles change when you change precision). For FLOPs accounting, we've chosen to model FP8 FLOPs as "half" that of BF16 FLOPs, because FP8 MMAs take twice as many elements as their BF16 counterparts. This is optimistic on Hopper: fine-grained scaling spends part of that 2× on scale handling outside the MMA. And it's also only an approximation, since you have to pay for quantization (not modeled above). You also have a number of degrees of freedom on what precision you keep your intermediate activations in: as a simplification, we simply assume the input activations will be saved in whatever precision the upcoming GEMM is expecting to see, but this isn't a law of nature. In particular, the saved input to SwiGLU isn't directly going to a GEMM, so we could save it in whatever format we want.
• Most of our memory savings arise from choosing to save expert activations in FP8 rather than BF16 for backwards. This is quite beneficial, but it does mean we need to dequantize-requantize in backwards. If you solely want to run the FFN GEMMs in FP8 but don't care about the memory savings, it's also an option to immediately save both the regular and transposed activations for backwards; you can see this via the e4m3ᵀ dual stash (expert inputs) checkbox.
• DeepSeek-V3 has a peculiarity where they have a one off 12-bit FP format specifically for the attention output, as they needed to save some more memory there. This is pretty unusual, and some reproductions don't bother, saving the attention output in BF16. You can toggle this with the E5M6 attn-out stash checkbox.

Fitting it

We can now step through an iterative process by which we can get DeepSeek-V3 to fit on 2048 H100s without OOMing.

In the end, we are left with around 18 GiB of slack for intermediates like communicator workspaces and allocator fragmentation. The rest of the battle is making your actual implementation fit inside this remaining space!

Conclusion

This concludes our roofline memory model for DeepSeek-V3 on H100. Below, I have a version of the model that has all of the knobs we've introduced earlier in this post, so you can play around with any other scenarios you might be interested in. I have one parting thought: it is comparatively simpler to debug OOMs in LLM pretraining, because the memory usage of your trainer doesn't depend on doing distributed communications at scale (assuming your memory allocations and deallocations are deterministic--a very good idea!) In modern PyTorch, you can use fake ProcessGroup (fake PG) to do a very high fidelity simulation of the memory usage of any given rank in an arbitrarily large training config on a single GPU (NB: with pipeline parallelism, you do have to check every pipeline stage individually). It's a good candidate for agentic hill-climbing! Even still, though, I sincerely believe that having the roofline is still very useful, since sometimes the reason you are OOMing because of something very dumb, and the roofline will tell you exactly what it is.