Megatron Internals III: Mixed Precision, Loss Scaling, and Grad Clipping
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. ...