全栈4比特训练让大模型预训练更稳定高效
Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention
- 分模块设计4比特方案，覆盖投影、优化器、注意力等关键环节
- 30亿参数训练损失仅比半精度高0.838%，零样本困惑度低0.1个百分点
- 支持单张5090显卡运行，优化器内存降低42.5%，加速2.8倍
近期的四比特预训练研究主要集中于Transformer线性投影的优化，而对优化器状态、计算过程以及低精度注意力前向反向路径的探索仍不充分。本文提出 extbf{Full-Stack FP4}，一个模块化四比特框架，分别针对投影、AdamW状态、Root/Muon计算和注意力设计独立方案。 extbf{LoRA-SVD}在BF16中保护紧凑投影子空间，同时保持完整形状的四比特计算，使仅线性层损失差距从 extbf{1.40 extbackslash%}降至 extbf{0.61 extbackslash%}。有序平方根、分块均值与哈达玛流水线实现稳定的四比特动量存储；依赖形状的系数与截断机制稳定直接的四比特Root迭代；混合精度注意力将对softmax敏感的操作保留在BF16。在使用640亿词元的30亿参数模型预训练中，BF16与Full-Stack FP4的损失分别为 extbf{2.267}和 extbf{2.286}，差距为 extbf{0.838 extbackslash%}；平均零样本困惑度分别为26.675和26.665，后者准确率平均低0.10个百分点。在单张RTX 5090上的原生四块测量显示，相较于优化后的BF16，Root计算速度提升2.50–2.83倍，且优化器峰值内存降低37.9–42.5 extbackslash%。
原文摘要 · Abstract (English)
Recent NVFP4 pretraining work has primarily optimized Transformer linear projections, leaving persistent optimizer states, optimizer computation, and low-precision attention forward--backward paths less explored. We present \textbf{Full-Stack FP4}, a modular NVFP4 framework with separate recipes for projections, AdamW states, Root/Muon computation, and attention. \textbf{LoRA-SVD} protects a compact projection subspace in BF16 while retaining full-shape NVFP4 computation, reducing the linear-only loss gap from \textbf{1.40\%} to \textbf{0.61\%}. An ordered square-root, tile-mean, and Hadamard pipeline enables stable NVFP4 AdamW momentum storage; shape-dependent coefficients and clipping stabilize direct NVFP4 Root iterations; and mixed-precision attention retains softmax-sensitive operations in BF16. On 3B pretraining with 64B tokens, BF16 and Full-Stack FP4 reach losses of \textbf{2.267} and \textbf{2.286}, a \textbf{0.838\%} gap. Their average zero-shot perplexities are 26.675 and 26.665, respectively, with Full-Stack FP4 averaging 0.10 percentage points lower in accuracy. Native four-block measurements on one RTX 5090 show 2.50--2.83$\times$ Root speedups over optimized BF16 and 37.9--42.5\% lower AdamW peak memory.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。