TLDR: MOSS is a new framework for training large language models using FP8 (8-bit floating point) format, which is more efficient than traditional methods. It introduces a two-level microscaling strategy for activations to maintain precision while reducing computational overhead, and an automatic scaling mechanism for weights that predicts scaling factors, eliminating costly runtime calculations. MOSS achieves BF16-level accuracy with significantly higher training throughput, making FP8 LLM training more practical and efficient.
Training large language models (LLMs) is a monumental task, requiring vast computational resources. The shift to lower precision formats like FP8 (8-bit floating point) offers a promising path to accelerate this process and reduce memory demands. However, the inherent reduction in numerical precision with FP8 can introduce significant challenges for maintaining stable and accurate training.
Existing FP8 training frameworks often employ mixed-granularity quantization, where activations use per-group quantization and weights use per-tensor or per-block quantization. While effective for accuracy, per-group quantization for activations can lead to substantial dequantization overhead. This happens because scaling factors are applied along the inner dimension of matrix multiplications, forcing slower CUDA Cores to perform dequantization within the main computation loop, thereby negating the performance benefits of FP8.
Furthermore, many current frameworks rely on ‘just-in-time’ scaling, which dynamically adjusts scaling factors based on the current data distribution. This online quantization method is inefficient for FP8 training as it involves multiple memory reads and writes, adding to the computational burden.
Introducing MOSS: A Novel Approach to FP8 LLM Training
To overcome these limitations, researchers have developed MOSS, a new FP8 training framework designed for both efficiency and numerical stability. MOSS introduces two key innovations:
1. Two-Level Microscaling for Activations: MOSS employs a progressive two-level quantization strategy for sensitive activations. This involves a high-precision global scale (FP32) combined with compact, power-of-two local scales (E8M0). This hierarchical approach balances precision and dequantization cost. Crucially, MOSS’s design shifts the costly FP32 dequantization operations out of the main matrix multiplication (GEMM) loop, confining them to the epilogue. This allows the core GEMM operations to run much faster on high-throughput Tensor Cores, significantly reducing overhead compared to previous methods like COAT.
2. Automatic Scaling for Weights: For weights in linear layers, MOSS leverages the predictable update behavior of Adam-like optimizers. By understanding that weight updates are inherently bounded by the learning rate, MOSS can predict and adjust scaling factors during training. This eliminates the need for expensive runtime ‘max-reduction’ operations, which typically involve reading all FP32 values from memory to compute scaling factors. MOSS’s automatic scaling introduces negligible and constant runtime overhead, making it far more efficient than just-in-time scaling.
Also Read:
- MoSKA: A New Architecture for Faster and More Efficient Long-Sequence LLM Inference
- QUARK: Accelerating Transformers with Quantization and Circuit Sharing
Performance and Results
MOSS has been rigorously evaluated by training 7-billion parameter models, including OLMo-7B and LLaMA-2-7B, on GPU-constrained settings. The results are compelling:
- Accuracy: MOSS achieves lossless accuracy, matching the performance of BF16 baselines on both pre-training and fine-tuning tasks. This demonstrates that MOSS preserves model accuracy and in-context learning capabilities at a level comparable to higher-precision training.
- Efficiency: MOSS delivers significant speedups. It achieves up to 34% higher end-to-end training throughput on OLMo-7B compared to BF16. Furthermore, it outperforms the state-of-the-art FP8 training framework, COAT, by 12.3%.
- GEMM Kernel Efficiency: An ablation study confirmed the superior efficiency of MOSS’s GEMM kernel design, showing performance comparable to Transformer Engine and significantly outperforming COAT.
The framework also generalizes effectively to downstream fine-tuning tasks, proving its versatility beyond just pre-training. By combining coarse-grained global scaling with lightweight power-of-two microscales and an intelligent automatic scaling mechanism, MOSS offers a practical and hardware-friendly solution for scalable FP8 training.
This advancement in FP8 training efficiency could significantly lower the computational costs and energy consumption associated with training large language models, making advanced AI research and development more accessible. For more technical details, you can refer to the original research paper. Read the full paper here.
