Unit 8.3 · about 5 h
223 / 324GPU memory: allocated, reserved, peak and memory models
OOM errors name a requested size but never the cause. The fix differs completely between live tensors, cached reserve and fragmentation, and only the counters tell them apart.
- derive
- measure
You can now
- Derive persistent training memory from parameter, gradient, optimizer-state and activation terms at stated precision.
- Read allocated, reserved and peak counters and explain each discrepancy between them.
- Separate live tensors, allocator reserve and outside-allocator memory when an OOM makes no sense.
Covers
- Allocated, reserved and peak memory
- Parameter, gradient, optimizer and activation memory models
- Temporary workspace, input and allocator reserve
Study, in this order
4 steps
- Caching-allocator memory managementPyTorch - checkpointing, numerical accuracy and CUDA semantics
Memory management:
memory_allocatedandmax_memory_allocatedfor tensors,memory_reservedandmax_memory_reservedfor the allocator,empty_cachereleasing only unused segments - Snapshot-based memory debuggingPyTorch - checkpointing, numerical accuracy and CUDA semantics
Generating a Snapshot with
_record_memory_historyand_dump_snapshot; Allocator State History: segments, blocks andactive_allocatedversus inactive states; Identifying Non-PyTorch allocationsDepth · Skip pinned host-memory capture. Use device-only snapshots for this unit.
- Profiler memory analysisPyTorch - reliability and performance recipes
Step 4, Using profiler to analyze memory consumption:
profile_memorywithselfversustotalmemory and theaten::emptyexample - Mixed-Precision Training Memory AccountantTensorTonic - mixed-precision memory and transformer FLOPs
Compute exact storage for parameters, gradients, saved activations, and optimizer state from tensor shapes and byte widths
Focus on
fp32Adam training holds about 16 bytes per parameter before activations: 4 for the weight, 4 for the gradient, 8 for the two moment buffers. Precision and optimizer choice change the constants, not the method.- Allocated sums live tensors; reserved sums allocator segments. A 4 MiB tensor can sit inside 20 MiB of reserve, and
nvidia-smireports closer to reserve. - Peak counters are sticky. Reset them immediately before the measured region or the peak belongs to an earlier phase.
empty_cachereturns unused segments to the driver. It never frees live tensors and never defragments them, so it cannot make a too-large tensor fit.
Reproduce closed-book
fp32 Adam memory for a 10M-parameter model as 160 MB (152.6 MiB) before activations. Then, in a fresh process on a CUDA GPU, allocate one \(1024 \times 1024\) fp32 tensor and read memory_allocated as 4.00 MiB and memory_reserved as 20.00 MiB. Run one matmul and read allocated again: it stays higher by the cuBLAS workspace (about 8 MiB on an RTX 3050 with PyTorch 2.14). State which number nvidia-smi tracks.Traps
- Reading
nvidia-smias live usage. It sees everything the driver gave the process: the CUDA context, cached segments and non-PyTorch allocations such as NCCL. - Counting the cuBLAS workspace as model memory. The CUDA notes say each handle and stream keeps its workspace until explicitly cleared, so take the baseline after a warm-up matmul, not before.
- Comparing a derived model against current allocated only. Peak across forward, backward and optimizer phases is the number that must fit, and it needs a reset counter.
Practice
TensorTonic: Mixed-Precision Training Memory Accountant
Solve the storage accounting from tensor shapes and byte widths
Solve it, then reconcile your formula with the compact model's profiler memory table and write down every included and excluded term.
Self-test
-
A 10M-parameter model trains in
fp32with Adam and no AMP. How many bytes do parameters, gradients and optimizer states hold?40 MB, 40 MB and 80 MB: 160 MB total before activations, workspace, inputs and reserve. -
After allocating one \(1024 \times 1024\)
fp32tensor, allocated reads 4.00 MiB and reserved reads 20.00 MiB. What is the extra 16 MiB?The unused rest of the 20 MiB segment the allocator reserved for this request, held for reuse, not live tensors.nvidia-smitracks closer to reserved than to allocated. -
Peak grows every loop iteration while current stays flat. What was forgotten?
reset_peak_memory_statsbefore the measured region. Max counters never decrease on their own. -
Will
empty_cachelet a run allocate one bigger tensor than the OOM size?No. It returns unused cached segments to the driver for other processes, but it frees no live tensors, so the docs say it cannot increase the memory available to PyTorch.
Completion criterion
Stop when
Your hand-derived parameter, gradient and optimizer bytes match the profiler within stated exclusions, and you can attribute any allocated-versus-reserved gap to segments, peak staleness or outside-allocator memory.
Mark passed, then open 8.4 Only mark this when you can do it without notes.Leads to
- 8.6 Activation checkpointing and recomputation Trade the activation term for recompute once the model is measured.
- 8.4 Precision formats, autocast and gradient scaling Halve the parameter, gradient and activation terms with precision choice.
Comes after
- 7.6 Parameter, FLOP and memory accounting for attentionReconcile the hand count with measured allocator categories.
Resources for this unit
- PyTorch - reliability and performance recipes Use the Profiler, Automatic Mixed Precision, model loading and Benchmark recipes, then the Performance Tuning Guide. Apply each to one compact model from an earlier section and keep a CPU fallback when CUDA is unavailable.
- PyTorch - checkpointing, numerical accuracy and CUDA semantics Use the official user-guide entries for torch.utils.checkpoint, numerical accuracy, CUDA semantics, reproducibility and serialization. Predict memory categories before measuring and define resume equivalence before running.
- TensorTonic - mixed-precision memory and transformer FLOPs Solve only Mixed-Precision Training Memory Accountant and Transformer Training FLOP Estimator. Reconcile both with the compact model's profiler and document every included and excluded term.
Also in this section
- PyTorch - compiler troubleshooting Read graph breaks, guards, recompilations and logging. Compile one stable fragment, intentionally trigger a graph break or recompilation, inspect it, and compare eager and compiled outputs and gradients.
- Deep Learning Training Debugger Use the failure-injection, measurement and evidence rules as the deeper standard. The full Tier 2 project is a later route and is not added to this stage milestone.
- Learn PyTorch - profiler walkthrough Use notebook 10 only if the official profiler route remains unclear. Transfer the relevant annotations into the compact model instead of keeping the tutorial code.
- TrenTorch - precision, checkpointing and resume Solve Loss Scaling, FP16/BF16 Representable Range, Gradient Checkpointing and Resume-from-checkpoint in NumPy, then reproduce each behavior with the PyTorch API in the compact model.
- Python os - `os.replace` Read
os.replacefor the atomic-rename step of checkpoint writes: save to a temp file in the same directory,flushandfsync, then replace the final name. - Hands-On Machine Learning with Scikit-Learn and PyTorch Read Appendix B, Reduced Precision Models and Mixed-Precision Training, for why
float16training underflows and whatfp32primary weights and loss scaling fix.