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）通过在每个 token 上仅激活部分神经网络专家，使 Transformer 模型能够在扩展参数规模的同时保持高效的预训练和更快的推理。

  • 项目
    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，这是一款 Mixture-of-Experts 模型，能够在仅使用 27 亿激活参数的情况下匹配 7B 密集模型的性能。