LLM Memory Profiling, in depth: predict the GPU budget, measure each phase, and attribute every gigabyte before the OOM

Four numbers called GPU memory

Four different numbers get called GPU memory, and confusing them is the first source of wrong conclusions. PyTorch does not ask the driver for memory every time it creates a tensor. Its caching allocator requests large segments with cudaMalloc, carves tensors out of them, and keeps freed blocks for reuse instead of returning them. So the memory tensors occupy and the memory the process holds are different quantities, and the driver sees a third.

Peak allocated is the number that matters for whether a step fits, because the peak is reached briefly and then released. A once-per-step memory_allocated() log taken after the optimizer step shows only the floor and will tell you the job has plenty of room right up to the moment it fails. The difference between reserved and allocated is cached memory; the difference between device used and reserved is memory outside PyTorch's allocator, which is typically several hundred megabytes of CUDA context per process plus whatever NCCL, cuBLAS workspaces and custom kernels took directly.

Predict the budget before you run

Before running anything, write down what the step should need. Five terms cover nearly every decoder-only transformer job.

• Model states. With mixed-precision AdamW, each parameter costs about 16 bytes: two for bf16 weights, two for bf16 gradients, four for an fp32 master copy and eight for the two fp32 Adam moments. Keeping weights and gradients in fp32 under autocast also lands at 16. This is the accounting the ZeRO paper used, and it is why sharding the optimizer state is the first lever for big models.
• Saved activations. Tensors the forward pass keeps for the backward pass. Korthikanti et al. estimated roughly s*b*h*(34 + 5*a*s/h) bytes per layer for a 16-bit transformer with sequence length s, batch b, hidden size h and a heads. Fused attention kernels such as FlashAttention do not store the s-by-s score matrix, which removes the second term. Treat the 34 as a floor: gated MLPs with wider intermediate layers save more.
• Logits and loss. The output projection produces tokens times vocabulary values. With modern 128K-token vocabularies this is frequently the single largest tensor in the step, and the cross-entropy computation commonly upcasts it to fp32 and materialises its gradient.
• KV cache (inference). Two tensors per layer per token: 2 * layers * kv_heads * head_dim * bytes.
• Overhead. CUDA context, communication buffers, temporary workspaces and allocator rounding. Budget one to two gigabytes and verify it.

Measure each phase

Prediction tells you what should happen; measurement tells you what did. Measure each phase separately by resetting the peak counter before it and reading it after. Run a few warm-up steps first: the optimizer allocates its moment tensors lazily on the first step(), so step 0 understates the floor and many jobs that fit for one step fail on the second.

import torch

def gb(x): return x / 2**30

def profile_step(model, batch, optimizer):
    torch.cuda.synchronize()
    out = {"floor": gb(torch.cuda.memory_allocated())}
    for phase in ("forward", "backward", "optimizer"):
        torch.cuda.reset_peak_memory_stats()
        if phase == "forward":
            loss = model(**batch).loss
        elif phase == "backward":
            loss.backward()
        else:
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
        torch.cuda.synchronize()
        out[phase + "_peak"] = gb(torch.cuda.max_memory_allocated())
    out["reserved"] = gb(torch.cuda.memory_reserved())
    free, total = torch.cuda.mem_get_info()
    out["outside_torch"] = gb(total - free - torch.cuda.memory_reserved())
    stats = torch.cuda.memory_stats()
    out["alloc_retries"] = stats.get("num_alloc_retries", 0)
    return out

Compare each phase peak with its predicted term. The num_alloc_retries counter is an early warning: it increments when the allocator could not find a free block, flushed its cache and retried cudaMalloc. A job with a rising retry count is running at the edge, and each retry also costs a synchronisation that shows up as a throughput stall. The outside_torch figure is the gap no allocator tool can see.

Attribute the peak with a snapshot

When a phase peak exceeds its prediction, the allocator's history tells you which allocations made up the peak and which line of Python created each one. Recording is off by default because it captures a stack trace per allocation; turn it on for a handful of steps.

import torch

torch.cuda.memory._record_memory_history(max_entries=100_000)
for step, batch in zip(range(4), loader):          # 1 warm-up + 3 measured steps
    loss = model(**batch).loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
torch.cuda.memory._dump_snapshot("step_memory.pickle")
torch.cuda.memory._record_memory_history(enabled=None)   # stop recording

Drag the pickle into the viewer at pytorch.org/memory_viz, which runs locally. The Active Memory Timeline view draws every live tensor as a band over time, so a training step appears as a sawtooth; clicking the tallest band at the peak shows the stack that allocated it. The Allocator State History view shows segments and blocks, which is where fragmentation becomes visible. The leading underscore means the API is not stable, so check the docs for your installed version, and always set max_entries to keep the snapshot small. For the operator-level view that joins memory with time, the profiler's profile_memory=True option is covered in the PyTorch profiler guide.

Worked example: a 1B full fine-tune

Take a full fine-tune of Llama 3.2 1B: about 1.24 billion parameters, 16 layers, hidden size 2048, a vocabulary of 128,256 tokens, bf16 with AdamW, micro-batch 4 sequences of 4,096 tokens (16,384 tokens per step). The prediction:

Two observations fall out before a single kernel runs. First, the logits term is larger than the activations of all sixteen layers combined. A 1B model has a vocabulary-sized output projection that is as wide as that of an 8B model, so the loss dominates small models in a way parameter count hides. Second, the peak is not a sum of steady states: the loss spike sits on top of fully saved activations, at the exact moment forward ends and backward begins. That is where the timeline should peak, and a snapshot should confirm it by attributing the tallest bands to the language-model head and the cross-entropy call.

On an 80 GB card the job fits with room to spare. On a 48 GB card it fails, and halving the batch would 'fix' it by halving throughput. The account suggests better moves. A fused or chunked cross-entropy, which computes the loss over slices of the sequence without materialising the full fp32 logits (Liger Kernel's fused linear cross-entropy is one implementation), removes most of the red block. Activation checkpointing that keeps only each layer's input drops the saved activations to roughly 2.2 GB at the price of an extra forward pass, about a third more compute. Either change alone fits 48 GB at the original batch size. Measured numbers will land above these predictions, because the gated MLP saves more than the 34-byte coefficient assumes.

Inference: weights and the KV pool

Inference has a different shape. There are no gradients or optimizer state; the floor is the weights, and almost everything else is the KV cache. For Llama 3.1 8B (32 layers, 8 key-value heads of dimension 128, bf16), each token costs 2 x 32 x 8 x 128 x 2 = 131,072 bytes, or 128 KiB. The weights take about 16 GB. Serving engines such as vLLM reserve a fraction of the device (gpu_memory_utilization, 0.9 by default), run a profiling forward pass at the maximum batch shape to measure peak activation memory, and hand everything left to the KV cache pool. The KV cache sizing guide walks through turning that pool into concurrency.

Profiling a server therefore means checking three things: that the profiling pass really saw the worst case (a longer prompt or a larger sampling batch at runtime can exceed it), that memory outside the engine, such as a second process or a tokenizer using the GPU, has not eaten the headroom, and that the pool is not oversized for the traffic.

Finding leaks

A leak in a training loop is memory that grows step over step. Plot memory_allocated() after each optimizer step; a flat line is healthy, and a staircase is a leak. The causes are nearly always references the Python program keeps by accident.

• Accumulating tensors with graphs. total_loss += loss keeps every step's autograd graph alive. Use loss.item() or loss.detach() for logging.
• Hooks that save activations. Forward hooks added for debugging or feature extraction often store tensors in a dictionary that nobody clears.
• Holding an exception. Catching an out-of-memory error and retrying inside the except block keeps the traceback, and with it every frame's tensors, alive. Set a flag in the handler and retry outside it.

The snapshot finds leaks quickly: record ten steps, and the bands that start in step 2 and never end, with the stack that created them, are the leak.

Fragmentation and allocator settings

Sometimes the out-of-memory message itself says the problem is not size. It reports how much is reserved but unallocated, and when that is several gigabytes the allocator held enough free memory in total but not in one contiguous block. Variable sequence lengths are the usual cause: each batch leaves differently sized holes. The allocator's behaviour is controlled by PYTORCH_ALLOC_CONF (recent releases; PYTORCH_CUDA_ALLOC_CONF remains as an alias). Setting expandable_segments:True lets segments grow in place instead of allocating a new one per oversized request, which often removes the problem outright; max_split_size_mb stops large cached blocks being split for small requests; and garbage_collection_threshold makes the allocator release cached blocks before it reaches the limit. The mechanics are covered in the memory pool article. Padding or bucketing sequence lengths to a few fixed sizes attacks the cause rather than the symptom.

Failure modes

The recurring ways memory profiling itself goes wrong:

• Measuring the floor. Logging allocated memory once per step after the optimizer misses the peak entirely. Log the peak and reset it.
• Profiling step 0. Optimizer state does not exist yet, compilation caches are cold, and the allocator has not settled. Measure steps 2 to 5.
• Ignoring the longest batch. A profile on average-length data says nothing about the batch that packs to the maximum length. Profile with a synthetic worst-case batch.
• Trusting nvidia-smi alone. It includes the cache, so it reads high after a peak and never comes down, which looks like a leak but is not.
• Profiling one rank. In data-parallel or sharded jobs, the rank that holds an uneven shard, the pipeline's first stage or the one doing evaluation can peak higher. Collect peaks from every rank.
• Calling empty_cache() as a fix. It returns cached memory to the driver, which helps other processes but forces later re-allocation and does not lower the peak.

Choosing the fix

Once the account says which term is too big, choose the fix that targets that term and costs the least. A rough order:

Reducing the micro-batch is last on purpose: it costs throughput and hides the term that was actually wrong. The training step anatomy article shows how each phase's kernels create and free these tensors.

What to do next

1. Write the five-term budget for your current job and compare it with the device size before the next run.
2. Add per-phase peak logging with reset_peak_memory_stats() to steps 2 to 5 and record the allocator retry count.
3. Capture one snapshot of three steps and identify the allocation at the peak.
4. Run once with a synthetic maximum-length batch and confirm it fits with at least 5% headroom.
5. Graph allocated memory after each step for a few hundred steps and confirm it is flat.
6. If the peak sits at the loss, try a fused cross-entropy before touching batch size; if it sits in the forward pass, try selective checkpointing.
7. For serving, compare the KV pool's peak occupancy with its size and resize the pool or the replica count to match.