Visualize and understand GPU memory in PyTorch – Hugging Face Blog Summary

TL;DR

The post shows how to record and visualize GPU memory snapshots in PyTorch, explains what each part of a memory profile represents (model parameters, optimizer state, activations, gradients, optimizer intermediates), and gives a generalized formula for estimating peak memory usage during training.

🔎 The PyTorch visualizer

PyTorch provides a built‑in tool to record and visualize GPU memory usage. By calling torch.cuda.memory._record_memory_history before running code and torch.cuda.memory._dump_snapshot afterward, you obtain a profile.pkl file that can be viewed at https://pytorch.org/memory_viz. The visualizer displays a timeline of memory allocations and releases.

In the example with a simple nn.Linear layer, the graph shows:

  • Model creation allocates ~2 GB for weights and biases (float32).
  • Each input tensor adds ~200 MB.
  • Each forward pass output adds ~1 GB.
  • Activations from previous steps are retained until they are no longer needed for backpropagation, after which their memory is freed.
  • Reassigning variables releases the previously referenced tensors.

📊 Visualizing Memory During Training

For a realistic training loop with a large language model (Qwen/Qwen2.5-1.5B) and AdamW optimizer, the memory profile shows three spikes, one per training iteration.

The profile can be broken down as follows:

  1. Model Initialization – model parameters (blue) occupy memory and stay allocated until training ends.
  2. Forward Pass – activations (orange) are computed and stored layer by layer; they peak at the loss calculation.
  3. Backward Pass – gradients (yellow) are computed; activations are discarded, causing the orange zone to shrink.
  4. Optimizer Step – optimizer state (green) is initialized once; optimizer uses gradients to update parameters, temporarily storing optimizer intermediates (red). After the update, gradients and intermediates are freed.

The pattern repeats for each iteration, producing the observed spikes.

📐 Estimating Memory Requirements

The peak memory usage is the highest point in the profile, which may occur during the forward pass or the optimizer step depending on batch size.

A general expression that covers both cases is:

Total Memory = Model Memory + Optimizer State + max(Gradients + Optimizer Intermediates, Activations)

where each term is defined below.

Model Parameters

Model Memory = N × P

  • N = number of parameters
  • P = precision in bytes (e.g., 4 for float32)

For the Qwen2.5-1.5B example (1.5 B parameters, float32): Model Memory = 1.5 × 10⁹ × 4 bytes = 6 GB.

Optimizer State

For AdamW, which stores two moments per parameter: Optimizer State Size = 2 × N × P

Gradients

Gradients Memory = N × P (same size as model parameters).

Optimizer Intermediates

Optimizer Intermediates Memory = N × P (same size as model parameters).

Activations

Activation memory depends on batch size (B), sequence length (L), and the number of activations per token (A). Activation Memory = A × B × L × P

A can be measured directly with forward hooks, but the post provides a heuristic derived from a linear fit across models:

A = 4.6894 × 10⁻⁴ × N + 1.8494 × 10⁶

Using this heuristic lets you estimate activation memory without running a full forward pass.

Total Memory Formula (combined)

Substituting the component formulas gives:

Total Memory = N×P + 2×N×P + max(N×P, N×P, A×B×L×P)
             = 3×N×P + max(N×P, A×B×L×P)

Since the max term will always be at least N×P, the expression simplifies to:

Total Memory = 3×N×P + A×B×L×P

when activations dominate; otherwise the optimizer step dominates and the total is 4×N×P.

The post includes a small tool to plug in N, P, B, L and obtain an estimate.

🚀 Next steps

Understanding memory profiles equips you to look for ways to reduce usage, such as lowering batch size, using gradient checkpointing, or switching to mixed‑precision training. The suggestions in the TRL documentation’s “Reducing Memory Usage” section apply broadly to any PyTorch‑based training.

🤝 Acknowledgements

Thanks to Kashif Rasul for feedback and suggestions on the blog post.

Sources