
大模型训练已经跑到万卡集群规模,显存始终是那道隐形的门槛。
MoE 架构把大模型推到了万亿参数级别,也埋下了一颗定时炸弹:显存尖峰。某个微批次的 token 一旦集中涌向少数专家,激活显存瞬间飙升,一张卡 OOM,整个分布式任务直接挂掉 —— 几百步白跑。
重计算、激活卸载都是常用手段,但问题在于静态配置。激进配置平时 MFU 30%,尖峰必崩,一次冷重启 284 秒;全层重计算稳是稳,MFU 直接掉到 24.6%,算力长期闲置。怎么选都亏。
能不能让系统自己跟着显存压力走 —— 平稳时全速跑,尖峰来了临时加码,过去了再收回来?
清华大学翟季冬教授团队的答案是八卦炉2(BaGuaLu-2 ),相关论文已被SC26录用。运行时自适应,零配置干预 —— 门控算完、专家算之前这个窗口里,系统根据真实负载实时决策,不同算子分别选择保存、重计算或卸载,把显存峰值压在设备上限以内。不用停任务,不用改配置。
万卡智算中心实测:128 到 12,288 张海光 DCU,188B 到 1.04 万亿参数,平均吞吐比基线高 1.14×,12,288 卡下总吞吐 2.012 EFLOPS,平均 MFU 32.9%、最高 35%。
论文链接:https://arxiv.org/abs/2610.23064
一、问题背景:MoE 训练中的显存尖峰痛点
如果你训练过 MoE 模型,大概率遇到过这种情况:训练跑了几百个迭代都很平稳,突然某个微批次里 token 集中涌向少数几个专家,激活显存瞬间飙升,一张卡 OOM,整个分布式任务直接挂掉。
问题出在 “选择” 本身
MoE 的核心是路由器为每个 token 选择少数专家执行。现代负载均衡方法能抑制长期的专家拥塞,但无法预知当前批次的具体内容。实测数据:97.1% 的迭代中专家负载不均衡率低于 1.2×,但 2.9% 的异常迭代中最高达到 2.6×。就是这极少数迭代,反复打断训练进程。
流水线让问题更严重
流水线预热阶段只有前向、没有反向释放激活,多个微批次的激活同时驻留设备。如果其中一个批次把大量 token 路由到当前设备的少数专家,额外的 MoE 激活叠加在已驻留激活之上,局部偏斜直接变成设备级 OOM。
静态配置的两难
重计算和卸载是常用工具,但传统做法是在启动前写死配置。激进配置正常阶段 MFU 30.0%,但尖峰必 OOM,一次冷重启 284 秒;全层重计算的保守配置 MFU 降到 24.6%,算力长期闲置。无论怎么选,都在浪费资源。

图 1:MoE 负载尖峰会推高激活显存,静态策略不得不在 OOM 重启和保守运行之间反复切换。
二、设计思路:运行时自适应,零配置干预
八卦炉 2 的核心思路很直接:显存压力在变,策略也应该跟着变。平稳时少做干预保留高吞吐,尖峰时临时加强缓解,尖峰过去自动恢复 —— 不终止任务,不修改配置,不重启。这套机制叫 Flexible Activation Provisioning(灵活激活显存供给)。

图 2:平稳阶段采用较轻的显存缓解,尖峰到来时临时加大力度,显存曲线始终被控制在设备上限之内。
① 激活显存建模:可扩展的成本模型
系统把保存激活、重计算、卸载放进同一个成本模型,同时估算显存收益和时间开销,并考虑额外操作能否与原有计算通信重叠。同样大小的卸载,发生在不同流水线位置,暴露在关键路径上的时间可能完全不同。多模型多配置下,平均预测误差仅 1.56%。模型设计为可扩展,方便接入新的算子和策略。
② 自适应供给:算子级精细控制
决策点选在一个关键窗口:MoE 门控计算完成后、专家计算开始前。此时真实负载已知,风险尚未落地,仍有时间临时加强重计算或卸载。策略不再以整层为单位,不同算子分别选择保存、重计算或卸载,只选满足预算的最小组合。系统还会追溯激活来源,排除那些不同精度重复保存、一个卸载另一个仍驻留的冗余组合。

图 3:系统根据当前流水线状态和 MoE 负载,依次完成激活显存建模、自适应供给与双批次融合执行。
③ 双批次融合:动态规划调度,开销可忽略
重计算和卸载不是免费的,直接排在关键路径上就是用吞吐换安全。八卦炉 2 利用流水线空隙 —— 稳定阶段设备上同时存在当前前向、当前反向、前一前向卸载和下一反向恢复,这些操作属于不同微批次,保持依赖顺序即可交错执行。系统用轻量级动态规划重排计算流和通信流,把额外工作塞进硬件空泡;连续算子合并、预算状态缓存,大多数迭代直接复用调度结果。最终决策开销低于 10 ms,而一次迭代超过 18 秒。

图 4:先通过动态策略压低显存峰值,再交错调度相邻批次,将额外计算和数据搬运藏入硬件空泡。
三、实测效果:多规模验证,代码已开源
全规模性能领先
在某万卡智算中心,模型从 188B 到 1.04 万亿参数,集群从 128 到 12,288 张海光 DCU(BW1000)。相比 Megatron-Kuai 和 AdaPipe,平均提升 1.14× 和 1.13×(已含 OOM 重启损失)。12,288 卡下总吞吐 2.012 EFLOPS,平均 MFU 32.9%,最高 35%。

图 5:从 128 到 12,288 张加速卡,八卦炉 2 在所有测试规模上均保持吞吐领先。
消融实验
从全重计算基线出发:细粒度激活策略 1.16× → + 动态管理 1.22× → + 重叠调度 1.38×。每个模块都有明确收益。

图 6:细粒度激活策略、动态显存管理和重叠调度依次叠加,最终形成 1.38× 的整体收益。
开源与快速上手
代码仓库:GitHub — thu-pacman/BaGuaLu
已开源内容:激活显存分析模型、调度器集成的激活卸载与回载等核心代码
低门槛集成:提供适配样例和中立评测基准,方便集成到现有训练流水线
零配置使用:开箱即用,支持科研实训和快速验证
欢迎提交 Issue 和 PR,参与技术共建。仓库公开范围不包含论文介绍的全部运行时机制,核心模块已可独立使用。
四、团队与社区
八卦炉 2 由清华大学翟季冬教授团队主导,清程极智、北京科技大学及独立研究者联合研发,团队致力于大模型训练和推理系统的开源共建,欢迎开发者和研究者共同探索前沿技术。
赤兔推理引擎:https://github.com/thu-pacman/chitu
八卦炉训练系统:https://github.com/thu-pacman/BaGuaLu
该团队的评测技术亦应用在智算中心的能耗控制方面,对于使用国产算力较多的金融云效果将更为明显。
欢迎点击 “阅读原文” 或打开 https://arxiv.org/abs/2610.23064,阅读完整论文!
清程极智现高薪直聘 AI infra 工程师,欢迎投递!

© THE END
转载请联系本公众号获得授权
投稿或寻求报道:liyazhou@jiqizhixin.com