Back to Learning
AI Next Labs / Learning
Foundations: where memory and compute go in LLM training
Transformer shapes, the 16 bytes per parameter, activation memory, FLOPs, gradient accumulation and recomputation, with step-through diagrams.
Published
Written by Jaydeep Raijada.
Every parallelism strategy exists because something doesn't fit on one GPU or doesn't run fast enough there. This lesson builds the mental model they all rely on: what a transformer computes, what has to be kept in memory while it trains, and how much of each there is.
Part 0 of 8 in the 5D Parallelism series. Previous: Overview. Next: Data parallelism.
Work through the interactive lesson below: step through each diagram, try the controls, and check yourself with the questions at the end.
References
- Attention Is All You Need (Paper). Vaswani et al., 2017. The transformer architecture.
- Mixed Precision Training (Paper). Micikevicius et al., 2017. FP16 compute with an FP32 master copy and loss scaling.
- Reducing Activation Recomputation in Large Transformer Models (Paper). Korthikanti et al., 2022. Source of sbh(34 + 5as/h) and selective recomputation.
- Training Deep Nets with Sublinear Memory Cost (Paper). Chen et al., 2016. The original gradient checkpointing paper.
- Scaling Laws for Neural Language Models (Paper). Kaplan et al., 2020. Includes the 6·P FLOPs-per-token approximation.
- PaLM: Scaling Language Modeling with Pathways (Paper). Chowdhery et al., 2022. Defines MFU and contrasts it with HFU.
- PyTorch CUDA semantics: memory management (Docs). How the caching allocator reserves and reuses GPU memory.
- The Ultra-Scale Playbook (Guide). Hugging Face. Memory anatomy and recomputation with interactive estimates.