随着大语言模型和生成式 AI 模型规模持续增长,训练所需 GPU 时长和工程迭代时间也不断攀升。NVIDIA Hopper 和 Blackwell 架构 GPU 通过引入 FP8 和 NVFP4 等低精度算子支持,为解决这一问题提供了新途径。NVIDIA 发布的技术指南展示了如何利用 Transformer Engine 优化低精度训练。
Transformer 模型的大部分训练时间消耗在 GEMM(通用矩阵乘法)上,低精度格式主要通过加速矩阵乘法来提升训练速度。然而,模型配置本身并不告诉你哪些 GEMM 在运行。开发者需要将 Transformer 配置和批次大小转化为精确的 M×K×N 矩阵形状,然后跨精度进行基准测试,以确定架构的最优精度。
NVIDIA 提供的微基准测试工具可根据模型超参数直接推导 GEMM 形状,在多种精度下进行基准测试并计算完整加速比分析。通过 autocast 模式和 prequantize 模式的对比,可以清晰了解动态量化开销与纯 GEMM 内核性能之间的差异。测试结果显示,在 NVIDIA B300 上,NVFP4 相比 MXFP8 在大规模 GEMM 上可带来 1.46x 至 1.66x 的加速。
这套优化方法论帮助团队在投入大规模训练之前快速评估不同精度的收益,从而做出更明智的训练策略决策。
WeChat
Profile