# 為什麼訓練 MoEs 如此困難 最近,我發現自己想要一個小型的、專注於研究的訓練庫, 可以快速輕鬆地進行小實驗。這些實驗範圍 從嘗試新的注意力架構(MLA、SWA、NSA、KDA - 全部可插拔)到多精度訓練,最近還有使用「新」優化器的多優化器設置。我嘗試了三個主要的 競爭者(Nemo、Megatron 和 Torchtitan),但由於許多不同的原因,它們 都不符合我的需求,並且設置、使用和穩定運行都相當痛苦。我再次懷念我在谷歌的工具, 並且為此重寫我的生產訓練堆棧(這是專門為大型基礎設施監控和穩定性量身定做的)也感覺像是浪費時間,會使舊的和新的庫都變得更糟。 然而,這讓我思考,為什麼訓練前沿質量的「小型」 MoEs(例如總參數少於 20B)如此困難?我想要的庫/庫為什麼還不存在?經過一段時間的思考,我能想到的大多數挑戰歸結為三個不同的方面: - FLOPS / FLOP 效率 - 負載平衡 / 路由穩定性 - 數據質量和數量 FLOPS 訓練密集模型現在相當簡單。訓練 動態大多是耦合的,如果你的架構中有足夠的參數, 模型幾乎會學會,儘管你犯了很多錯誤(這讓我多次受傷)。[DeepSeek 風格的超稀疏]( MoEs 不同,因為你的訓練動態在某種程度上是解耦的。只有一部分 MLP 在給定的標記上是活躍的, 隨著訓練的進行,活躍的專家會隨著時間的推移而變化和演變。這就是為什麼多輪訓練和數據重述對 MoEs(尤其是較大的 MoEs)如此有效的原因。 你會獲得巨大的推理效率提升和小的訓練效率提升,但代價是解耦的訓練動態(使得可預測和穩定訓練變得困難),而且你必須投入更多的 FLOPS 以確保學習到某種最佳的路由策略,並且參與各種策略的專家得到充分訓練。好吧,這就是 FLOPS / FLOP 效率挑戰出現的地方。第一個挑戰是,根據本質,超稀疏的 MoEs 需要大量的 HBM 來加載 專家,但這意味著你需要很多 GPU,因此系統中有很多閒置的 GPU。FSDP(和各種其他分片拓撲)大多是密集訓練的遺物,並且在充分利用所有這些閒置的 FLOPS 方面做得不好。這導致大多數人的超稀疏 MoE 訓練的 MFU(低)單位數字。雖然有幾種方法可以解決這個問題(未來會有更多關於這方面的內容),但我專注於兩個具體的方面: - 新的分片拓撲(一種新穎的專家並行調度系統,保持 GPU 忙碌) - 混合精度訓練 負載平衡 / 路由穩定性 我將把新的分片拓撲留給專門的寫作,但混合精度訓練是顯而易見的。你有這麼多閒置的 HBM,所以降低專家精度,將其減少一半或四分之一(分別為 fp8 和 nvfp4)。這在理論上是好的,但在實踐中,混合精度訓練通常需要更多的 HBM,因為你必須保持主權重和梯度在更高的精度中,然後將主權重量化到其較低的精度表示並緩存它們以供下一次前向傳遞使用。因此,它有助於推理效率(隨著越來越多的 FLOPS 用於 RL 和推理,這是一個真正的勝利),但代價是訓練期間需要更多的 HBM 和更多的 FLOPS。減少混合精度開銷應該是一個特定的重點。然而,任何你觸及的東西,無論是降低權重的精度和準確性,最終都會導致訓練動態的其餘部分不穩定。對於 MoE,這通常首先顯示出來的是路由穩定性。 ...