论文
用于快速 FP4 预训练的格式感知融合
Format-Aware Fusion for Fast FP4 Pretraining
摘要
四位浮点 (FP4) 张量核心可加速矩阵乘法,但规模计算、操作数打包、布局构建和保存的后向状态可能会消除增益。我们提出了 format-aware fusion,它为本地 \mxfp{}、全局 \nvfp{} 和协作线程数组本地 \nvfp{} 共同设计每个量化生成器及其尺度域和消费者布局。我们使用 bfloat16 输出投影和编译的交叉熵通过 1600 亿个词元评估 Llama-3-family 8B 预训练。在匹配的相同加速器探针中,bfloat16 和 Transformer Engine \nvfp{} 达到 18.8K 和 27.6K 词元/秒/GPU,而我们最快的自定义路线达到 37.9K。采用行梯度随机舍入和固定符号 32 值 Hadamard 权重梯度预处理的 \mxfp{} 达到 37.2K 个词元/秒/GPU(86.3% bfloat16 模型 FLOP 利用率),并比原始 bfloat16 训练损失端点高出 2.11%。具有四个最终 bfloat16 块的 Transformer Engine 配方在 27.1K tokens/s/GPU 时比 bfloat16 高出 0.87\%。下游排名与训练损失排名不同,表明 FP4 结果共同取决于规模合约、操作数和执行路径。
