mlc-ai/pith-train

Compact and Agent-Native MoE Training System

解決的問題

生產級 MoE（專家混合）訓練框架通常過於複雜，依賴重型編譯擴展和間接設計模式，使得人類與 AI 編碼代理難以修改或優化。PithTrain 提供一個高性能量的訓練系統，其程式碼量緊湊到足以讓代理或人類從頭到尾閱讀，同時維持生產級功能。

如何運作

這是一個面向代理的框架，分為三層：用於訓練迴圈的應用層、用於模型協定與分散式訓練的引擎層，以及用於底層內核的運算元層。它實作 4D 並行策略（流水線、資料、上下文與專家並行），並使用 DualPipeV 排程器實現前向-後向執行與 P2P 通訊的重疊。系統支援 FP8 訓練，並與 DeepGEMM 和 FlashAttention 等函式庫整合。

適用對象

需要訓練 MoE 模型（如 Qwen 或 DeepSeek 架構）的研究人員與開發者，希望使用一個容易導航、修改與擴充的程式碼庫，特別是在使用 AI 編碼助手時。

特色

  • 面向代理的設計：約 11K 行 Python 程式碼的緊湊程式碼庫，極少隱式間接引用。
  • 生產級效能：支援 4D 並行與計算-通訊重疊。
  • FP8 訓練：內建 FP8 線性與量化支援，提升效率。
  • 彈性檢查點：支援可重分片檢查點與轉換為 Hugging Face 格式。

相關

  • Dispatch
    專家混合（MoE）說明

    專家混合（MoE）讓 transformer 模型在擴展參數規模的同時，透過僅在每個 token 上啟用部分神經網路專家，維持高效的預訓練與更快的推論速度。

  • 專案
    microsoft/Tutel

    Tutel 是 Microsoft 開源的高性能 MoE 函式庫，適用於大型語言模型。它提供動態並行、低精度推理（NVFP4、MXFP4、FP8、BF16）、支援高達 1M token 的上下文，以及視覺擴展功能，並提供 Docker 映像檔來運行 DeepSeek-V3.2、Kimi-K3、GLM-5.x、Qwen-3 和 GPT-OSS 等模型，支援 NVIDIA 和 AMD GPU。可以透過 pip 從 Git 儲存庫安裝，或從源碼構建，然後運行提供的 Docker 容器，或使用 Python API 進行自定義 MoE 路由。

  • 專案
    deepseek-ai/DeepEP

    用於 MoE 模型的高性能通訊函式庫，提供優化的專家並行 GPU 核心，以最大化吞吐量並最小化 SM 資源占用。

  • 專案
    PaddlePaddle/ERNIE

    一組提供文本與視覺語言理解方面最先進性能的大規模多模態 MoE 模型與開發工具包。

  • Dispatch
    Qwen1.5-MoE-A2.7B 版本說明

    Qwen 推出 Qwen1.5-MoE-A2.7B，一款混合專家模型，其效能與 7B 密集模型相當，但僅使用 27 億個啟用參數。