Reproduce “Hexa-MoE: Efficient and Heterogeneous-aware Training for Mixture-of-Experts” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Shuqing Luo, Jie Peng, Pingzhi Li, Hanrui Wang, Tianlong Chen

HEXA-MoE introduces a framework for Mixture-of-Experts (MoE) training designed to be completely static, memory-minimized, and heterogeneous-aware. This approach achieves up to 4.3 times faster training speed and a 48% reduction in memory footprint compared to existing MoE libraries, while effectively utilizing varied computing resources.

Key insights

  • Refactoring core MoE operations into specialized, in-place expert-specific GPU kernels significantly enhances computational efficiency by removing token padding/discarding overheads.
  • A dynamic pipeline-shared cache and communication-computation overlapping effectively mitigate memory bottlenecks and communication latency in data-centric MoE training.
  • Proportional workload distribution based on inverse device latency enables efficient and adaptable MoE training across heterogeneous computing environments, improving resource utilization.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.