[论文] MemFine：内存感知细粒度调度用于 MoE 训练

Source: arXiv - 2511.21431v1

概览

训练大规模的 Mixture‑of‑Experts（MoE）模型正面临硬性瓶颈：令牌的动态路由导致严重的负载不平衡，进而使 GPU 内存使用量激增。MemFine 引入了一种内存感知、细粒度的调度系统，将令牌流和专家工作负载切分为更小的“块”，使训练器只在需要时重新计算，同时保持在普通 GPU 的内存限制之内。其结果是一个更具可扩展性的 MoE 训练流水线，能够在之前因内存不足而无法运行的硬件上运行。

关键贡献

• 基于块的分解： 将令牌分配和专家计算拆分为可管理的块，支持选择性重新计算。
• 理论内存模型： 提供闭式表达式用于预测内存消耗，并实时指导调度器的决策。
• 动态调度算法： 在运行时优化内存节省与计算吞吐之间的权衡，无需手动调参。
• 实证收益： 相比基线全重新计算方法，激活内存最高可降低 48 %，吞吐提升 4.4 %。
• 硬件友好设计： 适用于内存受限的普通 GPU，免除对高端硬件或大幅降低批量大小的需求。

方法论

1. 令牌与专家块划分 – MemFine 不再把整个批次视为整体，而是将输入的令牌流和专家集合划分为更小的块。每个块可以独立处理，从而限制单次操作的峰值内存需求。
2. 块级重新计算 – 对于在反向传播时通常需要保存激活的层，MemFine 有选择地丢弃激活，仅在反向传播时重新计算必要的块。内存模型会引导此过程，确保重新计算成本不超过内存节省。
3. 基于内存模型的调度器 – 轻量的分析模型估算每种块配置的内存占用。调度器随后选取满足 GPU 内存预算且最大化吞吐的配置。该决策在每个训练步骤都会更新，以适应 MoE 中不断变化的令牌路由模式。
4. 实现细节 – 作为 PyTorch 等主流深度学习框架的插件，直接替换标准 MoE 调度器，用户只需进行极少量代码修改即可使用。

结果与发现

• 内存节省： 通过将激活内存削减近一半，原本在 16 GB GPU 上会崩溃的模型现在能够稳定训练。
• 吞吐提升： 细粒度调度仅带来极小的重新计算开销，而内存导致的停顿减少带来的收益更大。
• 准确率： 由于 MemFine 只重新计算被丢弃激活的精确前向路径，几乎不影响最终模型质量。

实际意义

• 成本效益的扩展： 企业可以在不购买昂贵多 GPU 服务器的情况下，将 MoE 模型扩展到数十亿参数；单卡 16‑32 GB GPU 已足以支撑多数工作负载。
• 运维更简洁： 无需手动调节专家容量或批量大小以适配内存——MemFine 的调度器自动完成，降低工程师工作量。
• 更广泛的可及性： 硬件预算有限的研究者和初创公司也能尝试此前难以触及的前沿 MoE 架构。
• 集成路径： MemFine 直接插入现有 MoE 库，开发者只需几行代码即可获得即时的内存优势，无需重写模型逻辑。

局限性与未来工作

• 在极大批量上的重新计算开销： 虽然调度器已尽量减轻，但非常大的批量仍可能导致显著的重新计算延迟。
• 硬件多样性： 当前评估主要针对 NVIDIA GPU；将该方法迁移到 TPU 或 AMD GPU 可能需要额外调优。
• 动态专家数量： MemFine 假设每层的专家数量是固定的，处理训练过程中增删专家的模型仍是未解难题。
• 未来方向： 作者计划基于运行时剖析探索自适应块大小，将 MemFine 与混合精度训练结合，并开源更通用的调度器 API，以促进社区广泛采用。

作者

• Lu Zhao
• Rong Shi
• Shaoqing Zhang
• Yueqiang Chen
• Baoguo He
• Hongfeng Sun
• Ziqing Yin
• Shangchao Su
• Zhiyan Cui
• Liang Dong
• Xiyuan Li
• Lingbin Wang
• Jianwei He
• Jiesong Ma
• Weikang Huang
• Jianglei Tong
• Dongdong Gao
• Jian Zhang
• Hong Tian

论文信息

• arXiv ID: 2511.21431v1
• 分类: cs.DC
• 发表时间: 2025 年 11 月 26 日
• PDF: Download PDF