Collage: Light-Weight Low-Precision Strategy for LLM Training

TL;DR

This paper introduces Collage, a low-precision training strategy for large language models that maintains accuracy while reducing memory and computation costs, enabling faster training with less hardware resources.

Contribution

Collage employs multi-component float representation to compensate for numerical errors in low-precision training, improving efficiency without sacrificing model performance.

Findings

Achieves similar or better training performance compared to mixed-precision methods.

Reduces memory usage by up to 23%.

Speeds up training by up to 3.7 times.

Abstract

Large models training is plagued by the intense compute cost and limited hardware memory. A practical solution is low-precision representation but is troubled by loss in numerical accuracy and unstable training rendering the model less useful. We argue that low-precision floating points can perform well provided the error is properly compensated at the critical locations in the training process. We propose Collage which utilizes multi-component float representation in low-precision to accurately perform operations with numerical errors accounted. To understand the impact of imprecision to training, we propose a simple and novel metric which tracks the lost information during training as well as differentiates various precision strategies. Our method works with commonly used low-precision such as half-precision ($16$-bit floating points) and can be naturally extended to work with even lower precision such as $8$-bit. Experimental results show that pre-training using Collage removes the requirement of using $32$-bit floating-point copies of the model and attains similar/better training performance compared to $(16, 32)$-bit mixed-precision strategy, with up to $3.7\times$ speedup and $\sim 15\%$ to $23\%$ less memory usage in practice.