dMX: 面向低精度浮点格式的可微分混合精度分配

# dMX: 用于低精度浮点格式的可微混合精度分配
来源: https://arxiv.org/html/2606.04115
Giuseppe Franco Ian Colbert Pablo Monteagudo\-Lago Felix Marty Nicholas Fraser AMD

###### 摘要

将大型语言模型 (LLMs) 量化到低精度浮点表示对于高效部署至关重要，但在所有层统一应用单一比特宽度在性能与精度权衡方面并非最优。本文介绍了 dMX，一种针对可学习浮点位宽分配的可微混合精度量化框架。我们特别研究了其在开放计算项目 (OCP) 标准定义的微缩放浮点 (MXFP) 数据类型系列中的应用。每层位宽分配被公式化为一个连续优化问题，其中每层的浮点格式由一个标量参数参数化，将多变量设计空间折叠为单个可学习偏移量。在训练期间，此偏移量取连续值，避免了离散量化格式之间的突然振荡。基于温度的退火调度逐步离散化学习到的偏移量，确保最终配置映射到硬件兼容的 MXFP 格式，且训练与推理行为之间无突变。一项目标感知正则化项将平均位宽引导至用户指定的预算，作为推理成本的粗粒度代理，平衡模型质量与部署效率。我们在不同系列的 LLM（如 Llama、Qwen3 和 SmolLM2）上进行了实验，在 WikiText-2 上评估困惑度，并在四个零样本推理基准上评估准确率。在这些设置下，dMX 一致地产生了帕累托占优模型，并优于基于 Kullback-Leibler (KL) 散度的层选择启发式方法，有效地在模型质量与平均位宽之间进行权衡。

## 1 引言

大规模部署大型语言模型 (LLMs) 需要大幅减少内存占用和计算成本，这一挑战随着每一代新模型架构的出现而加剧。量化是实现此目标的主要工具之一：用低精度表示替代高精度参数，同时力求保持模型质量。全量化感知训练 (QAT) 原则上可以恢复大部分精度损失，但对于现代 LLM 而言，其成本仍然高得惊人。这一实际限制推动了更具表现力的数据类型的发展，这些类型可用于训练后量化 (Frantar et al., 2022 (https://arxiv.org/html/2606.04115#bib.bib55); Lin et al., 2024 (https://arxiv.org/html/2606.04115#bib.bib56); Xiao et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib57))。

参照图说明图1：尾数和指数位宽使用连续值与离散的 E2M1 配置时量化网格的比较。在此背景下，FP8、FP6 和 FP4 等浮点格式在相同标称位宽下提供了比整数格式更大的动态范围 (van Baalen et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib51))。相对于按张量或按通道的方案，按组缩放进一步提高了粒度 (Rouhani et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib26))，而动态激活量化则适应推理时每个输入的统计特性。开放计算项目 (OCP) 发布的微缩放 (MX) 标准 (Darvish Rouhani et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib48)) 正式确定了一系列块缩放浮点格式（即 MXFP8、MXFP6 和 MXFP4），这些格式现已得到最新加速器架构的支持。

尽管取得了这些进展，但在所有层统一应用单一低精度格式仍然会导致显著的精度损失。例如，在整个网络中部署 MXFP4 通常会使模型质量降至可接受阈值以下 (Egiazarian et al., 2025 (https://arxiv.org/html/2606.04115#bib.bib37); Sanjeet et al., 2026 (https://arxiv.org/html/2606.04115#bib.bib1))。这一限制催生了混合精度量化，其中不同层或算子获得不同的位宽，从而使模型能够在精度与效率之间达到更有利的平衡。这种方法带来了自身的挑战：跨层的位宽分配是一个组合优化问题 (Cohen et al., 2024 (https://arxiv.org/html/2606.04115#bib.bib60); Wu et al., 2018 (https://arxiv.org/html/2606.04115#bib.bib11))，并且估计量化误差如何在层间相互作用是困难的 (Akbulut et al., 2026 (https://arxiv.org/html/2606.04115#bib.bib59); Ranjan and Savakis, 2025 (https://arxiv.org/html/2606.04115#bib.bib58))。贪婪方法通常评估孤立地量化每一层所带来的性能下降 (Dong et al., 2019 (https://arxiv.org/html/2606.04115#bib.bib6))，或者仅通过局部敏感性代理部分地近似跨层交互 (Dong et al., 2020 (https://arxiv.org/html/2606.04115#bib.bib7); Akbulut et al., 2026 (https://arxiv.org/html/2606.04115#bib.bib59))。它们的决策很大程度上依赖于所选的敏感性指标，并且未能考虑量化误差通过网络非线性累积的情况。

基于梯度的方法提供了一种全局替代方案 (Wu et al., 2018 (https://arxiv.org/html/2606.04115#bib.bib11); Yang and Jin, 2021 (https://arxiv.org/html/2606.04115#bib.bib13))，因为它们通过反向传播联合优化位宽参数，从而捕捉贪婪方法所忽略的跨层依赖关系。可学习的位宽表示此前主要针对整数量化进行研究 (Yang and Jin, 2021 (https://arxiv.org/html/2606.04115#bib.bib13); Huang et al., 2022 (https://arxiv.org/html/2606.04115#bib.bib14); Yang et al., 2021 (https://arxiv.org/html/2606.04115#bib.bib15))；相比之下，很少有工作探索浮点格式的类似思路。在浮点设置中，优化必须应对数据类型本身的结构，这涉及到指数和尾数位宽以及相关量（如指数偏置）。当前硬件仅支持一组有限的格式组合，通常限于 (MX)FP8、(MX)FP6 和 (MX)FP4，并对特殊值表示和指数偏置有预定义选择 (Darvish Rouhani et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib48))。这些约束意味着连续搜索空间最终必须坍缩到一小部分离散的可接受配置，这使得搜索的参数化以及离散化机制成为关键的设计决策。虽然将在第4节 (https://arxiv.org/html/2606.04115#S4) 中对相关工作作更全面的分析，但据我们所知，这是第一个提出面向 MX 格式的基于梯度的位宽分配框架的工作。

本文介绍了一个基于梯度的框架，用于学习逐层浮点位宽。流程概览如图2 (https://arxiv.org/html/2606.04115#S2.F2) 所示。主要贡献如下：

- • 一种用于浮点量化的可微位宽参数化方法，实现了对逐层精度分配进行端到端的基于梯度的优化。
- • 在校准期间，位宽的浮点格式的连续表示，与基于舍入和直通估计器 (STEs) (Bengio et al., 2013 (https://arxiv.org/html/2606.04115#bib.bib38)) 的离散替代方案相比，产生了更平滑的优化曲面。
- • 一种基于温度的退火机制，逐步离散化学习到的位宽，确保收敛到兼容硬件的格式，同时最小化校准时间与推理时间模型行为之间的差距。
- • 通过针对浮点量化的混合精度量化，支持针对特定硬件配置。我们的实验主要关注 MXFP 配置，特别是 MXFP8/MXFP4 和 MXFP6/MXFP4 格式对。
- • 与现有的训练后量化 (PTQ) 算法兼容，能够集成基于梯度的权重优化方法和基于旋转的预处理技术 (Liu et al., 2024 (https://arxiv.org/html/2606.04115#bib.bib36))。

## 2 方法论

参照图说明

图 2：dMX 流程概览。所有蓝色元素突出了本文的主要贡献。预训练的 LLM (左侧) 包含每层 i 的学习到的连续偏移量 βi，该偏移量参数化该层使用的位宽。在前向传播过程中，这些偏移量被映射到离散格式分配 β^=F(β,T)。任务损失和用户定义的对 β 的正则化项 R 共同驱动梯度更新 (红色循环)，而温度退火调度逐步锐化映射，使每个 βi 收敛到兼容硬件的 MXFP 格式 (蓝色循环)。所提出的方法基于三个观察。第一个观察是，位宽在优化过程中无需是整数值。当位宽被视为连续量时，量化方程仍然定义良好，且没有数学约束强制要求离散位宽。

第二个观察直接源于：尽管训练期间允许连续位宽，但最终的分配必须在部署时映射到离散的、硬件支持的格式。对于浮点量化，这种映射比整数设置更具限制性。MX 格式系列仅支持一组特定的配置，每种配置在指数和尾数之间有固定的位分配 (Darvish Rouhani et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib48))。

第三个观察是，梯度可以像针对任何其他可学习参数一样，相对于位宽进行传播，前提是位宽本身作为浮点变量存储。这一思路已成功应用于整数位宽学习 (Yang and Jin, 2021 (https://arxiv.org/html/2606.04115#bib.bib13); Huang et al., 2022 (https://arxiv.org/html/2606.04115#bib.bib14))；将其扩展到浮点量化在很大程度上仍未探索。

这些观察引发了三个设计挑战。第一，如何参数化浮点表示，使得所有进入量化器的量相对于位宽是可微的。第二，如何确保学习到的位宽在训练结束时收敛到硬件兼容的值。第三，如何在用户定义的边界内约束位宽。以下各小节依次讨论每个挑战。

### 2.1 连续浮点位宽参数化

浮点量化的标准公式将浮点值表示为

x(FP)=(−1)S × 2u−b × (1+∑_(i=1)^m M_i ⋅ 2^(−i)), (1)

其中 S∈{0,1} 是符号位，u 是无符号指数且满足 0 ≤ u < 2^e，b = 2^(e−1) − 1 是指数偏置，e 和 m 分别表示指数和尾数的位数，M_i∈{0,1} 是第 i 个尾数位。此公式根据非规格化值的情况进行相应调整。该格式的可表示范围由下式界定

q_max(FP) = (2 − 2^(−m)) ⋅ 2^(2^e − b − 1), q_min(FP) = −q_max(FP). (2)

给定张量 X，元素 x∈X 和粗尺度 s = t / q_max(FP)，其中 t = max(|X|)，每个量化元素 x_q 通过缩放、四舍五入到最近可表示值以及裁剪得到：

x_q = clip( ss ⌊ (x/s) / ss ⌉ ; q_min(FP), q_max(FP) ), (3)

其中 ss 是一个细粒度的 2 的幂尺度，源自每个元素的量级，⌊⋅⌉ 表示四舍五入到最近整数。对于固定输入 x，当 e 和 m 被视为实值参数时，量化输出 x_q(e, m; x) 相对于 e 和 m 是可微的，因为范围边界和偏置随它们连续变化。将公式 (3) 推广到连续值以及位宽的推导见附录 A (https://arxiv.org/html/2606.04115#A1)，关于使用连续值时量化网格如何变化的示例如图1 (https://arxiv.org/html/2606.04115#S1.F1) 所示。

虽然指数和尾数位宽原则上可以独立学习，但这可能产生与现有硬件不兼容的配置。为避免这种情况，通过将每种支持的格式表示为相对于共享基线的偏移来简化格式空间。MXFP8、MXFP6 和 MXFP4 表示分别对应于 E4M3、E2M3 和 E2M1 (Darvish Rouhani et al., 2023 (https://arxiv.org/html/2606.04115#bib.bib48))；本文不考虑 MXFP8 的 E5M2 配置，因为它通常不用于前向传播 (Micikevicius et al., 2022 (https://arxiv.org/html/2606.04115#bib.bib25))。我们可以定义一个通用的 MXFP(β) 配置为

E(2+β)M(1+β), (4)

其中 β 是一个可学习参数。当 β=2 时，我们得到 MXFP8 配置 (即 E4M3)，而 β=0 给出基线 MXFP4 配置。在训练期间，β 可以取 [0,2] 中的任何连续值，从而允许配置之间平滑过渡，并使优化对底层表示的突然变化不那么敏感。这可以应用于 MXFP8/MXFP4 混合精度配置 (e 和 m 均被优化) 以及 MXFP6/MXFP4 配置 (仅优化 m)。最后，将权重和激活的位宽绑定，将问题简化为每层一个可学习标量。

### 2.2 基于温度的退火

训练期间连续表示有利于优化稳定性，但推理需要离散位宽以映射到硬件支持的格式。核心难点在于管理这种过渡而不引入不稳定性；同时，避免训练时间和推理时间行为之间的不连续性至关重要。一种朴素方法是在每次前向传播中使用带有直通估计器 (STE) (Bengio et al., 2013 (https://arxiv.org/html/2606.04115#bib.bib38)) 的舍入算子，这可能导致单个层在整个训练过程中在 MXFP4 和 MXFP8 之间振荡。这种振荡会在损失曲面中产生突然变化，通常导致较差的最终性能，如第3.1节 (https://arxiv.org/html/2606.04115#S3.SS1) 中的实验所确认。

为避免这些振荡，我们建议通过一个温度调控的 sigmoid 函数来平滑位宽偏移量。设计目标是映射在训练早期大致表现为恒等映射，允许优化器探索连续位宽空间 (图 2