Transformer Engine 讓 MoE 訓練效能提升逾十倍

核心訊息

在大型 AI 訓練中，Mixture of Experts（MoE） 已成為重要架構。透過 NVIDIA 的 Transformer Engine 搭配 JAX，DeepSeek‑V3 在 NVIDIA GB200 上的訓練效能從 103 TFLOPS/GPU 提升至 1,068 TFLOPS/GPU，相當於 10.4 倍 的加速。

背景與技術挑戰

MoE 以「條件運算」取代傳統的單一密集前饋網路（FFN），將大量小型 專家網路 與一個學習式 路由器 結合，僅啟動前 K 名最適合的專家。此設計在 訓練計算量 上相較於等效的密集模型大幅降低，DeepSeek、Qwen、Mixtral 等皆證明了其高效能。

然而，規模化 MoE 訓練會遭遇以下瓶頸：

Token 路由：路由器在訓練過程中會產生偏斜的分配，導致每個專家收到的 token 數量差異極大。
專家分派與彙總（dispatch & gather）：需要在多 GPU 之間傳遞不等長的資料，若未優化會讓 all‑to‑all 通訊佔用 84% 的 kernel 時間。
Ragged GEMM：因為每個專家的 token 數量不一致，無法形成傳統的矩形矩陣乘法（GEMM），傳統庫存的最佳化假設被打破。

這些問題使得 GPU 計算資源常被 通訊等待 所拖慢，效能提升受限。

Transformer Engine 的解決方案

Transformer Engine 為 MoE 設計了 原生支援 ragged 版面的專用 kernel，可直接在不規則張量上執行矩陣運算，避免了額外的資料填充或重排。結合 JAX 的高階抽象，開發者只需在程式碼層面呼叫相應 API，即可獲得：

Token 路由 的低延遲實作
專家分派/彙總 的高效 all‑to‑all 通訊
Ragged GEMM 的最佳化計算路徑

實測顯示，這套優化把 DeepSeek‑V3 的每顆 GPU 計算吞吐量提升至 1,068 TFLOPS，相較於未優化的 103 TFLOPS，提升幅度超過十倍。

未來展望與讀者啟示

MoE 的 dropless 與 capacity‑based 兩種路由策略各有取捨：dropless 讓每個 token 必定被選定的專家處理，雖有助於模型品質，卻對系統資源要求更高；capacity‑based 則透過限制每個專家的負載，減輕不均衡帶來的效能衝擊。Transformer Engine 已支援兩者的核心運算，未來隨著 MegaBlocks 等研究將專家運算重新表述為 block‑sparse 矩陣乘法，預計能進一步縮減記憶體占用與通訊開銷。

對於從事大規模模型訓練的工程師而言，採用 NVIDIA Transformer Engine + JAX 已不再是選項，而是提升 MoE 訓練效能的必要路徑。隨著硬體與軟體的持續協同優化，MoE 有望在保持高效能的同時，推動更大規模、更高品質的 AI 模型快速落地。