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.
- Separate parameters, gradients, master weights, optimizer moments, live activations and temporary buffers.
- Set precision and sharding rules independently for each class of state.
- Choose where optimizer computation occurs, keeping its mutable state and transfer requirements consistent.
- Model recomputation and offload as extra work or movement rather than free capacity.
- Save version-consistent durable checkpoints and measure recovery alongside steady step time.
Science and assumptions
State inventory before sharding
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?
- Gradients and optimizer state
- Only the immutable tokenizer vocabulary
- 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.