Quartet II：通过改进的无偏梯度估计在NVFP4中进行精确的LLM预训练 Quartet II: Accurate LLM Pre-Training in NVFP4 by Improved Unbiased Gradient Estimation
通过MS-EDEN无偏量化技术，在NVFP4格式下实现高效且精确的LLM预训练。
前置知识
NVFP4格式
NVFP4是NVIDIA Blackwell GPU支持的一种4位微缩放浮点格式。它将值映射到E2M1浮点编码（2位指数，1位尾数），并配备两级缩放因子：每16个值一个E4M3缩放因子，以及每个张量一个FP32全局缩放因子用于范围扩展。这种格式在保持动态范围的同时，允许张量核心加速矩阵乘法，相比FP8可提供2-4倍的吞吐量提升。
理解NVFP4格式是理解本文量化方案的基础，因为它定义了量化操作的约束条件和硬件支持特性。
微缩放格式（Microscaling Formats）
微缩放格式是一种低精度数值表示方法，其中每个小数据块（如16个元素）共享一个高精度缩放因子。这种设计在保持较低存储精度的同时，通过共享缩放因子来保留数据的动态范围。NVFP4和MXFP4都是微缩放格式的实例，它们通过缩放因子将4位浮点值的范围扩展到适合深度学习训练的程度。
微缩放格式是本文量化方案操作的核心数据结构，理解其工作原理对于把握MS-EDEN的创新点至关重要。
无偏量化（Unbiased Quantization）
无偏量化是指量化操作的期望值等于原始值，即对于随机性ω，有E[Q(x, ω)] = x。在深度学习训练中，无偏梯度估计确保优化方向在统计意义上是正确的，避免系统性偏差累积导致训练发散。随机舍入（SR）是实现无偏量化的常用方法，但它会显著增加方差。
无偏性是本文方法的核心保证，它确保了量化训练的长期稳定性和收敛性。
随机舍入（Stochastic Rounding, SR）
随机舍入是一种概率舍入方法，其中数值以与其到最近量化级别的距离成比例的概率向上或向下舍入。例如，值0.3有70%的概率舍入到0，30%的概率舍入到1。这种方法在期望上保持无偏性，但在低精度下会引入显著方差，因为每个元素的舍入决策都是独立的。
随机舍入是现有NVFP4训练方法的主要无偏技术，本文通过MS-EDEN提供了更优的替代方案。
随机化Hadamard变换（Randomized Hadamard Transform, RHT）
RHT是一种正交变换，通过将数据旋转到随机方向来平滑异常值分布。在量化中，RHT可以减少离群值对量化精度的影响，因为旋转后的数据通常具有更均匀的幅度分布。RHT在分组（如128个元素）上应用，使用伪随机种子确保可重复性。
RHT是MS-EDEN实现无偏性的关键组件，它通过旋转操作确保量化误差在统计上相互抵消。
EDEN重缩放
EDEN是一种分布式优化中的无偏量化技术，它结合随机化旋转和校正缩放因子来实现无偏性。具体来说，它计算校正因子S = ⟨x, x⟩ / ⟨RHT(x, ω), Q(RHT(x, ω))⟩，使得E[RHT^{-1}(S·Q(RHT(x, ω)), ω)] = x。这种方法比随机舍入具有更低的均方误差，但需要高精度存储缩放因子。
EDEN是MS-EDEN的理论基础，理解其原理有助于把握本文如何将其扩展到微缩放格式。
Four Over Six（4/6）启发式
Four Over Six是一种NVFP4量化算法改进，它为每个数据块评估两个潜在的缩放因子（4.0和6.0），并选择均方误差更低的那个。这种方法可以改善前向传播的量化精度，但如果直接与随机舍入结合用于反向传播，会引入偏差，因为选择最低MSE缩放因子的行为本身就会引入系统性偏差。
4/6启发式是本文前向传播方案的重要组成部分，理解其特性和限制有助于把握Quartet II的设计决策。
研究动机
训练现代大型语言模型的计算成本呈指数级增长，对领域可持续性构成挑战。预训练过程由密集矩阵乘法（GEMM）主导，这些操作占用了大部分计算资源。降低这些GEMM的精度是降低端到端训练成本最直接的手段之一。虽然从FP16/BF16到FP8的混合精度训练已经成熟，但向4位微缩放格式（如NVFP4和MXFP4）的过渡面临关键挑战：在保持FP16/FP8质量优化的同时，在4位精度下执行大部分操作。现有量化训练方法（如随机舍入）虽然能保持无偏性，但会显著增加方差，导致精度损失。具体来说，随机舍入的均方误差比舍入到最近（RTN）高约2.5倍（如表1所示，SR的MSE为23.5×10^{-3}，而RTN为9.0×10^{-3}）。这种精度损失在长期预训练中会累积，导致验证损失增加和收敛不稳定。
本文的目标是本文的具体目标是改进NVFP4格式下量化训练的现状，通过重新审视无偏梯度估计问题，实现更精确的量化训练。具体来说，作者希望：1) 提出一种新的无偏量化原语，显著降低量化误差同时保持无偏性；2) 构建一个完整的NVFP4线性层计算图，结合改进的前向传播和反向传播量化；3) 在端到端LLM预训练中验证方法的有效性；4) 提供高效的内核实现，使方法在NVIDIA Blackwell GPU上实用。
与已有工作不同的是，本文的独特切入角度是重新考虑反向传播中的无偏量化策略。现有方法都依赖元素级随机舍入来实现无偏性，但这种方法在低精度下会引入显著方差。作者借鉴分布式优化中的EDEN技术，提出将随机性从单个FP4值转移到微缩放因子上。这种转移是通过MS-EDEN实现的：它首先对数据应用随机化旋转，然后使用RTN量化（比SR误差更低），最后通过在FP8缩放因子上应用随机舍入来实现无偏性。这种方法的关键洞察是：NVFP4格式中FP8缩放因子的精度（E4M3）虽然不足以直接存储EDEN的校正因子，但可以通过随机舍入来无偏地表示这些因子。这种设计在保持无偏性的同时，将量化误差降低了2倍以上（如表1所示，MS-EDEN的MSE为9.4×10^{-3}，而SR为23.5×10^{-3}）。
核心方法
Quartet II的方法整体思路是结合前向传播和反向传播的优化策略，构建一个完整的NVFP4线性层计算方案。前向传播旨在最大化表示容量，通过最小化激活和权重的量化误差（通常通过均方误差衡量）。反向传播对偏差特别敏感，因此需要无偏（或仔细控制的）梯度量化来确保稳定收敛。具体来说，前向传播使用RTN量化结合Four Over Six启发式选择最优缩放因子，这可以改善量化精度而不引入偏差。反向传播则使用MS-EDEN进行无偏量化，它通过在微缩放因子上应用随机舍入来实现无偏性，同时保持较低的量化误差。整个方案的关键创新在于将前向传播的表示容量优化与反向传播的无偏性保证结合起来，实现了比现有方法更精确的量化训练。
本文的核心创新点是MS-EDEN，这是一种针对微缩放格式的新无偏量化技术。与现有方法的本质区别在于：现有方法（如随机舍入）在单个FP4值上引入随机性来实现无偏性，但这会显著增加方差；而MS-EDEN将随机性转移到微缩放因子（FP8缩放因子）上。具体来说，MS-EDEN首先对数据应用随机化Hadamard变换（RHT），然后使用RTN量化（比SR误差更低），最后通过在FP8缩放因子上应用随机舍入来实现无偏性。这种设计的关键优势是：1) RTN量化比SR具有更低的均方误差（表1显示RTN的MSE为9.0×10^{-3}，而SR为23.5×10^{-3}）；2) FP8缩放因子上的随机舍入引入的方差比4位量化本身小一个数量级；3) 通过RHT旋转，量化误差在统计上相互抵消，确保无偏性。此外，MS-EDEN与Four Over Six启发式在前向传播中协同工作，进一步提升了表示容量。
方法步骤详情
Quartet II的方法步骤如下：前向传播阶段：1) 对输入激活X和权重W应用RTN量化到NVFP4格式，使用原生NVFP4缩放（每16个元素一个FP8缩放因子）和每个张量一个FP32缩放因子用于范围扩展；2) 应用Four Over Six启发式，为每个数据块评估缩放因子4.0和6.0，选择均方误差更低的那个；3) 执行NVFP4张量核心矩阵乘法Y = XW^T，保存量化后的权重和激活供反向传播使用。反向传播阶段：1) 生成分组RHT旋转矩阵，使用伪随机种子；2) 将保存的量化权重和激活进行反量化、转置，然后使用MS-EDEN重新量化：首先对当前块应用RHT（128个元素为一组），然后使用RTN量化到NVFP4（允许裁剪），接着计算EDEN校正因子S = ⟨x_RHT, x_RHT⟩ / ⟨x_RHT, x_RTN⟩，最后在FP8缩放因子上应用随机舍入来合并校正因子；3) 对梯度E和E^T也应用MS-EDEN量化；4) 执行NVFP4张量核心矩阵乘法计算权重梯度和输入梯度；5) 旋转操作沿GEMM内维度应用，因此乘积输出不需要进一步处理，旋转自然抵消；6) 将结果馈送到优化器步骤和反向传播中。当GEMM中的两个张量都被量化时，沿内维度在128个元素的组上应用RHT进行异常值平滑。
技术新颖性
Quartet II的技术新颖性体现在多个层面。首先，MS-EDEN是第一个针对微缩放格式的无偏量化技术，它解决了EDEN校正因子无法直接存储在NVFP4粗粒度缩放表示中的问题。通过将校正因子合并到FP8缩放因子中并应用随机舍入，MS-EDEN在保持无偏性的同时实现了更低的量化误差。其次，作者发现Four Over Six启发式如果直接与随机舍入结合用于反向传播会引入偏差，因此将其限制在前向传播中使用，这是一个重要的设计决策。第三，MS-EDEN要求旋转和缩放校正沿GEMM内维度应用，这创造了权重重新量化的需求。但作者通过分析证明，即使需要权重重新量化，MS-EDEN的量化误差仍然低于没有权重重新量化的随机舍入（如表1所示，MS-EDEN的MSE为9.4×10^{-3}，而SR为23.5×10^{-3}）。第四，作者提出了后验范围对齐的硬件感知实现启发式，通过在第一个内核中使用扩展范围代理（E8M3缩放）和在第二个内核中应用EDEN校正，避免了双重加载和旋转，将理论带宽节省提高了约20%。最后，整个Quartet II计算图结合了前向传播的表示容量优化和反向传播的无偏性保证，实现了比现有方法更一致的梯度估计。
实验结果
本文的核心发现通过广泛的消融实验和端到端训练验证。首先，在30M到200M参数的Llama风格模型上，作者系统评估了不同反向传播量化方案。图1显示，MS-EDEN在所有方案中都优于随机舍入（SR），特别是在完全量化的方案(e)中（Q(E)·Q(W^T)^T, Q(E^T)·Q(X^T)^T），MS-EDEN的验证损失增加比SR低至少20%。其次，前向传播量化实验（图2）表明，Four Over Six启发式与原生NVFP4缩放协同工作，将性能差距减少了约一半。第三，完全NVFP4量化实验（图4）显示，Quartet II相对于NVIDIA基线、FourOverSix和TetraJet-v2，一致性地将损失降低了至少20%。在更大规模的Nanochat实验中，560M参数模型在11B tokens上训练，Quartet II将预训练损失差距（相对于BF16）减少了15-25%，从1.56%（NVIDIA）降至1.26%。1.9B参数模型在38B tokens上训练，损失差距从1.92%（NVIDIA）降至1.44%。在硬件加速方面，Quartet II在RTX 5090上实现了超过4倍的线性层加速（即使对于小尺寸），在B200上实现了高达2.5倍的加速（对于22B参数模型）。端到端训练在1.1B参数模型上实现了1.85倍的实际吞吐量提升，在3.3B到11B参数模型上实现了1.48到1.68倍的加速。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| C4验证损失（30M参数模型，800 tokens/参数） | 相对于BF16的损失增加 | 0.003 | NVIDIA: 0.007, TetraJet-v2: 0.005 | 比NVIDIA低57%，比TetraJet-v2低40% |
| C4验证损失（200M参数模型，800 tokens/参数） | 相对于BF16的损失增加 | 0.02 | NVIDIA: 0.06, TetraJet-v2: 0.04 | 比NVIDIA低67%，比TetraJet-v2低50% |
| Nanochat预训练（560M参数，11B tokens） | 相对于BF16的BPB增加 | 1.26% | NVIDIA: 1.56%, 4/6: 1.51%, TetraJet-v2: 1.72% | 比NVIDIA低19%，比4/6低17%，比TetraJet-v2低27% |
| Nanochat预训练（1.9B参数，38B tokens） | 相对于BF16的BPB增加 | 1.44% | NVIDIA: 1.92%, 4/6: 1.76%, TetraJet-v2: 1.72% | 比NVIDIA低25%，比4/6低18%，比TetraJet-v2低16% |
| 线性层加速（RTX 5090，22B参数模型） | 相对于BF16的加速比 | 4.2x | 纯GEMM天花板: ~6x | 达到理论峰值的70% |
| 端到端训练加速（B200，11B参数模型） | 相对于BF16的加速比 | 1.68x | BF16: 1.0x | 68%吞吐量提升 |
局限与改进
本文的局限性主要体现在以下几个方面：首先，实验验证主要集中在LLM预训练任务上，使用的是Llama风格架构和C4数据集，方法的泛化性到其他模型架构（如CNN、RNN）和其他任务（如微调、推理）尚未充分验证。其次，硬件支持仅限于NVIDIA Blackwell GPU一代，这限制了方法的适用范围，因为其他硬件平台可能不支持相同的微缩放格式或张量核心指令。第三，虽然MS-EDEN降低了量化误差，但它需要额外的计算开销进行随机化旋转和重新量化，这在小模型上可能抵消精度收益。第四，作者承认从基本统计属性（如高斯源上可测量的更低MSE）来看，方法应该能泛化到其他问题，但这可能需要额外的超参数调整，如随机化旋转的时间和空间粒度。第五，实验中使用的模型规模相对较小（最大1.9B参数），对于当前更大的模型（如70B+参数）的有效性需要进一步验证。最后，虽然方法提供了内核实现，但实际部署中的工程复杂性（如内核融合、内存管理）可能影响实际性能。
独立分析的弱点
尽管Quartet II取得了显著进展，但仍存在一些弱点。首先，方法仅支持NVFP4格式，而其他微缩放格式（如MXFP4）可能需要不同的量化策略。虽然作者提到NVFP4在精度上优于MXFP4，但实际部署中硬件支持可能多样。其次，MS-EDEN依赖随机化Hadamard变换（RHT），这需要额外的计算和内存开销。在小模型上，这种开销可能抵消量化带来的收益，如图6所示，对于800M参数模型，B200上的实际加速有限。第三，权重重新量化的需求增加了反向传播的计算复杂性。虽然作者证明MS-EDEN即使需要权重重新量化也优于SR，但这仍然是额外的成本。第四，Four Over Six启发式仅用于前向传播，因为它如果用于反向传播会引入偏差。这限制了其在反向传播中的潜在收益。第五，内核实现依赖于特定的硬件指令（如mma.m16n8k16），这降低了可移植性。改进方向包括：1) 探索适用于其他微缩放格式的MS-EDEN变体；2) 优化RHT实现以减少计算开销，例如通过内核融合；3) 研究是否可以在不牺牲无偏性的前提下避免权重重新量化；4) 设计适用于反向传播的无偏缩放选择算法；5) 开发更通用的内核实现，支持多种硬件平台。
未来方向
作者提出的未来研究方向包括：1) 将方法扩展到其他问题和模型架构，这可能需要额外的超参数调整，如随机化旋转的时间和空间粒度；2) 探索在更大规模模型（如70B+参数）上的应用，验证方法的可扩展性；3) 优化内核实现以进一步减少量化开销，特别是在小模型和小批量大小情况下。基于本文成果可延伸的方向包括：1) 将MS-EDEN应用于其他低精度训练场景，如INT4训练或混合精度训练；2) 探索MS-EDEN在分布式训练中的潜力，因为EDEN最初就是为分布式优化设计的；3) 研究MS-EDEN与其他训练稳定技术（如梯度裁剪、学习率调度）的协同作用；4) 将方法应用于其他领域，如计算机视觉或强化学习中的低精度训练；5) 开发自适应量化策略，根据训练阶段动态调整量化参数；6) 探索MS-EDEN在推理加速中的应用，虽然本文focus在训练上，但无偏量化技术也可能改善推理效率。
复现评估
本文的复现评估如下：代码已开源在GitHub（https://github.com/IST-DASLab/Quartet-II），这大大降低了复现门槛。数据方面，C4数据集公开可用，Nanochat使用的FineWeb-Edu数据集也是公开的。算力要求方面，实验使用了8块GPU（具体型号未明确说明，但应该是NVIDIA Blackwell GPU），对于大规模训练可能需要更多资源。复现难度中等：虽然代码可用，但需要NVIDIA Blackwell GPU支持，这限制了可复现性。此外，内核实现依赖特定的硬件指令，可能需要调整才能在其他硬件上运行。作者提供了详细的超参数设置（附录B和C），这有助于复现实验结果。总体而言，对于拥有合适硬件的研究团队，复现是可行的；但对于没有Blackwell GPU的团队，复现存在硬件障碍。作者还提供了内核基准测试细节（附录D），这有助于验证性能声称。
论文图表
