FoldMoE: Efficient Long Sequence MoE Training via Attention-MoE Pipelining
Guichao Zhu, Lintian Lei, Yuhao Qing, Yichao Fu, Fanxin Li, Dong Huang, Zekai Sun, Heming Cui
Abstract
Training LLMs with Mixture-of-Experts (MoE) architecture on long sequences poses significant challenges due to the all-to-all communication bottleneck of expert parallelism. While existing approaches attempt to hide the communication costs in computation through token-level pipelining within MoE layers, their effectiveness is limited by the insufficient computation. We present FoldMoE, a high-performance MoE training system that enables token-level overlapping across entire Transformer blocks through novel attention-MoE pipelining. We propose an efficient pipeline schedule, and a novel token buffering design to decouple attention and MoE layer partitioning, along with a time-uniform micro-batching strategy for enhanced efficiency. Evaluations on GPT-MoE models with sequences up to 32K tokens show that FoldMoE achieves up to 1.49x and 2.72x speedup over state-of-the-art token-level overlapping and non-overlapping baselines respectively.
BibTeX
@inproceedings{zhu-etal-2025-foldmoe,
title = "{F}old{M}o{E}: Efficient Long Sequence {M}o{E} Training via Attention-{M}o{E} Pipelining",
author = "Zhu, Guichao and
Lei, Lintian and
Qing, Yuhao and
Fu, Yichao and
Li, Fanxin and
Huang, Dong and
Sun, Zekai and
Cui, Heming",
editor = "Che, Wanxiang and
Nabende, Joyce and
Shutova, Ekaterina and
Pilehvar, Mohammad Taher",
booktitle = "Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
month = jul,
year = "2025",
address = "Vienna, Austria",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2025.acl-long.186/",
doi = "10.18653/v1/2025.acl-long.186",
pages = "3705--3717",
ISBN = "979-8-89176-251-0"
}