The training of large-scale Mixture of Experts (MoE) models faces a critical memory bottleneck due to severe load imbalance caused by dynamic token routing. This imbalance leads to memory overflow on GPUs with limited capacity, constraining model scalability. Existing load balancing methods, which cap expert capacity, compromise model accuracy and fail on memory-constrained hardware. To address this, we propose MemFine, a memory-aware fine-grained scheduling framework for MoE training. MemFine decomposes the token distribution and expert computation into manageable chunks and employs a chunked recomputation strategy, dynamically optimized through a theoretical memory model to balance memory efficiency and throughput. Experiments demonstrate that MemFine reduces activation memory by 48.03% and improves throughput by 4.42% compared to full recomputation-based baselines, enabling stable large-scale MoE training on memory-limited GPUs.


翻译:大规模专家混合(Mixture of Experts,MoE)模型的训练面临严峻的内存瓶颈,其根源在于动态令牌路由导致的严重负载不均衡。这种不均衡使得容量有限的GPU出现内存溢出,制约了模型的可扩展性。现有的负载均衡方法通过限制专家容量来缓解问题,但会损害模型精度,且在内存受限的硬件上仍可能失效。为此,我们提出MemFine,一种面向MoE训练的内存感知细粒度调度框架。MemFine将令牌分布与专家计算分解为可管理的块,并采用分块重计算策略;该策略通过理论内存模型进行动态优化,以平衡内存效率与训练吞吐量。实验表明,与基于完全重计算的基线方法相比,MemFine将激活内存降低了48.03%,并将吞吐量提升了4.42%,从而能够在内存有限的GPU上实现稳定的大规模MoE训练。

0
下载
关闭预览

相关内容

【CVPR2025】基于组合表示移植的图像编辑方法
专知会员服务
8+阅读 · 2025年4月5日
专知会员服务
30+阅读 · 2020年9月18日
【NeurIPS2019】图变换网络:Graph Transformer Network
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
VIP会员
最新内容
深入Project Maven:为何人工智能在战场上依然失灵
专知会员服务
3+阅读 · 47分钟前
锻造未来士兵:外骨骼、基因工程与赛博格
专知会员服务
0+阅读 · 56分钟前
《无人机蜂群通信技术研究》50页
专知会员服务
4+阅读 · 今天14:55
战力倍增器:自主武器系统与乌克兰及加沙冲突
人工智能赋能战场情报:提速决策进程
专知会员服务
3+阅读 · 7月17日
《拥抱新兴技术:面向未来军官的教育革新》
专知会员服务
7+阅读 · 7月17日
相关VIP内容
【CVPR2025】基于组合表示移植的图像编辑方法
专知会员服务
8+阅读 · 2025年4月5日
专知会员服务
30+阅读 · 2020年9月18日
相关基金
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员