Stochastic Rounding for LLM Training: Theory and Practice

TL;DR

This paper introduces stochastic rounding for large language model training, providing theoretical insights and practical benefits such as improved efficiency, stability, and reduced memory usage in low-precision training strategies.

Contribution

It offers the first theoretical analysis of stochastic rounding with Adam optimizer and extends BF16+SR to distributed training, demonstrating superior performance in large-scale models.

Findings

BF16 with SR outperforms traditional mixed precision strategies.

Achieves up to 1.54x higher throughput and 30% less memory usage.

Empirical validation on models up to 6.7B parameters.

Abstract

As the parameters of Large Language Models (LLMs) have scaled to hundreds of billions, the demand for efficient training methods -- balancing faster computation and reduced memory usage without sacrificing accuracy -- has become more critical than ever. In recent years, various mixed precision strategies, which involve different precision levels for optimization components, have been proposed to increase training speed with minimal accuracy degradation. However, these strategies often require manual adjustments and lack theoretical justification. In this work, we leverage stochastic rounding (SR) to address numerical errors of training with low-precision representation. We provide theoretical analyses of implicit regularization and convergence under the Adam optimizer when SR is utilized. With the insights from these analyses, we extend previous BF16 + SR strategy to be used in distributed settings, enhancing the stability and performance for large scale training. Empirical results from pre-training models with up to 6.7B parameters, for the first time, demonstrate that our BF16 with SR strategy outperforms (BF16, FP32) mixed precision strategies, achieving better validation perplexity, up to $1.54\times$ higher throughput, and $30\%$ less memory usage.