Reproduce “Collage: Light-Weight Low-Precision Strategy for LLM Training” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Tao Yu, Gaurav Gupta, Karthick Gopalswamy, Amith Mamidala, Hao Zhou, Jeffrey Huynh, Youngsuk Park, Ron Diamant and 2 more

The COLLAGE framework introduces a low-precision strategy for training Large Language Models (LLMs) by integrating Multi-Component Floats (MCF) into the optimizer, eliminating the need for high-precision master weights. This approach achieves up to 3.7x faster training throughput and an average of 15-28% reduction in peak GPU memory usage, while matching or exceeding the accuracy of state-of-the-art mixed-precision methods.

Key insights

  • Compensating for roundoff errors in low-precision arithmetic through Multi-Component Floats (MCF) enables robust LLM training without relying on high-precision master weights.
  • The `β2` coefficient in Adam-like optimizers is particularly sensitive to low-precision rounding, and its representation as an MCF expansion is crucial for stability with aggressive values like `0.999`.
  • The proposed "effective descent quality" (EDQ) metric effectively quantifies information loss during optimizer updates, providing a diagnostic tool to understand and improve precision strategies.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.