返回首頁

AI 基礎設施/低精度訓練

Transformer Engine 2.19 接上 Rubin 與二維 MXFP8,並封鎖可能靜默算錯的 cuBLAS 演算法

NVIDIA 的低精度訓練函式庫新增 Rubin SM107a、MXFP8 二維權重縮放及 MoE 通訊路徑。版本同時停用一個會令 B300 與 Rubin 分組 GEMM 靜默損毀資料的 cuBLAS 13.7 演算法。

Strubbl · CC BY-SA 4.0 · Image source
zh-Hant

NVIDIA Transformer Engine 2.19 把下一代 Rubin GPU 的軟體支援推進到訓練核心:以 CUDA 13.4 以上編譯時可辨識 SM107a,PyTorch 端則新增實驗性的「分方向混合量化」,允許矩陣的 rowwise 與 columnwise 路徑採用不同 quantizer。這比整個張量套用單一精度策略更有彈性,尤其適合前向、反向或通訊階段具有不同誤差容忍度的工作負載。

較具體的新介面是 `MXFP8BlockScaling(enable_2d_quantization=True)`。一般 MXFP8 為每 32 個連續值保存一個 E8M0 scale;二維模式進一步針對權重配置區塊尺度,目標是降低方向性量化造成的誤差,但必須由使用者主動啟用。2.19 也把 MXFP8 帶入 NCCL-EP 的 dispatch forward 與 combine backward,支援預先量化的 `GroupedLinear` 和融合 grouped MLP,直接對準大型 MoE 訓練的專家通訊成本。

注意力路徑亦有變動:cuDNN FP8 fused attention 現支援 packed THD 輸入與 context parallelism,並加入實驗性 `GatedDeltaNetAttention` 核心;部分 FlashAttention 後端可在 `torch.compile(fullgraph=True)` 下工作。這些功能不是無條件加速:FP8 THD sink-softmax backward 要求 cuDNN 9.26 以上,GatedDeltaNet 在 `head_dim=32` 時已知反向傳播可能出現 NaN。

升級的安全面同樣重要。NVIDIA 停用了 cuBLAS 13.7 在 B300/Rubin 上一個可能產生靜默資料損毀的 grouped-GEMM 演算法;若正在驗證新硬體,應把版本固定與數值回歸測試視為必要工作。NCCL-EP 的 `ep_bootstrap`、`EpBuffer` 介面及 `mhc_fused_projection` 輸出型別也有破壞性變更。官方未提供 2.19 整體模型吞吐或收斂基準,因此工程團隊仍需按自身矩陣形狀、通訊拓撲與模型品質量測收益。

來源

  1. Transformer Engine v2.19 Release Notes
  2. Release v2.19.0
  3. Transformer Engine 2.19 Common API