Quartet: Native FP4 Training Can Be Optimal for Large Language Models
Quartet: Native FP4 Training Can Be Optimal for Large Language Models
标题:Quartet:原生FP4训练可成为大语言模型的最优选择
作者/机构:Roberto L. Castro (ISTA), Andrei Panferov (ISTA), Soroush Tabesh (ISTA), Oliver Sieberling (ETH Zürich), Jiale Chen (ISTA), Mahdi Nikdan (ISTA), Saleh Ashkboos (ETH Zürich), Dan Alistarh (ISTA & Red Hat AI)
A1 主要贡献
本文研究了在NVIDIA Blackwell新架构支持下,使用FP4低精度格式直接训练大语言模型(LLM)的挑战和机遇。
核心问题:尽管FP4等极低精度计算能显著提升吞吐量和能效,但现有的FP4训练算法在应用于当前的大语言模型时,会面临严重的精度下降和稳定性问题,或者需要回退到更高精度进行部分计算。这导致了新兴的低精度硬件能力与实现精确、端到端量化模型训练所需的算法支持之间存在差距。
研究目标:本文旨在填补这一差距,通过对硬件支持的FP4训练进行系统性研究,设计出一种能够在所有主要计算(即线性层)中都使用低精度,同时保持高精度的端到端FP4训练方法。
创新点与主要贡献:
1. 提出一种新的量化训练方法比较框架:通过引入“诱导缩放定律”(induced scaling law)来比较不同的量化训练方法。该定律通过两个关键参数来量化性能:参数效率 effN(与前向传播精度相关)和数据效率 effD(与反向传播精度相关)。一个方法若能同时提升这两个指标,则被认为是更优的。
2. 揭示了效率参数与量化误差和偏差的关联:研究发现,参数效率 effN 与训练方法的前向压缩误差直接相关,而数据效率 effD 与其梯度估计器的偏差相关联。作者引入一种新的“失准度量”(misalignment metric)来衡量这种偏差。这些度量结合实际硬件的加速比,可以预测在给定计算和数据预算下,训练模型达到目标精度的“最优”低精度设置。
3. 设计并实现了Quartet算法:基于上述分析框架,本文设计了一种名为Quartet的算法,专门用于MXFP4原生训练。该算法通过结合先前最先进的量化感知训练(QAT)【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】和量化反向传播优化技术【索引44,Training llms with mxfp4,2025,arXiv】,最大化了参数效率和数据效率。其核心技术贡献是为新的Blackwell架构开发了一个高度复杂的、高效的GPU实现。
4. 实验验证Quartet的有效性:通过在C4数据集上预训练Llama系列模型,实验证明:1)在不同的计算预算和模型规模下,Quartet的精度均优于先前的方法;2)其快速实现使其性能超越了高度优化的FP8核心。这首次证明了MXFP4在实践中可以在精度与速度的权衡上成为一种“最优”的训练方案,具备与FP16和FP8竞争的潜力。
A3 背景知识
量化网格。量化是将高精度的内部模型状态(如权重、激活值或梯度)映射到一个低精度的离散集合,即量化网格。该网格可以是均匀的(如整数 quantization),也可以是非均匀的(如浮点(FP)quantization,其数值间距大致呈指数分布)。由于原始值与网格的尺度可能不同,通常会存储一个高精度的尺度因子s。对于向量x,量化过程可表示为q(x) = round(x/s; grid),而原始值可近似重构为x̂ = s · q(x)。尺度的常见选择是设为x中的最大绝对值(absmax)以避免裁剪,或通过优化以最小化均方量化误差【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】。
量化粒度。除了网格选择,量化方法在尺度的粒度上也有所不同。一个尺度值可以被整个张量共享【索引3,Halo: Hadamard-assisted lower-precision optimization for llms,2025,arXiv】,也可以在每行或每列间共享【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】,或者在更细粒度的自定义块上共享,如2D块【索引54,Jetfire: efficient and accurate transformer pretraining with int8 data flow and per-block quantization,2024,arXiv;索引29,Deepseek-v3 technical report,2024,arXiv】或1D块【索引33,Ocp microscaling formats (mx) specification version 1.0,2023,Open Compute Project;索引44,Training llms with mxfp4,2025,arXiv】。值得注意的是,最新的Blackwell GPU架构【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】引入了对MXFP4/6/8和NVFP4格式的硬件支持。MXFP格式【索引33,Ocp microscaling formats (mx) specification version 1.0,2023,Open Compute Project】在每32个元素的1D块上共享一个FP8的2的幂次尺度,而NVFP4【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】使用FP8(E4M3)尺度和16个元素的1D块。
舍入。量化通常涉及舍入操作。例如,确定性地舍入到最近的网格点(deterministic rounding),可以得到最低的均方误差(MSE)。相比之下,随机舍入(stochastic rounding)引入了随机性,根据输入值与邻近网格点的距离以一定概率向上或向下舍入。虽然随机舍入可能引入更高的MSE,但它有助于控制偏差,这对于维持迭代优化算法的收敛性至关重要【索引1,Qsgd: Communicationefficient sgd via gradient quantization and encoding,2017,Advances in neural information processing systems】。
异常值缓解。量化神经网络时的一个关键问题是网络权重、激活值和梯度中存在大的异常值【索引16,LLM.int8(): 8-bit matrix multiplication for transformers at scale,2022,arXiv】。缓解这类异常值的一种标准方法是使用Hadamard变换【索引40,Distributed mean estimation with limited communication,2017,International conference on machine learning;索引9,Quip: 2-bit quantization of large language models with guarantees,2023,arXiv;索引3,Halo: Hadamard-assisted lower-precision optimization for llms,2025,arXiv;索引2,Quarot: Outlier-free 4-bit inference in rotated llms,2024,arXiv;索引44,Training llms with mxfp4,2025,arXiv】:给定一个向量$x \in R^d$,其Hadamard变换$h(x)$定义为$h(x) = H_d x$,其中$H_d \in R^{d \times d}$是归一化的Hadamard矩阵,其元素来自$\{±1\}$。Hadamard矩阵具有递归结构$H_d = \frac{1}{\sqrt{2}} H_2 \otimes H_{d/2}$,这使得当d是2的幂时可以进行高效计算【索引20,Unified matrix treatment of the fast walsh-hadamard transform,1976,IEEE Transactions on Computers】。已有为GPU优化的快速Walsh-Hadamard变换(FWHT)实现【索引15,Fast hadamard transform in cuda, with a pytorch interface,2024,Dao-AILab;索引41,Hadacore: Accelerating large language models with fast hadamard transforms,2024,PyTorch Team】。当d不是2的幂时,输入向量x通常被零填充到下一个2的幂,或使用分组Hadamard变换,即将x分割成大小相等的块(每个块的长度都是2的幂),然后对每个块独立应用Hadamard变换。
Blackwell架构支持。NVIDIA的第五代Tensor Cores在Blackwell架构【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】中提供了原生的4位浮点数执行能力。这些核心支持不同的块缩放格式,如MXFP4【索引33,Ocp microscaling formats (mx) specification version 1.0,2023,Open Compute Project】和NVFP4【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】,其峰值吞吐量大约是FP8/FP6的两倍,单个B200 GPU的密集FP4计算峰值可达18 PFLOPS【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】。有趣的是,本文研究表明,目前MXFP4是唯一在Blackwell上支持所有前向和反向低精度乘法所需布局的微缩放格式【索引42,CUTLASS,2025,http://github.com/NVIDIA/cutlass Cores在硬件中处理所需的即时重缩放,无需在CUDA层面进行基于软件的重缩放。更多细节在4.4节提供。
LLM预训练。本文在Llama-2【索引43,Llama 2: Open foundation and fine-tuned chat models,2023,arXiv】架构上预训练Transformer模型【索引48,Attention is all you need,2023,arXiv】,模型参数范围为3000万、5000万、1亿、2亿(非嵌入参数),数据与参数比率从25倍(接近计算最优【索引25,Training compute-optimal large language models,2022,arXiv】)到800倍(极端数据饱和)。此外,还选择性地将模型大小扩展到约70亿参数以验证训练稳定性。所有模型均在C4【索引17,Documenting large webtext corpora: A case study on the colossal clean crawled corpus,2021,arXiv】数据集的训练集上训练,并以C4验证集损失作为主要指标。使用AdamW优化器【索引30,Decoupled weight decay regularization,2019,arXiv】,权重衰减为0.1,梯度裁剪为1.0,学习率预热10%并采用余弦调度。本文为一个小型未量化基线模型确定了最优学习率,然后根据非嵌入参数数量的反比进行缩放,并将其用于所有评估的量化方案。所有超参数在附录A.1中呈现。
A2 方法细节
4.1 成分1:通过诱导缩放定律比较量化训练方法
大语言模型的可预测伸缩性。大语言模型(LLM)在模型大小和数据量上跨越多个数量级时表现出的可预测伸缩能力,是当前人工智能领域扩展格局的基石【索引26,Scaling laws for neural language models,2020,arXiv】。数学上,这表示预期损失是模型和数据参数的函数,通常以参数化函数的形式描述。该函数可通过一组训练运行进行拟合,然后用于确定最优的计算训练方案【索引25,Training compute-optimal large language models,2022,arXiv】或推断模型性能【索引19,The llama 3 herd of models,2024,arXiv】。
引入精度维度的缩放定律。本文研究了评估损失与前向和反向传播所用精度(分别表示为$P_{forward}$和$P_{backward}$)相关的缩放定律。为此,作者提出了一种具有以下函数形式的缩放定律:
$$L(N, D, P_{forward}, P_{backward}) = \left( \frac{A}{(N \cdot \text{eff}_N(P_{forward}))^\alpha} + \frac{B}{(D \cdot \text{eff}_D(P_{backward}))^\beta} \right)^\gamma + E$$其中A, B, α, β, γ是描述损失随模型参数数量N和训练语料库大小D普遍伸缩的常数。
关键新增参数:参数效率和数据效率。此公式的关键新增部分是拟合参数effN($P_{forward}$)和effD($P_{backward}$)。effN($P_{forward}$)代表前向传播所用精度$P_{forward}$的参数效率,而effD($P_{backward}$)代表在可能不同的精度$P_{backward}$下进行反向传播的“数据效率”。(这两个因子自然都在区间(0, 1]内,当使用全精度时达到值1)。具体来说,该参数化假设前向传播精度的影响主要体现在可训练参数上,即将精度降低到$P_{forward}$会使模型的“有效”参数数量降低为$N \cdot effN(P_{forward}) \le N$。这遵循了将前向传播量化效应建模为参数数量的乘法因子的普遍趋势【索引23,Scaling laws for sparsely-connected foundation models,2024,International Conference on Learning Representations;索引28,Scaling laws for precision,2024,arXiv;索引24,Compression scaling laws:unifying sparsity and quantization,2025,arXiv;索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】。对于数据项,作者假设降低反向传播精度主要影响数据项D,因此为了达到相同的损失,实际上需要更多的数据,具体来说是需要增加一个因子$1/effD(P_{backward})$。这是作者提出的一种新颖的建模反向传播量化的方式,与优化理论结果【索引1,Qsgd: Communicationefficient sgd via gradient quantization and encoding,2017,Advances in neural information processing systems】以及观察到的性能差距(见图1(a))相一致。作者在附录A.2中提供了实验数据来证明这些假设,并与替代的缩放定律【索引28,Scaling laws for precision,2024,arXiv】进行了比较。
使用缩放定律作为比较标准。实验上,作者观察到不同的量化训练方法,例如STE【索引6,Estimating or propagating gradients through stochastic neurons for conditional computation,2013,arXiv】与QuEST【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】,会诱导出不同的缩放定律,特别是不同的效率参数。通常,缩放定律用于在不同参数和数据规模下推断模型性能,而本文提出使用缩放定律来比较不同的训练方法。具体而言,如果量化训练方法A同时提供了比方法B更高的参数效率effN和更高的数据效率effD,则认为方法A优于方法B。
4.2 成分2:混合精度引导了推理与训练的权衡
基于缩放定律的最优精度选择。上述缩放定律表明,给定一组缩放参数和我们希望模型达到的目标损失,我们可以直接求解出能够匹配该损失的“最优”前向和反向精度。然而,正如Sardana等人【索引38,Beyond chinchilla-optimal: Accounting for inference in language model scaling laws,2025,arXiv】所指出的,在实践中,我们往往更看重推理成本而非训练成本,这在确定“最优”训练精度时可能会导致不同的结果。由于推理延迟仅取决于前向传播(约占训练计算量的33%),而反向传播消耗剩余的约66%,因此可能需要分别分析这些权衡。
指导原则。具体来说,我们可以陈述一组简单的指导原则:
* 前向传播。低精度在前向传播中带来了一个权衡:参数效率降低,但推理速度提高。例如,我们可以训练一个参数量N更大的模型,但将其前向传播量化到更低精度,从而获得更好的权衡。因此,应选择$P_{forward}$来优化这一权衡。
* 反向传播。类似地,通过量化反向传播带来的训练加速可以抵消数据效率effD的降低:我们可以在相同的计算预算下,用更多数据来训练一个量化程度更重的模型。因此,应选择$P_{backward}$来优化这一权衡。
与先前工作的对比。这与之前的工作形成对比,后者通常要求低精度不带来任何精度损失(例如,Chmiel等人【索引11,Accurate neural training with 4-bit matrix multiplications at standard formats,2024,arXiv】)。这种要求不必要地将这些权衡简化为仅选择最快的无损精度。本文认为,缩放定律分析能够提供一种更精细的方法,以决定“最优”的前向和反向精度组合。
示例速度模型。为了说明这一点,我们假设一个与硬件无关的“按位运算”(BOPS)模型,该模型指出加速比与数据类型位宽成反比。表1列出了相对于FP8基线的加速比:
评估有效损失。然后,给定一个前向传播计算预算$N_{max}$和一个训练预算$N_{max}D_{max}$,有效损失将由以下公式给出:
$Loss(N_{max} \cdot sp_{fw}, D_{max} \cdot sp_{tr} / sp_{fw}, P_{fwd}, P_{bwd})$,我们使用公式(1)中的缩放定律进行评估,这导致了图1(a)中的拟合结果。可以看出,$sp_{fw}$和$sp_{tr}$如何作为乘法因子作用于effN和effD,并直接抵消了次优的参数和数据效率。
最优区域图解。图1(b)–(c)展示了最优区域:具体来说,它告诉我们对于哪些模型大小(Y轴)和相应的相对训练计算量(X轴),FP4相对于FP8是最优的(红色区域 vs. 橙色区域)。绿色区域是使用我们的MXFP4实现进行训练将是该指标下最优的区域。在图4中,我们证明了我们用于比较的验证损失与下游性能是一致的,这意味着最优性也同样传导到了下游任务。
总结。总而言之,成分2指出,低精度的影响应在计算预算下进行分析;然后,通过拟合缩放定律,可以揭示在何种情况下,给定的精度对于任一传播过程是最佳选择。
Table 1: 相对于FP8基线,前向(spfw)、反向(spbw)的加速比;sptr是spfw和spbw的加权调和平均数,权重分别为1/3(前向)和2/3(反向)。
4.3 成分3:最小化前向传播误差和无偏梯度估计
确定最佳量化训练方法。以上两个成分应能帮助我们从现有方法中确定“最佳”的量化训练方法,重点关注硬件支持的MXFP4【索引33,Ocp microscaling formats (mx) specification version 1.0,2023,Open Compute Project】格式。
前向传播量化。如第2节所述,现有的QAT(仅前向)方法可分为“噪声注入”法【索引5,Uniq: Uniform noise injection for non-uniform quantization of neural networks,2021,ACM Transactions on Computer Systems (TOCS)】和“误差最小化”法,例如【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】。根据上述讨论(成分1和2),我们关注前向传播,旨在寻找能最大化参数效率因子effN的方法。为此,我们实现了四种标准的QAT方案:1) 随机舍入(SR)与标准的每组AbsMax归一化【索引44,Training llms with mxfp4,2025,arXiv】;2) 普通的四舍五入(RTN)量化与每组AbsMax归一化;3) 可学习尺度裁剪(LSQ)与RTN量化【索引18,Learned step size quantization,2019,arXiv;索引53,Training Transformers with 4-bit Integers,2023,Advances in Neural Information Processing Systems (NeurIPS)】;4) Hadamard归一化后接基于RMSE的裁剪(QuEST)【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】。为公平起见,我们在量化前对每种方案的权重和激活值都应用了Hadamard变换。我们按照4.1节的方法比较这些方案:使用每种技术训练模型,进行缩放定律拟合,并记录它们产生的effN因子。作为补充信息,我们还展示了在拟合随机高斯数据时表示的均方误差(MSE)。结果见表2的第一行/列。
前向传播结果分析。表2中的结果显示,在所有现有方法中,QuEST具有最佳的参数效率effN。此外,effN似乎与MSE高度相关,正如Panferov等人【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv;索引35,Unified scaling laws for compressed representations,2025,arXiv】所建议的。此外,这些结果与Chmiel等人【索引11,Accurate neural training with 4-bit matrix multiplications at standard formats,2024,arXiv】的分析一致,他们确定对于前向传播,确定性的RTN总是优于随机舍入。
反向传播:一种新颖的误差-偏差权衡。上述发现在反向传播量化中并不适用,因为优化理论表明,无偏的梯度估计对于收敛至关重要【索引1,Qsgd: Communicationefficient sgd via gradient quantization and encoding,2017,Advances in neural information processing systems】。这导致了我们在前向传播中可以获得的误差最小化与给定方法在反向传播中引入的偏差之间的权衡。我们通过一种新颖的对不同量化方法间的梯度对齐性分析来研究这种权衡。
梯度偏差研究。为了研究梯度偏差,我们遵循了【索引45,Drive: One-bit distributed mean estimation,2021,arXiv;索引46,Eden: Communication-efficient and robust distributed mean estimation for federated learning,2022,arXiv】的分析,他们研究了带随机旋转的RTN量化,并用随机Hadamard变换(表示为H)来近似。他们表明,虽然RHT使量化在方向上无偏,但它在幅度上增加了偏差。为了解决这个问题,他们提出了一种方法,通过以下依赖于输入(X)和随机性(ξ)的组级重缩放因子S,使得后RHT向量的RTN投影(表示为Q)无偏:
$$\mathbb{E}_{\xi}[Q(X, \xi)]=X \text { if } Q(X, \xi)=S \cdot \operatorname{RTN}(\widehat{H}(X, \xi)), \text { where } S:=\frac{\langle X, X\rangle}{\langle\widehat{H}(X, \xi), \operatorname{RTN}(\widehat{H}(X, \xi))\rangle}$$实用性与分析。不幸的是,他们的重缩放方法与MXFP4格式的粗粒度组级缩放不兼容,因此我们无法在实践中使用。然而,我们仍然可以利用他们的方法,通过简单地研究不同量化器对应的$1 − E [1/S]$的期望值来衡量其失准程度,我们称之为投影幅度失准。这个因子与不同方案的MSE一起呈现在表2中。关注随机舍入(SR)与带AbsMax的四舍五入(RTN),可以看出SR以高误差换取了完美的对齐。
连接对齐度量与训练动态。为了将这些量与训练动态联系起来,我们分析了失准和误差对一个30M参数Llama模型反向传播量化的累积效应。在图2(a)和(c)中,我们绘制了层间激活梯度的对齐度量——余弦相似度和投影幅度失准——作为反向传播“深度”的函数。我们可以再次观察到相似度和幅度失准之间的权衡。最后,图2(c)将这些量与最终模型质量(与全精度模型的损失差距)在数据与参数比增加的情况下联系起来。
反向传播结果分析。有趣的是,我们观察到余弦相似度(以及引申的MSE)对初始收敛和较短的训练运行有很大影响,而投影幅度失准对较长的运行有更大影响。具体来说,虽然RTN反向传播量化对于较短的训练可能更可取,但随机舍入(SR)对于数据更饱和的模型表现始终更好。在此设置中,拐点大约在D/N = 400的数据与参数比率处。
总结。我们的分析勾勒出了一个新的权衡:前向传播的参数效率(等同于量化MSE)与反向传播的数据效率(我们等同于新的失准度量)。接下来,我们将采用一种“两全其美”的方法,旨在执行一个最小化MSE的前向传播(基于QuEST【索引34,Quest: Stable training of llms with 1-bit weights and activations,2025,arXiv】)和一个无偏的反向传播(基于随机舍入【索引44,Training llms with mxfp4,2025,arXiv】)。我们接下来要解决的新挑战是这种方法的一个极其高效的、面向GPU的实现。
4.4 成分4:为精确量化训练提供快速GPU支持
Quartet概述。我们将之前的讨论整合到算法1中,该算法旨在执行精确的训练,同时将线性层的所有三个矩阵乘法都在低精度下执行。前向传播应用一个固定的Hadamard变换$H_g$(块大小g等于量化组大小)和QuEST投影到低精度,然后用MXFP4核将得到的张量相乘。反向传播用一个相同的块级随机Hadamard变换$H_g$来去相关化相乘的张量,应用无偏的随机舍入(SR)到MXFP4,在MXFP4中执行两个梯度GEMM,进行重缩放以补偿SR的范围匹配,应用QuEST掩码($M_x, M_w$)并反转Hadamard变换$H_g$。
成本和格式专用化。上述流程的主要额外成本来自Hadamard旋转及其逆变换:具体来说,相比标准训练增加了两次Hadamard/逆Hadamard变换。我们的关键观察是,由于MXFP4已经将32个连续的权重(在1D中)分组并共享尺度,我们可以也应该以相同的组大小应用Hadamard旋转及其逆变换。通过快速的Hadamard实现,理论成本为$O(g \log g)$——对于$g \le 256$来说,与GEMM相比可以忽略不计。
GPU核支持。尽管上述蓝图看起来简单,但在Blackwell GPU上高效实现它——为了利用快速的MXFP4支持——是极具挑战性的。举例来说,直接实现上述模式会比FP16未量化训练还要慢,更不用说优化的FP8了。我们的快速实现建立在CUTLASS 3.9【索引42,CUTLASS,2025,http://github.com/NVIDIA/cutlass。
阶段1:融合量化相关操作。首先,我们观察到,由于组大小较小,Hadamard变换可以实现为相应输入矩阵和一个固定的32×32 Hadamard矩阵之间的直接GEMM(见第3节),产生FP32输出,并存储在GPU共享内存(SMEM)中。这使我们能够利用CUTLASS的多级分块模板来优化数据移动,从而高效地实现Hadamard操作。所有后续操作都通过一个自定义的CUTLASS epilogue集成,该epilogue利用先前存储在内存层次结构较高层的中间结果,并在寄存器文件(RF)中本地操作。在此阶段,利用Blackwell的新硬件支持,使用PTX指令将FP32值下转换为FP4(E2M1)。为了构建最终的MXFP4格式,我们计算形状为1×32的缩放因子。这些尺度以8位E8M0格式表示。最后,计算裁剪掩码,并将三个结果张量(值、尺度和掩码)写入全局内存(GMEM)。在整个过程中,数据存储都经过优化,以使用尽可能宽的内存指令。
阶段2:专用GEMM核。Blackwell引入了tcgen05.mma指令,该指令原生支持带缩放因子的矩阵乘法,形式为$D = C + (A \times SFA) \cdot (B \times SFB)$。这些缩放因子沿GEMM的内(K)维度应用。对于MXFP类型,矩阵A和B沿K维度的每32个元素共享一个相应的缩放因子。这意味着一个M × K的矩阵A与一个大小为M × ⌈K/32⌉的尺度矩阵SFA相关联。我们的专用核基于CUTLASS的块缩放GEMM,用于窄精度。作为此实现的一部分,我们还包含了必要的函数,以重组在阶段1中生成的缩放因子,使其与该架构所需的布局对齐【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】。
实现贡献。据我们所知,我们的实现是第一个在Blackwell架构上高效支持微缩放格式量化相关操作的实现。我们将其作为“QuTLASS”的一部分发布,这是一个开源库,可在此处访问。
A4 实验环境与实验结果
实验环境
- 数据集:使用C4数据集【索引17,Documenting large webtext corpora: A case study on the colossal clean crawled corpus,2021,arXiv】的训练集进行预训练,并在其验证集上报告损失作为主要评估指标。下游任务评估使用了HellaSwag【索引57,Hellaswag: Can a machine really finish your sentence?,2019,arXiv】,WinoGrande【索引37,Winogrande: An adversarial winograd schema challenge at scale,2019,arXiv】和ARC-easy【索引13,Think you have solved question answering? try arc, the ai2 reasoning challenge,2018,arXiv】。
- 模型架构:采用Llama-2【索引43,Llama 2: Open foundation and fine-tuned chat models,2023,arXiv】架构的Transformer模型,非嵌入参数量覆盖3000万至70亿。具体参数如层数、嵌入维度、注意力头数见附录A.1。
- 硬件配置:预训练实验在包含8块NVIDIA H100 GPU的数据中心级服务器上进行。速度测试在单块消费级NVIDIA RTX 5090 GPU上完成。
- 软件配置:使用AdamW优化器【索引30,Decoupled weight decay regularization,2019,arXiv】,学习率采用10%预热的余弦衰减策略。实现基于CUTLASS 3.9【索引42,CUTLASS,2025,http://github.com/NVIDIA/cutlass。
实验结果
缩放定律拟合:
实验首先按照4.1节描述的方法,通过一个两阶段过程拟合了公式1中的缩放定律。第一阶段,在FP8基线运行上拟合通用参数A, α, B, β, E, γ。第二阶段,对每个量化方案单独拟合其参数效率effN和数据效率effD。
精度比较:
本文将Quartet与四种最近的、在4位精度下进行全量化(前向和反向)训练的方法进行了比较:
1. LUQ【索引11,Accurate neural training with 4-bit matrix multiplications at standard formats,2024,arXiv】:采用无偏量化,将4位权重/激活值与随机下溢和对数随机舍入配对。
2. HALO【索引3,Halo: Hadamard-assisted lower-precision optimization for llms,2025,arXiv】:使用Hadamard旋转来缓解异常值,评估其最精确的FP4 HALO-2设置。
3. Jetfire【索引55,Jetfire: Efficient and Accurate Transformer Pretraining with INT8 Data Flow and Per-Block Quantization,2024,Proceedings of the 41st International Conference on Machine Learning (ICML)】:在32×32块内进行量化,本文将其从INT8适配到FP4。
4. LSS【索引53,Training Transformers with 4-bit Integers,2023,Advances in Neural Information Processing Systems (NeurIPS)】:用于INT4训练,结合了基于Hadamard的前向传播和“杠杆分数”采样的INT4梯度。
精度讨论 (见表3):
如表3所示,在所有tokens-to-parameters比率下,Quartet都获得了最低的验证损失,并且优势显著。在100倍的tokens-per-parameter比率下,Quartet相比最强的基线LUQ-INT4,相对损失降低了10%。当移植到FP4时,Jetfire和HALO出现了严重的性能下降且不稳定。LSS仅在短时训练中具有竞争力,在超过50倍数据量后便会发散。总体而言,Quartet展现出显著更高的参数效率(effN=0.64)和数据效率(effD=0.94),这意味着达到相同损失,它大致需要少15%的参数和少5倍的数据。
稳定性和下游任务性能 (见图3(c)和图4):
图3(c)展示了Quartet在训练大两个数量级的7B参数模型时的稳定性。此外,在100M到1.6B参数的Llama模型上,下游基准测试(HellaSwag, WinoGrande, ARC-easy)的结果与C4验证损失表现出良好的一致性,证明了验证损失能够有效反映模型的泛化能力(图4)。
加速比结果 (见图3(a,b)):
在NVIDIA RTX 5090 GPU上评估了实现的效率。结果显示(图3a,b),Quartet的加速效果随模型尺寸(算术强度)增加而提升。
* 前向传播(推理):相较于FP8最高加速2倍,相较于BF16最高加速4倍。
* 反向传播:相较于FP8最高加速1.5倍,相较于BF16最高加速2.6倍。
* 整体训练:相较于FP8,整体训练加速最高可达约1.6倍;相较于BF16,最高可达2.9倍。
A5 结论
本文提供了一套用于建模、比较和设计大语言模型全量化训练方案的指导原则。遵循这些原则,本文提出了一种新的、最先进的全MXFP4训练算法——Quartet。
局限性:Quartet目前是为特定的数据类型(MXFP4)和计算架构(NVIDIA Blackwell)设计的。其某些方面依赖于硬件对特定操作(如随机舍入)的支持,这在其他格式中可能缺失。
未来工作:计划将该方法推广到其他替代格式,并探索其在更大规模的分布式模型执行中的应用。
A6 附录
A.1 训练超参数
表4列出了模型特定的超参数。表5列出了所有实验中共享的超参数。
A.2 缩放定律拟合
分阶段拟合。我们分两个阶段拟合缩放定律:
阶段1:与先前的工作【索引8,Distillation scaling laws,2025,arXiv】相同,我们对BF16基线运行拟合了以下形式的未量化缩放定律:
这些运行的模型参数N ∈ [30M, 50M, 100M, 200M],数据/参数比D/N ∈ [25, 50, 100, 200, 400, 800](见图1(a)),使用Huber损失(δ = 10⁻⁴)对L的对数进行拟合。表6显示了拟合结果。
阶段2:使用阶段1中固定的拟合参数,我们使用相同的损失函数拟合额外的effN和effD参数。对于4.2节中独立比较的方法,我们分别为仅前向和仅反向量化独立拟合了effN和effD。对于第5节中的端到端4位比较,我们对表3中的设置联合拟合了这些参数。
替代形式。我们额外对固定γ = 1【索引25,Training compute-optimal large language models,2022,arXiv】和β = 1【索引26,Scaling laws for neural language models,2020,arXiv】的缩放定律形式进行了拟合。拟合结果与主要使用的Busbridge等人【索引8,Distillation scaling laws,2025,arXiv】的模型一起呈现在图5中。
A.3 性能分解
性能分解分析。图6展示了Llama-7B模型中三种线性层形状的运行时构成细分,以MXFP4前向传播为例。每个子图显示了三个关键内核阶段所花费的总运行时间的百分比:矩阵乘法、量化相关操作以及为mma指令【索引32,Nvidia blackwell architecture technical brief,2024,NVIDIA Corporation】重排缩放因子。
核配置比较。该图比较了三种内核配置。左侧子图显示了我们用于量化相关操作的融合内核,使用了基本的32×32线程块瓦片大小。中间子图将此瓦片大小增加到128×32,从而实现了更高效的量化阶段。右侧子图包含一个自定义的Triton内核,通过优化MXFP重排阶段进一步提高了性能。所有结果都归一化为100%。如图所示,调整量化内核显著减少了在量化阶段所花费的时间比例——尤其对于大型矩阵形状。增加线程块瓦片大小会导致每个块有更多的活动线程束,从而增强算术强度并实现更好的延迟隐藏。在基于CUTLASS的实现中,这种变化影响了多级分块策略(线程块、线程束和指令级分块),该策略旨在优化通过共享内存和寄存器的数据移动【索引42,CUTLASS,2025,http://github.com/NVIDIA/cutlass。
A.4 端到端预填充加速
端到端预填充加速。图7展示了在固定序列长度为256的7B参数模型上,MXFP4相对于FP8的推理预填充(prefill)加速比随批次大小变化的函数。结果表明,在所有批次大小下,使用MXFP4都实现了性能的持续提升,加速比逐渐增加,在批次大小为128时达到峰值1.41倍,随后趋于平稳。
A.5 训后量化结果
与训后量化(PTQ)的比较。我们将应用训后量化(PTQ)的结果与在最大的7B模型上使用MXFP4格式的QUARTET进行了比较。对于PTQ基线,我们与QUAROT【索引2,Quarot: Outlier-free 4-bit inference in rotated llms,2024,arXiv】进行了评估,其中权重使用GPTQ【索引22,GPTQ: Accurate post-training compression for generative pretrained transformers,2022,arXiv】进行量化。为确保公平比较,我们对原始QUAROT方法进行了两项关键修改:
1. 注意力模块:我们移除了在线Hadamard变换的使用,而是在v_proj层的输出维度和out_proj层的输入维度上应用一个大小为128的固定Hadamard变换。由于我们在MXFP4格式中使用的组大小为32,此优化通过消除每个头的在线Hadamard计算来加速整个过程,而不影响准确性。
2. MLP下投影:对于MLP中具有非2的幂次维度的down_projection层,我们使用能够整除MLP中间维度的最大2的幂次大小来应用分组Hadamard变换。
结果分析。表7展示了PTQ方案(QuaRot)与QUARTET的比较。当应用于同一模型时,QUARTET的困惑度(PPL)比QuaRot低0.42点。值得注意的是,QUARTET比标准的QAT方法更高效,因为它同时量化了前向和反向传播。
A.6 计算资源
预训练实验在配备8块H100 NVIDIA GPU的数据中心级机器上进行,总计算量约为6000 GPU小时。尽管大多数实验不需要如此复杂的设置,但我们发现7B预训练实验对DRAM要求特别高,需要此类特定硬件。
加速比结果是在一台消费级NVIDIA RTX5090 GPU上获得的,总运行时间不到1小时。
💬 评论讨论
欢迎在这里分享您的想法和见解!