# Training state across HBM, host memory and storage Canonical: https://semiaifoundry.com/learn/atlas/training-tier-placement/ Revision: 2026-10-07 Rights: https://semiaifoundry.com/ai-atlas/licenses.html Software · Atlas v1.3.0 # Training state across HBM, host memory and storage Distinguish mutable training state, recomputation, offload and durable checkpoints. [Explore in Atlas](https://semiaifoundry.com/ai-atlas/#training-tier-placement)Tensor state over a training step and across memory tiers; abstraction Learning objectiveDistinguish mutable training state, recomputation, offload and durable checkpoints. ## The mechanism Training retains values for the forward pass, computes gradients and updates parameters and optimizer state. The fastest tier must therefore accommodate more than inference weights. Activation checkpointing stores selected intermediates and recomputes others; persistent checkpoints instead save restartable state outside the live execution path. State sharding changes ownership across devices, while CPU/NVMe offload moves both capacity and work across interfaces. A read-oriented flash tier cannot simply inherit the role of frequently updated optimizer or gradient storage: its write service, software support and movement schedule must be demonstrated. 1. Separate parameters, gradients, master weights, optimizer moments, live activations and temporary buffers. 2. Set precision and sharding rules independently for each class of state. 3. Choose where optimizer computation occurs, keeping its mutable state and transfer requirements consistent. 4. Model recomputation and offload as extra work or movement rather than free capacity. 5. Save version-consistent durable checkpoints and measure recovery alongside steady step time. ## Science and assumptions ### State inventory before sharding $$ M_{\mathrm{state}}=P(b_w+b_g+b_m+b_o)+A+O $$ Mstate = P (bw + bg + bm + bo) + A + O Symbols and units Mstate Live model/training memory estimate bytes P Parameter count parameters bw Stored weight bytes per parameter bytes/parameter bg Gradient bytes per parameter bytes/parameter bm Separate master-weight bytes per parameter bytes/parameter bo Optimizer-state bytes per parameter bytes/parameter A Retained live activations bytes O Other live workspace and metadata bytes Per-parameter byte terms represent the explicitly chosen training representation. bo includes all optimizer-state values per parameter. Assumptions Dense unsharded inventory; omit absent state such as a separate master copy, then apply actual partition/replication factors by tensor class. ## Engineering example Synthetic mixed-precision Adam inventory: 1 billion parameters × (2 weight + 2 gradient + 4 master + 8 moment bytes) = 16 GB before activations, workspace, sharding or replication. It is one representation, not a rule for every optimizer. ## Distinctions and tradeoffs Training repeatedly updates gradient/optimizer/parameter state, so read bandwidth alone cannot characterize an appropriate tier. Offload and recomputation enable larger models at the cost of movement or extra arithmetic; the best choice depends on overlap and recovery requirements. Useful training progress depends on completed and recoverable updates, not the largest parameter count that can be allocated. ## Specifications and interfaces Two checkpoint meanings Activation recomputation versus durable restart Saving an activation is not saving a recoverable training run. Offload boundary Data and computation placement A CPU optimizer also consumes CPU compute and memory bandwidth. ### Interfaces Forward/backward tensors Input Activations, gradients and parameter version Update and durable state Output New parameters, optimizer state and recoverable checkpoint ## Check your understanding ### Which state is inherently write-active during ordinary optimizer updates? 1. Gradients and optimizer state 2. Only the immutable tokenizer vocabulary 3. Only a finalized archival checkpoint Answer and reasoning Gradients and optimizer state Training repeatedly updates gradient/optimizer/parameter state, so read bandwidth alone cannot characterize an appropriate tier. ## Evidence and scope DeepSpeed documents optimizer offload to the host CPU; ZeRO-Infinity studies GPU/CPU/NVMe state management. Neither source establishes an HBF training implementation or its write endurance. Enable the optimizer-state option to budget weights and master/moment state; gradients, activations, workspaces and actual tier transfers are omitted. ## Sources - [DeepSpeed — ZeRO-Offload tutorial](https://www.deepspeed.ai/tutorials/zero-offload/) - [ZeRO-Infinity: Breaking the GPU Memory Wall](https://arxiv.org/abs/2104.07857) [Published lesson data](https://semiaifoundry.com/ai-atlas/lesson-data/training-tier-placement-fa59deb9ce7f62c5.json) · [Atlas releases and licensing](https://semiaifoundry.com/ai-atlas/licenses.html) ## Continue exploring ### Foundations - [Training state and activation rematerialization](https://semiaifoundry.com/learn/atlas/training-memory-state/) - [Place the working set, not just the model](https://semiaifoundry.com/learn/atlas/memory-working-set-placement/) ### Related lessons - [Checkpointing, failure and retained useful work](https://semiaifoundry.com/learn/atlas/checkpoint-retained-work/) - [Parallelism as a memory and communication budget](https://semiaifoundry.com/learn/atlas/parallelism-communication-budget/) - [HBF: a read-oriented tier and its evidence boundary](https://semiaifoundry.com/learn/atlas/hbf-read-tier-contract/)