HyperAIHyperAI

Command Palette

Search for a command to run...

Blackwell平台运用NVFP4加速JAX/MaxText训练

近日,NVIDIA正式发布基于JAX与MaxText框架的NVFP4混合精度预训练配方,专为NVIDIA Blackwell架构优化。面对大模型训练中的算力瓶颈与精度权衡,该方案采用创新的两级微缩放编码,将NVFP4低比特精度精准应用于Transformer的MLP层,并结合二维块量化、随机哈达玛变换与随机舍入技术,有效抑制量化噪声并保障训练收敛。 性能实测表明,在GB200与GB300超级芯片上预训练Llama 3系列模型,该方案较传统FP8基线实现1.31至1.73倍的计算吞吐加速,单卡算力峰值提升逾500至700 TFLOP/s。在相同超参数与批大小下,NVFP4的训练损失曲线与FP8高度一致,收敛偏差仅0.026 nats,证实了极低精度下的无损精度表现。目前该配方已开源至JAX-Toolbox仓库,提供标准容器化部署路径。该技术通过显著压缩训练步时与算力成本,为AI基础设施高效迭代超大规模前沿模型提供了关键引擎。

相关链接