隨著大語言模型和生成式 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