moe-engine: Sparse MoE Training Infrastructure for 10k-GPU Clusters

Project Basic Information

• Maintainer: Mattral
• Source Code: Composed-Mixture-of-Experts-Engine

Core Positioning

moe-engine is a sparse MoE training runtime infrastructure for ultra-large-scale GPU clusters, designed specifically for continuous node failure scenarios in 10k+ GPU clusters, aiming to achieve training stability without human intervention.

Key Features

• Supports 4D parallel strategy (DP+EP+TP+PP)
• Asynchronous hierarchical checkpoint mechanism
• TorchElastic fault tolerance recovery
• Fused Triton routing kernel optimization

In the field of large language model training, sparse MoE technology is an important path to break the computing power bottleneck. However, when scaling to 10k-GPU level, it faces a core challenge: node failures are no longer occasional events but a persistent norm. How to maintain end-to-end stability of the training process has become a core issue in infrastructure design.

moe-engine is born to address this challenge—it is not a model implementation but a production-grade runtime. Its core constraint is: when nodes keep dying in a 10k-GPU cluster, the system must remain alive for training without human intervention.

moe-engine uses a 4D parallel strategy to build a distributed training grid:

1. Data Parallelism (DP)：Implements fine-grained parameter sharding based on FSDP2, uses PyTorch 2.5+ DTensor abstraction, supports mixed-precision training, and balances memory efficiency and performance.
2. Expert Parallelism (EP)：Each EP rank owns part of the experts, executes all-to-all operations (token distribution/aggregation) via independent CUDA streams to overlap computation and communication.
3. Tensor Parallelism (TP)：Expert FFN uses column parallelism (gating/dimension up-projection) and row parallelism (dimension down-projection + all-reduce) strategies, with correctness verified in 2-rank environments.
4. Pipeline Parallelism (PP)：Uses 1F1B interleaved scheduling (warm-up → steady state → drain) to maximize pipeline utilization.

Fused Triton Routing Kernel

Routing is the bottleneck of MoE; the traditional process requires 3 HBM round trips. moe-engine's fused kernel compresses this to a single memory traversal:

• SRAM block computation (64×64 blocks) to complete matrix multiplication, softmax, top-K selection, and renormalization
• When K∈{1,2,4} and E≤256, selection sort outperforms bitonic sort to avoid memory bank conflicts
• For H=4096 and E=64 configuration, memory traffic is reduced by ~2.7x

Asynchronous Hierarchical Checkpoint

Non-blocking design ensures training is not blocked:

1. Sync layer: D2H copy of SHARDED_STATE_DICT snapshot (tens of milliseconds)
2. Host layer: Background thread writes to NVMe (O_DIRECT + 256MB chunks)
3. Persistent layer: Mirror to S3/MinIO after atomic renaming

TorchElastic Fault Tolerance Mechanism

Node failure recovery process:

1. Heartbeat detection to identify dead ranks
2. Evict faulty nodes and poll to reassign expert ownership
3. Restore state from the latest checkpoint
4. Continue training automatically without restart

Coordination backend: etcd for clusters with over 100 nodes, c10d for small-scale clusters.

Routing Kernel Throughput (CPU Reference Path)

Tokens (N)	Hidden (H)	Experts (E)	Top-K	Latency	Throughput
512	256	16	2	0.04 ms	12.8M tok/s
1024	512	32	2	0.12 ms	8.5M tok/s
2048	1024	64	2	0.47 ms	4.4M tok/s
4096	2048	64	4	1.83 ms	2.2M tok/s

Key Verifications

• Token Conservation: In 100 random seed tests, sum(dispatch_cnt) == N×K is strictly maintained with zero violations
• Load Balance: Default initialization imbalance ratio is 1.12 (max/avg), which can be optimized to 1.05 with z-loss regularization (1e-3)

1. Necessity of Fused Kernels: The bottleneck of MoE routing is memory bandwidth; fused operations reduce HBM round trips, and the benefits far exceed development costs
2. Value of Independent CUDA Streams: EP's all-to-all and FFN computation are independent; scheduling to different streams allows overlapping execution. With EP=8 and NVLink configuration, communication overhead is reduced by ~40%
3. Atomic Checkpoint Design: Partial checkpoints are catastrophic in distributed training; atomic renaming ensures integrity and should be promoted to all distributed persistence scenarios

Existing Limitations (v0.2 Version)

• Chaos testing: Node failure recovery (Scenario A) pass rate is ~85%; Gloo backend connectFullMesh has race conditions in container environments

Future Plans

• v0.3 version integrates Nsight/CUPTI performance analysis
• Continuously access clusters to obtain real multi-node performance data

moe-engine provides an excellent reference implementation for MoE training infrastructure. It proves that through well-designed 4D parallelism, fused kernels, asynchronous checkpoints, and automatic fault tolerance mechanisms, high-reliability end-to-end training can be achieved in 10k-GPU clusters. For teams building or optimizing large-scale MoE training systems, this project is a codebase worth in-depth study.