🤖 machine learning

Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention

本文介绍了 Full-Stack FP4，这是一个模块化框架，它将 NVFP4 量化从线性投影扩展到了包括稳定的优化器状态、Root/Muon 计算以及混合精度注意力机制在内的领域，在单张 RTX 5090 上实现了接近 BF16 的预训练性能，并获得了显著的显存节省与加速收益。

原作者： Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

发布于 2026-08-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个巨大的数字大脑如何说话、写作和思考。这个“大脑”是一个大语言模型（LLM），为了学习，它必须处理如山一般的海量信息。但问题在于：它学得越多，对计算机内存和电力的渴求就越强烈。这就像试图在一台破旧的小型笔记本电脑上运行一部超高清的大片；屏幕变得模糊，风扇发出尖叫。为了解决这个问题，科学家们一直尝试缩小计算机进行数学运算时使用的数字。他们不再使用极其精确、沉重的数字（就像一张细节丰富的照片），而是转向使用更小、更轻的数字（就像一张像素化的素描）。这被称为“低精度训练”。最近，一种名为“FP4”（4位浮点数）的新型微小数字出现了，它承诺让这些巨大的大脑变得更小、更快。但问题在于：虽然大脑的主体部分可以处理这些微小的数字，但那些“辅助”部分——比如它用来记忆学习内容的内存，以及它用来理解句子的复杂数学运算——在被迫使用如此小的数字时，会不断崩溃或变得混乱。这就像拥有一台赛车引擎，却试图用自行车链条来驱动它；引擎已经准备好了，但机器的其他部分却跟不上。

这篇论文介绍了一个名为 Full-Stack FP4 的巧妙新工具包，它充当了这些数字大脑的“首席机械师”。研究人员意识到，你不能对大脑的每一个部分都使用同样的微小数字规则。有些部分，比如主要连接部分，可以很好地处理微小数字。但其他部分，比如计算机用来记住昨天学到了什么的“记忆”，是非常敏感的，需要特殊的照顾。该团队创建了一个模块化系统，通过为大脑的不同部分使用不同的“配方”来进行运作。对于主要连接部分，他们使用了一种名为 LoRA-SVD 的技巧，这就像是在使用像素化方块的同时，保留最重要细节的高分辨率素描。对于记忆部分，他们发明了一种在将数字挤压进微小盒子之前将其“平滑化”的方法，以免信息丢失。对于帮助大脑专注于正确词汇的数学运算（注意力机制），他们决定将最敏感的部分保持在高清晰度，而让其余部分保持像素化状态。

当他们在 30 亿参数的模型（一个中等规模的数字大脑）上使用 640 亿个单词的训练数据进行测试时，结果令人印象深刻。“Full-Stack FP4”大脑的学习效果几乎与传统的、重型版本完全一致。它们之间的学习得分差异仅为微小的 0.838%，几乎察觉不到。事实上，当他们测试该微小数字版本在回答从未见过的问题时的表现时，尽管它使用的内存要少得多，但它预测正确词汇的能力甚至反而略好一些。在单块强大的显卡（RTX 5090）上，这种新方法使大脑在某些任务中的学习过程快了 2.5 到 2.8 倍，并且比标准高精度方法节省了 38% 到 42% 的内存。

研究人员谨慎地指出，虽然这种方法在 30 亿参数以内的模型上表现出色，但他们尚未证明它是否适用于规模更大的大脑或更长时间的训练。他们还发现，你不能把微小数字一股脑地套用到所有地方；如果你试图在没有这些特殊配方的情况下，强迫大脑最敏感的部分使用最小的数字，学习过程就会崩溃。但就目前而言，“Full-Stack”方法表明，我们可以通过针对大脑不同部分提供其真正所需的特定精度水平，而不是强行采用“一刀切”的方案，来构建更聪明、更快且更便宜的人工智能。

技术摘要：全栈 FP4：结合量化投影、优化器与注意力的稳定 LLM 预训练

1. 问题陈述

NVIDIA Blackwell GPU 的最新进展使得原生 NVFP4（4-bit）Tensor Core 支持成为大语言模型（LLM）预训练的一个实际目标。虽然现有的研究已成功证明了 Transformer 线性投影在 4-bit 下的稳定预训练，但在其他关键训练模块方面仍存在显著差距。目前大多数 NVFP4 配方仍对优化器状态、优化器计算和注意力机制保留较高精度（如 BF16 或 FP8）。

主要挑战在于这些模块具有不同的数值结构，统一的量化规则无法处理：

线性投影： 随着训练的进行，梯度信号逐渐减弱，而量化噪声保持不变，从而降低了持续低精度训练的效能。
AdamW 状态： 二阶矩动量是非负的、重尾的、持久的，并且对分母值非常敏感，直接进行 4-bit 存储容易导致失真。
Root/Muon 优化器： 这些优化器依赖于重复的 Newton–Schulz 矩阵乘法，低精度的迭代误差可能会累积并导致收敛不稳定。
注意力机制： 该算子包含对 Softmax 敏感的路径（例如 $PV$, $dOV^\top$ ），这些路径需要高精度；而其他路径（例如 $QK^\top$ ）则对量化更具鲁棒性。

因此，目前缺乏一个统一的、模块化的框架，能够在不依赖混合精度回退的情况下，实现覆盖整个训练栈（投影、优化器和注意力）的稳定、原生 NVFP4 预训练。

2. 方法论：全栈 FP4

作者提出了 Full-Stack FP4，这是一个由四个独立的、可组合的配方组成的模块化框架，旨在解决每个训练模块特定的数值难题。

2.1 LoRA-SVD 线性投影

为了缓解训练后期梯度信号退化的问题，作者提出了 LoRA-SVD。

机制： 该方法将权重参数化为 $W = W_{res} + \beta L_2 L_1$ 。它在 BF16 中维护一个紧凑的主成分子空间（ $L_2 L_1$ ），同时以 NVFP4 计算主残差（ $W_{res}$ ）。
反向传播： 子空间约束的反向传播使用 Cholesky-QR 约束，以确保 BF16 低秩分支的更新与主成分子空间保持一致，同时残差部分更新 NVFP4 组件。
重对齐： 子空间每 2,048 步进行一次随机 SVD 重对齐，以重置组件，确保 BF16 子空间始终保持最优。

2.2 量化状态 AdamW (Q-AdamW)

为了稳定存储持久且重尾的动量状态：

一阶矩 ( $m_t$ )： 在应用随机舍入（Stochastic Rounding, SR）和自适应 4/6 位量化之前，应用 Hadamard 混合以平滑分布。
二阶矩 ( $v_t$ )： 在 NVFP4 存储前采用三阶段预处理流水线：
1. 平方根： 缩短重尾。
2. 分块均值（Tile Mean）： 将正向公共部分（以 BF16 存储）与近似零均值的残差分离。
3. Hadamard 变换： 将局部能量分散到残差中，使其更具对称性。
重构： 通过反转变换和均值分离来重构状态，然后进行平方。该流水线避免了对基于 SVD 的子空间操作的需求。

2.3 NVFP4 Root 优化器 (Q-Root)

针对 Root 优化器中使用的 Newton–Schulz (NS) 迭代：

直接 NVFP4 执行： 不采用混合精度路径，而是直接在 NVFP4 中执行 NS 矩阵乘法。
稳定性： 通过形状相关系数（针对特定矩阵长宽比优化的 $a, b, c$ ）和 p99 离群值裁剪来实现稳定性。这种方法限制了连续 NS 步骤中的误差放大，而无需辅助的混合精度路径。

2.4 混合精度注意力 (Q-Attn)

为了平衡精度与效率：

精度拆分： 框架将 NVFP4 分配给鲁棒路径（ $QK^\top$ 和 $dS$），并将 BF16 保留给对 Softmax 敏感的路径（$PV$, $P^\top dO$ , 以及 $dOV^\top$ ）。
一致性： 前向和反向传播之间复用量化后的 $Q, K, V$ 视图，以确保张量一致性并避免冗余的内存传输。

3. 核心贡献

LoRA-SVD 投影： 一种可合并的重参数化方法，将仅线性部分的损失差距从 1.40% 降低到 0.61%，保护了 BF16 中的紧凑子空间。
量化状态 AdamW： 首次展示了在集成预训练运行中稳定的 NVFP4 AdamW 动量状态存储，利用变换状态流水线将二阶矩重构误差降低了约 60%。
NVFP4 Root： 通过形状感知系数和裁剪实现的直接低精度 Newton–Schulz 路径，将迭代误差从 ~10% 降低至 ~3%。
可训练混合精度注意力： 一种保护 Softmax 敏感路径为 BF16，同时量化合适的正向和反向路径的策略，实现了一致的 NVFP4 预训练算子。

4. 实验结果

作者在训练了 64B token 的 3B 参数 Transformer 上评估了该框架。

训练损失： BF16 基准损失为 2.267，而 Full-Stack FP4 达到 2.286，相对差距为 0.838%。
下游性能： 零样本评估显示，Full-Stack FP4 的平均困惑度为 26.665（对比 BF16 的 26.675），平均准确率仅比 BF16 基准低 0.10 个百分点。
消融实验：
- 单独使用 LoRA-SVD 将仅线性部分的损失差距从 1.40% 降至 0.61%。
- 这些模块化配方可以独立启用或组合使用，其中“Full-Stack”代表聚合覆盖。
原生硬件效率 (RTX 5090)：
- Root 加速： 在 Root 优化器步骤中，Full-Stack FP4 相比优化的 BF16 实现了 2.50×–2.83× 的加速，相比 NVIDIA Transformer Engine (TE) 实现了 1.83×–2.17× 的加速。
- 显存节省： 与 BF16 相比，该框架减少了 37.9%–42.5% 的 AdamW 峰值显存。
- 时间节省： 在考虑减少的步时以及匹配 BF16 损失所需的 token 略微增加后，估算的等效损失匹配时间节省为 56.9%–61.9%。

5. 重要性与局限性

重要性：
本文证明了对 NVFP4 进行“全栈”预训练是可行的，涵盖了投影、优化器状态、优化器计算和注意力机制。它确立了原生 4-bit 预训练可以在提供显著显存和吞吐量优势的同时，实现与 BF16 相当的性能（损失差距 <1%）。这项工作提供了一套模块化配方，允许从业者循序渐进地采用低精度。

局限性：

规模： 实验仅限于 3B 参数 及 64B token 的模型。论文并未建立其在 7B+ 规模或更长训练周期下的表现。
效率测量： 原生硬件效率是在单个 RTX 5090 上的 四个解码器块 上测量的，而非完整的分布式模型。
估算： 时间节省计算是基于固定 Batch Size 的跨设置估算，未考虑利用释放出的显存增加 Batch Size 带来的潜在吞吐量提升。
算子融合： 部分算子（LoRA 反向、注意力）尚未完全融合，且实现尚未充分利用所有 FlashAttention-4 的优化。
格式特定性： 这些配方和测量结果是针对 NVFP4 特定的，并未声称可迁移至 MXFP4。