摩尔线程开源万卡级训练栈:1 万亿参数 MoE 预训练与强化学习

一句话结论:摩尔线程在 10 月 28 日开源的 Megatron-MUSA-Titan 把 NVIDIA Megatron-LM 的训练栈整体搬上了国产全功能 GPU,让 1 万亿参数规模的 MoE 模型在万卡级集群上跑预训练与强化学习时可以共用同一套底层代码。

训练一个万亿参数的 MoE 模型,难点从来不在写模型定义,而在集群里的四件事同时不出错:算力用得满不满、专家并行通信堵不堵、中间激活撑不撑得住显存、流水线阶段之间空等多久。10 月 28 日,摩尔线程把针对这四个环节的系统优化整体开放为 Megatron-MUSA-Titan 项目,仓库地址为 MooreThreads/Megatron-MUSA-Titan,代码创建于 2025 年 3 月,最近一次提交在 2026 年 10 月 16 日。

这个项目的定位很清楚,它不是一套模型,也不是一个训练算法。它是摩尔线程全功能 GPU 上的训练基础设施,官方给出的目标场景是 1 万亿参数以上模型与 10000 块以上摩尔线程 GPU 组成的集群,覆盖预训练与强化学习两个阶段。对开发者最实际的意义在于,NVIDIA 的 Megatron-LM、slime、miles 三个官方仓库可以原样克隆下来跑在摩尔线程 GPU 上,不需要改动上游代码。

万卡级训练栈到底包含哪些东西

仓库当前的规模是一个可以在 GitHub 页面上直接数出来的数字,305 个文件,其中顶层目录分为 musa_titan、examples、docs、test 四块。musa_titan 下面是 Megatron 的适配代码,按预训练和强化学习分成两条路径;examples 下是可直接启动的训练配方;docs 下的特性说明文档逐条列出了每个优化项对应的环境变量和命令行参数。

支撑这个栈的外部组件是一串可以逐个核对的开源仓库。torch_musa 提供 PyTorch 对 MUSA 平台的支持,MUSA C 编译器对 CUDA C 的兼容让绝大多数算子代码不用重写;MT-TransformerEngine 在 NVIDIA TransformerEngine 基础上加了 FP8 训练实现;DeepEP 与 ACE 提供 MoE 的 token 分发通道,Mooncake 提供 RDMA 流水线传输。这些仓库加起来的社区体量,与 Megatron-MUSA-Titan 自身形成了鲜明对比。

组件 仓库 Star 数 在训练栈中的角色
Megatron-LM NVIDIA/Megatron-LM 18286 训练主体框架,官方仓库直接可用
SGLang sgl-project/sglang 38310 强化学习阶段的 rollout 引擎
slime THUDM/slime 8750 强化学习算法框架之一
Mooncake kvcache-ai/Mooncake 6907 RDMA 流水线并行通信与计算重叠
torch_musa MooreThreads/torch_musa 518 PyTorch 对 MUSA 的原生集成
MT-TransformerEngine MooreThreads/MT-TransformerEngine 13 FP8 训练配方与 FP8 算子实现
Megatron-MUSA-Titan MooreThreads/Megatron-MUSA-Titan 17 本次开源的 MUSA 训练基础设施

这组数字说明一件事。训练栈真正难的部分从来没有被单独封装成一个库,它一直散落在框架、编译器和通信库三处,这次开放的是把这些部件在国产硬件上打通并调优后的整体结果。也正因为刚公开,这个仓库的 Star 数只有 17,Fork 5,社区验证还在早期阶段。

四个环节的五项关键优化

万亿参数 MoE 训练的瓶颈可以拆成计算、通信、显存、调度四类。Megatron-MUSA-Titan 针对每一类都给出了对应的实现,并且全部以开关形式写在配置里,没有编译期的隐藏逻辑。

环节 优化项 依赖 启用方式
计算 blockwise FP8 训练 MT-TransformerEngine –fp8-format e4m3、–fp8-param-gather、–fp8-recipe mxfp8
通信 DeepEP over ACE token 分发 ACE 版本的 DeepEP USE_DEEPEP_ACE=1 配合 –moe-token-dispatcher-type flex
通信 专家并行通算重叠 DeepEP 与 ACE –overlap-moe-expert-parallel-comm
显存 VPP 与 1F1B 下的激活卸载 虚拟流水线并行 –offload-moe-fc1-input、–offload-moe-fused-swiglu-input
调度 流水线通信与计算重叠 Mooncake RDMA PP_USE_MOONCAKE=1 及 Mooncake 参数

FP8 这一项值得单独说一句。MT-TransformerEngine 把上游的 MXFP8BlockScaling 映射到自研的 MTFP8BlockScaling 实现,并提供了与 DeepSeek-V3 同款 FP8 训练策略的 MTFP8BlockScalingRecipeState。当前特性表列出五项已支持的能力:per-tensor FP8、per-block FP8、带 FP8 的张量并行重叠、MoE 重计算、zero bubble,FP8 alltoall 仍标记为 Coming Soon。

激活卸载有一个容易被忽略的硬性前提,官方文档明确写了当前配置要求 PP size 大于等于 3。因为卸载动作要挂在流水线的节奏上,流水线并行度太低就没有可重叠的阶段可以腾显存。

文档里还有一处细节说明了作者对真实集群的熟悉程度。Mooncake 的 P2P slot 大小在文档示例中给的是 134217728 字节,即 128 MiB,并特别注明这只是示例值,要按流水线配置里最大的激活或梯度传输量来调整。对照 Kimi-K2 的 1T 启动脚本,里面实际写的是 58720256 字节。默认值和实际生产取值不一样,这类差异通常只有跑过集群的团队才会注意到。

1 万亿参数配方里能读出什么

仓库里最能说明能力边界的是那份 Kimi-K2 1T MoE 配方脚本,它是官方文档反复指向的万卡训练脚本参考。脚本里的并行配置和模型结构都写在明处,任何人都能逐行核对。

配置项 取值 说明
num-layers 61 与 Kimi-K2 结构一致
hidden-size 7168 隐藏层维度
num-attention-heads 64 注意力头数
seq-length 4096 训练序列长度
vocab-size 163840 词表大小
ffn-hidden-size 18432 前馈网络中间维度
q-lora-rank / kv-lora-rank 1536 / 512 多头潜在注意力低秩压缩
默认并行度 TP 1、PP 31、EP 8 单节点 8 卡,PP 为主要切分维度
global-batch-size 1024 micro-batch-size 为 1
train-samples 2441406200 训练样本总量
学习率 2e-4 WSD 衰减风格

把并行度摊开看能得出一个直观判断。在这套默认配置里,张量并行是 1,专家并行只有 8,剩下 31 层全部切给流水线并行。这说明训练的重点被放在跨节点的流水线效率上,而不是单卡内的算力切分,与前面提到的 Mooncake RDMA 流水线重叠、PP size 大于等于 3 的激活卸载限制是自洽的。

脚本里的通信参数同样值得看。MCCL_IB_GID_INDEX 设为 3,MCCL_IB_TC 设为 136,MCCL_IB_QPS_PER_CONNECTION 为 16,MCCL_BUFFSIZE 为 20971520,MCCL_CROSS_NIC 为 0。INFiniBand 的这些参数会直接决定跨节点 all-to-all 的吞吐上限,写在脚本而不是留在文档里,说明这套配置是从真实集群上跑出来的。

除了 Kimi-K2,配方目录还覆盖 deepseek-v2、deepseek-v3、ling-1t、llama2、llama3、mixtral、qwen3。其中 ling-1t 的配置用了 256 个专家,路由 topk 为 8、分组 topk 为 4、路由函数为 sigmoid,同样开启 blockwise FP8。DeepSeek-V3 的 lite 版本则把 ffn-hidden-size 降到 2048 配 129280 词表,是给单节点或小规模集群做的复现入口。

预训练与强化学习怎么共用一套底座

这是 Megatron-MUSA-Titan 与常规训练适配项目差别最大的地方。仓库把 RL 阶段的组织方式定成三层分工:Megatron 负责权重更新与训练,slime 或 miles 负责强化学习算法本身,SGLang 承担 rollout 样本生成。RL 侧代码放在 musa_titan/rl 下的 slime_musa 与 miles_musa 两个子目录,共享预训练那套 Megatron 适配逻辑,不再复制一份。

接入方式刻意做得很轻。RL 侧需要在进程环境里设置 SLIME_MUSA_PATCH_AT_STARTUP=1 或 MILES_MUSA_PATCH_AT_STARTUP=1,Python 入口处调用 musa_titan.apply_rl_patch()。仓库根目录还放了一个 sitecustomize.py,会自动为 Ray 的 worker 进程和子进程调用同一个补丁函数,因此不需要在每个节点上手工注入。

RL 的示例配方目前提供 slime 和 miles 两套,模型统一用 Qwen3-30B-A3B,附了标准版与 16k 长上下文档位。启动方式是 Ray 作业,runtime_env.yaml 模板里不写死任何机器路径,路径由 set_env.sh 注入,模板本身刻意保持可移植。

生态侧有一个容易被忽略的进展。MUSA 已经合入 slime 与 RL-Kernel 的 main 分支,miles 的对应 PR 仍在推进中。也就是说摩尔线程对训练框架的改造不再是长期分叉,改动正在回流主线,这决定了第三方训练项目迁移到国产 GPU 时要维护的补丁量会越来越少。

30 分钟定位万卡集群的瓶颈

集群规模一大,性能分析本身就变成瓶颈。单个算子耗时、通信链路等待、流水线空泡、少数慢节点,任何一个都可能拖住整体,而靠人工翻日志在万卡规模下基本无效。

摩尔线程为此做了 Distributed Holistic Trace Analysis 工具,读取各 GPU 的 PyTorch trace,定位高开销算子、通信瓶颈、流水线空泡率与掉队节点,并给出流水线并行的关键路径分析。官方给出的数字是可以在 30 分钟内完成一次万亿参数模型的性能分析。对照仓库里的说明,它包含两个互补的部分,单卡侧的调用栈分析用于把算子开销对应回 Python 代码,分布式侧则跨 rank 聚合 trace。

这个工具的代码作者是 GitHub 用户 yongpan0709,以独立仓库形式发布,不在 Megatron-MUSA-Titan 主仓内。用法上有一点需要注意,单卡调用栈分析需要在采集 trace 时把 torch.profiler 的 with_stack 设为 True,trace 里才会带上调用栈信息。分布式 trace 分析同样可以用于非摩尔线程的加速卡。

迁移到这套栈需要做哪几件事

文档给出的路径很短。装好摩尔线程驱动、MUSA 工具包、torch_musa 与 MT-TransformerEngine,然后把 Megatron-LM 与本仓库同时放进 PYTHONPATH,环境变量 ACCELERATOR_BACKEND 设成 musa,在 Python 入口导入 Megatron 之前调用 apply_pretrain_patch(),RL 阶段换成 apply_rl_patch()。运行时用 torchrun 加 nnodes 与 node_rank 启动,SSH 多节点场景需要在模型目录里准备 hostfile,每个节点一行地址。

文档也主动划了边界,明确写着这些配方是面向具体集群配置的起点,而不是对已公布的万卡基准的一键复现。README 里的原话是 recipes are starting points for cluster-specific configuration, not one-command reproductions of the reported 10,000-GPU benchmark。这句话很重要,它把可验证的部分和需要使用者自己补的部分划清了。

还有两个环境变量容易踩坑。文档指出环境变量必须在启动每个训练 worker 的 Bash 脚本里设置,在启动节点上 export 并不会通过 SSH 或 Ray 自动转发。示例脚本也可能覆盖已导出的值,必要时得改 worker 侧的脚本。这条提示对多节点作业是硬性要求。

放在国产算力版图里看它的位置

把 Megatron-MUSA-Titan 放到更长的链条里看,它填的是训练软件栈这一段。硬件侧的国产替代正在密集推进,阿里巴巴在云栖大会上给出了真武 V900 芯片与 50 万卡集群的路径,众星微则用 12 款互连芯片补上国产算力在互连环节的短板。集群硬件的规模问题之外,把 Megatron 这类主流训练框架完整跑起来的软件适配此前缺少可公开引用的样本。

商汤大装置走的是另一条路,用 PD 分离把异构集群的吞吐拉高 5 倍,属于推理侧与调度侧的工程。相比之下 Megatron-MUSA-Titan 聚焦的是训练侧的完整落地。同一条链上还有 KRAFTON 开源的 Raon-1 mini 这类把 MoE 训练配置做成可复现脚本的项目,以及阶跃 Step 5 Preview 开放权重后 6000 亿稀疏 MoE 的公开权重工作,它们共同构成的是模型侧而非基础设施侧的可复现性。

把时间线摆在一起看更清楚。Kimi K3 的技术报告解读里已经出现了 2.8 万亿参数 MoE 与 90% 编码缓存命中率这类超大规模训练细节,NVIDIA 侧的 Megatron-LM 社区在同期持续迭代。摩尔线程这次开放的是把同样量级的训练在国产硬件上落地所需的全部适配层,而且是可查、可改、可提交 Issue 的形式。

常见问题

Megatron-MUSA-Titan 和直接用 Megatron-LM 有什么区别

框架本身没有区别,训练主体仍然是 NVIDIA 的 Megatron-LM。区别集中在四处针对国产硬件的系统优化:blockwise FP8 训练、基于 ACE 的 DeepEP 专家并行通信、挂载在 VPP 与 1F1B 调度上的激活卸载,以及借助 Mooncake RDMA 的流水线通算重叠。加上 torch_musa 与 MUSA C 提供的兼容层,官方 Megatron-LM、slime、miles 仓库可以原样使用。

这个仓库现在能拿来做什么,不能做什么

能做的事情是完整的:跑预训练配方、开强化学习作业、用 trace 工具做集群性能定位。不能做的事情是复现官方公布的万卡基准数字,因为公开的配方只提供起点,集群规模、并行度、EP 配置与数据路径都需要按自己的环境重新调。仓库本身刚公开,Star 17、Commit 集中在 2026 年 7 月到 10 月,稳定性与社区验证都还在早期。

FP8 训练要额外装什么

必须装 MT-TransformerEngine,它提供优化后的 FP8 实现,是 blockwise FP8 训练的前提条件。Megatron 侧需要加上 –transformer-impl transformer_engine、–fp8-format e4m3、–fp8-param-gather 与 –fp8-recipe mxfp8 这组参数。当前 FP8 alltoall 在特性表里仍是 Coming Soon,涉及 alltoall 的密集通信场景需要另做评估。

小结

Megatron-MUSA-Titan 的价值不在于又一个训练框架,而在于把国产全功能 GPU 拉进万亿参数 MoE 训练的完整链路,并且把这条链路上最难调的四类工程问题各自做成了可开关的配置项。四个 1T 级别的预训练配方、两套强化学习框架适配、一个 30 分钟出结论的集群性能分析工具,共同构成了这套底座的可用性拼图。

需要留意的是它刚开源,公开可核对的性能数据目前只有官方口径的高 MFU 与长稳训练描述,缺少第三方复现。真正决定它能否被更多团队采用的,是上游社区的合入进度与社区贡献密度,这两项在接下来的几个月里会比任何一份基准报告更有参考价值。

相关阅读