Optimizer

ZeRO: Partitioning Optimizer State, Gradients, and Parameters Plain DDP is clean but wasteful. Every data-parallel rank stores the same parameters, gradients, fp32 master weights, and optimizer moments. At small scale that redundancy is convenient. At LLM scale it is the memory wall. ZeRO, the Zero Redundancy Optimizer, keeps data-parallel semantics and removes the redundant storage one category at a time. ZeRO-1 shards optimizer state. ZeRO-2 shards optimizer state and gradients. ZeRO-3 shards optimizer state, gradients, and parameters (arXiv:1910.02054). ...

Megatron Internals III: Mixed Precision, Loss Scaling, and Grad Clipping Mixed precision training is not just “turn on fp16.” It is a state machine: fast low-precision tensors do forward and backward, stable fp32 tensors receive optimizer updates, and every rank agrees whether the step is valid before any shard changes. The classic recipe comes from Micikevicius et al., Mixed Precision Training: use lower precision where hardware is fast, keep fp32 master weights for updates, and use loss scaling when fp16 gradients underflow. Megatron wraps that recipe around tensor, pipeline, and data parallelism. ...