Visualize and understand GPU memory in PyTorch

Source: HuggingFace Date: 2024-12-24 URL: https://huggingface.co/blog/train_memory

Summary

Educational tutorial: practical guide to GPU memory profiling and estimation during PyTorch training. Covers using torch.cuda.memory._record_memory_history() and the pytorch.org/memory_viz tool. Provides memory formulas for each component (model parameters, activations, gradients, AdamW optimizer state). Key formula: Total = Model Memory + Optimizer State + max(Gradients/Intermediates, Activations). Benchmark example: Qwen2.5-1.5B at batch size 16 / seq 256 uses 6GB for model weights and 5M activations per input token. Notes that peak memory location shifts with batch size.

Implications

Transformers library trajectory. This post is diagnostic infrastructure for practitioners trying to fit models into GPU memory budgets. Activation memory following a linear relationship with model parameters (A ≈ 4.69×10⁻⁴ × N + 1.85×10⁶) is a practical heuristic that helps teams estimate memory requirements before spinning up training runs.

Open-weights ecosystem health. GPU memory is the binding constraint for most open-weights fine-tuning work. Posts that teach practitioners to measure rather than guess their memory usage reduce the trial-and-error cycle for new fine-tuning experiments — a real time and cost savings for the community.