面向格式的融合：实现快速 FP4 预训练

arXiv:2610.00053v1 公告类型：new 摘要：四位浮点（FP4）Tensor Cores 可加速矩阵乘法，但尺度计算、操作数打包、布局构建以及保存的反向传播状态可能抵消这一优势。我们提出\emph{面向格式的融合}，针对每个量化生产者及其尺度域和消费者布局进行协同设计，以支持原生 \mxfp{}、全局 \nvfp{} 以及协作线程阵列本地 \nvfp{}。我们使用 bfloat16 输出投影和编译后的交叉熵，对 Llama-3 系列 8B 模型进行了 1600 亿 token 的预训练评估。在匹配的同加速器测试中，bfloat16 和 Transformer Engine \nvfp{} 的速度分别达到每 GPU 18.8K 和 27.6K tokens/s，而我们最快的自定义路径达到 37.9K。采用行梯度随机舍入和固定符号 32 值 Hadamard 权重梯度预条件的 \mxfp{} 达到每 GPU 37.2K tokens/s（bfloat16 模型 FLOP 利用率为 86.3%），最终训练损失比原始 bfloat16 训练损失终点高 2.11%。采用四个末端 bfloat16 模块的 Transformer Engine 配方达到每 GPU 27.1K tokens/s，最终结果比 bfloat16 高 0.87%。下游排名与训练损失排名存在差异，表明 FP4 的效果取决于尺度约定、操作数和执行路径的共同作用。