Reproduce “Stochastic Rounding for LLM Training: Theory and Practice” on OpenResearch

About this paper

Kaan Ozkara, Tao Yu, Youngsuk Park

Large language model training with BF16 precision and stochastic rounding (SR), when optimized with higher learning rates, achieves up to 2.6% better validation perplexity, 1.54x faster throughput, and 30% memory reduction compared to standard mixed precision methods. This strategy effectively enhances training efficiency while improving final model accuracy.

Key insights

• Stochastic rounding implicitly regularizes the loss function with a quantization error penalty term, suggesting that higher learning rates are necessary to prevent training stagnation.
• SR introduces a strictly smaller convergence error in the Adam optimizer compared to nearest rounding, allowing its inherent noise to be absorbed within Adam's convergence bounds with appropriate hyperparameters.
• The additive noise from SR contributes to greater training stability and robustness, enabling the use of higher learning rates that typically destabilize standard mixed precision training.

Summary by alphaXiv.