MoR: Mixture Of Representations For Mixed-Precision Training

TL;DR

MoR introduces a dynamic, property-aware approach to mixed-precision training by coupling Group Amax Mantissa (GAM) scaling with a Mixture-of-Representations framework that selects FP8 or BF16 representations at tensor or sub-tensor granularity. The GAM scaling decouples mantissa and exponent, storing a single group mantissa and per-block exponents to preserve the tensor’s maximum value with FP32 precision while minimizing metadata. Empirically, tensor-level MoR achieves model-quality parity with BF16 while quantizing a large majority of tensors to FP8 (e.g., over 98% in experiments), across partition strategies and training configurations; sub-tensor MoR shows potential but can exhibit task-specific overfitting when using richer representations. Overall, MoR demonstrates robust, dynamic quantization decisions that can substantially reduce precision overhead without sacrificing downstream performance, offering a path toward more memory- and compute-efficient large-model training and enabling exploration of even lower-precision formats such as NVFP4.

Abstract

Mixed-precision training is a crucial technique for scaling deep learning models, but successful mixedprecision training requires identifying and applying the right combination of training methods. This paper presents our preliminary study on Mixture-of-Representations (MoR), a novel, per-tensor and sub-tensor level quantization framework that dynamically analyzes a tensor's numerical properties to select between a variety of different representations. Based on the framework, we have proposed and experimented concrete algorithms that choose dynamically between FP8 and BF16 representations for both per-tensor and sub-tensor level granularities. Our universal approach is designed to preserve model quality across various quantization partition strategies and datasets. Our initial findings show that this approach can achieve state-of-the-art results with 98.38% of tensors quantized to the FP8 format. This work highlights the potential of dynamic, property-aware quantization while preserving model quality. We believe this approach can generally improve the robustness of low precision training, as demonstrated by achieving FP8 accuracies that are on par with existing approaches without the need for fine-grain partitioning, or can be used in combination with other training methods to improve the leverage of even lower precision number formats such as NVFP4.

Figures (21)

• Figure 1: A diagram of the Transformer block architecture.
• Figure 2: Illustration of tensor-level MoR. A single decision (E4M3 or BF16) is made for the entire tensor $X$. This decision is based on a global relative error, which can be computed using various internal partitioning strategies for the quantization step (e.g., per-block, no-partition, per-channel). The process uses the GAM scaling algorithm from Algorithm \ref{['algo:gam_scaling']} to determine a shared mantissa and per-partition exponents (purple squares).
• Figure 3: An illustration of sub-tensor MoR. Blocks within matrices A and B are quantized to different formats (E4M3, E5M2, BF16). The corresponding GEMM operation may require upcasting blocks (e.g., $B_{31}$) to a higher precision (BF16) before multiplication.
• Figure 4: The fake quantization workflow. The input and output data type are kept in BF16 precision, but the process simulates the precision loss of the target format.
• Figure 5: Training loss, validation loss, and the parameter L2 Norm using the first training configuration. In the legend, from top to bottom, the curves represent: BF16 Baseline, MoR Per-Block, MoR Per-Channel, and MoR Per-Tensor.