Prime-rl 現在對於 MoE 提供了廣泛的支持,無論是 RL 還是 SFT,我們已經用它訓練了 100B+ 的模型。 我們支持: * Qwen3 a3-30b * GLM 系列和 Moonlight * 我們正在添加 gpt oss 系列。 我們最終重寫了大部分建模代碼,以使其能夠與 torch compile 一起工作,同時仍然與 hugging face 生態系統兼容。
我們還新增了對 glm 系列進行微調的支持,因為我們發現這使得 muon 預訓練模型的訓練更加穩定。這與 moonlight 論文中的發現有很好的相關性。 我們正在使用 Microsoft 在 dion 中的實現。
在性能方面,我們與 torchtitan moe 訓練不相上下(主要是因為我們以 torchtitan 作為優化我們代碼的參考)。 我們在 512 h200 GPU 上使用 muon 進行 65K 序列長度訓練時,達到了 100B glm air moe 的 20% mfu。在並行性方面,我們發現 torchtitan EP 對於稀疏 MoE 的優化並不多,並且在足夠大的批次大小下並未帶來任何有意義的加速。 我們似乎可以很快飽和 MoE 的 group mm kernel 性能,而使用 EP 進行更好的負載平衡因為全到全的阻塞而產生的回報遞減。 我們相信,在開源堆棧中,MoE 有很多改進的空間。迫切需要優化的內核來加快 MoE 的速度,並以非常細粒度的方式隱藏通信和計算。這對於 blackwell 來說似乎更為必要。
啊,還有大部分的工作都是由 @jackminong 完成的。
18.03K