MemFine: Memory-Aware Fine-Grained Scheduling for MoE Training
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866918285384089600 |
|---|---|
| author | Zhao, Lu Shi, Rong Zhang, Shaoqing Chen, Yueqiang He, Baoguo Sun, Hongfeng Yin, Ziqing Su, Shangchao Cui, Zhiyan Dong, Liang Li, Xiyuan Wang, Lingbin He, Jianwei Ma, Jiesong Huang, Weikang Tong, Jianglei Gao, Dongdong Zhang, Jian Tian, Hong Shen, Hui Luo, Zongtai Sun, Zhaoqun Niu, Hongxing Sun, Yue |
| author_facet | Zhao, Lu Shi, Rong Zhang, Shaoqing Chen, Yueqiang He, Baoguo Sun, Hongfeng Yin, Ziqing Su, Shangchao Cui, Zhiyan Dong, Liang Li, Xiyuan Wang, Lingbin He, Jianwei Ma, Jiesong Huang, Weikang Tong, Jianglei Gao, Dongdong Zhang, Jian Tian, Hong Shen, Hui Luo, Zongtai Sun, Zhaoqun Niu, Hongxing Sun, Yue |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2511_21431 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | MemFine: Memory-Aware Fine-Grained Scheduling for MoE Training Zhao, Lu Shi, Rong Zhang, Shaoqing Chen, Yueqiang He, Baoguo Sun, Hongfeng Yin, Ziqing Su, Shangchao Cui, Zhiyan Dong, Liang Li, Xiyuan Wang, Lingbin He, Jianwei Ma, Jiesong Huang, Weikang Tong, Jianglei Gao, Dongdong Zhang, Jian Tian, Hong Shen, Hui Luo, Zongtai Sun, Zhaoqun Niu, Hongxing Sun, Yue Distributed, Parallel, and Cluster Computing 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. |
| title | MemFine: Memory-Aware Fine-Grained Scheduling for MoE Training |
| topic | Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2511.21431 |