Scaling MoE Training: Key Techniques for Trillion-Parameter Model Adaptation

Introduction: Engineering Challenges of Open-Weight Multimodal Models

As the scale of open-weight models continues to rise, trillion-parameter native multimodal models are becoming the new technological frontier. These models typically share three defining characteristics: parameter counts reaching the trillion scale, a sparse Mixture-of-Experts (MoE) architecture, and support for ultra-long contexts along with native image understanding. The combination of these factors means that model adaptation is no longer a simple matter of operator substitution, but a full-stack systems engineering effort spanning computation, communication, and GPU memory.

Architecturally, next-generation flagship models commonly adopt a Stable Latent MoE sparse expert structure: total parameters can approach 2.8 trillion, with nearly 900 routed experts, yet only a dozen or so experts are activated per token. This “large parameter, small activation” design offers clear inference efficiency advantages, but it places extremely high demands on training systems in terms of resource scheduling and communication aggregation.

On the attention front, a hybrid attention architecture has become the mainstream approach: approximately three-quarters of the attention blocks use linear attention (e.g., KDA-style state-recurrent structures), while the remainder incorporate Top-K sparse attention. Linear attention replaces full attention matrices with state recurrence and chunked computation, significantly reducing the computational complexity of long sequences. Sparse attention, meanwhile, relies on index-based filtering and non-contiguous memory access. The interplay of these two types makes operator fusion and execution scheduling considerably more difficult.

On top of this, the model must support contexts of up to one million tokens and native image understanding, while maintaining capability in scenarios such as long-horizon autonomous coding, complex agent tasks, 3D creation, financial research, and deep knowledge work. This means the training system must not only “run,” but run reliably and support rapid iteration.

Three Core Challenges: Computation, Communication, and Memory

Computation: Fusion Challenges for Heterogeneous Attention Structures

The state-recurrent nature of linear attention introduces serial dependencies: each time step’s state update depends on the previous step’s result, and chunked computation requires a two-level reduction within and across blocks. Sparse attention introduces index computation, sparse gather/scatter operations, and non-contiguous memory access. The execution patterns of these two operator types are fundamentally different. If conventional small-operator composition is used, frequent kernel launch overhead and intermediate tensor spills to memory will severely degrade end-to-end performance.

The engineering response is “large-scale fusion”: consolidating multiple forward computation stages into a single kernel, and similarly combining stages in the backward pass to reduce launch count and memory round-trips. For the sparse attention path, fused AscendC-style operators (e.g., index loading, sparse focused attention, and their variants) replace small-operator pipelines, achieving more than 2x performance improvement in practice compared to previous implementations.

Communication: Load Imbalance in Sparse Routing

Sparse routing across nearly 900 experts naturally leads to token distribution imbalance: popular experts receive far more tokens than the average, while less-requested experts see small batch sizes. The direct consequences are high variance in per-expert batch sizes and imbalanced inter-card communication, often resulting in a long-tail effect where some devices wait while others idle. If the parallel strategy is not designed carefully, the All-to-All communication in expert parallelism can become the overall throughput bottleneck.

Memory: Combined Pressure from Weights and Optimizer States

Storing 2.8 trillion parameters in BF16 alone requires roughly 5.6 TB. When optimizer states, gradients, and activations are added, single-GPU memory is completely insufficient. This requires the system to have comprehensive memory management capabilities, including parameter sharding, dynamic aggregation, and activation recomputation. Critically, the sharding strategy must be co-designed with the communication topology to avoid “saving memory at the cost of losing bandwidth.”

Building Full-Stack Training Capabilities

To address these challenges, engineering practice typically builds capability across five dimensions: model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management. An agent-driven automated development workflow enables day-level onboarding and rapid iteration for new models.

Hybrid Parallelism: Combining FSDP2 with Expert Parallelism

For trillion-parameter MoE models, the industry standard is a hybrid architecture of “data/parameter sharding + expert parallelism.” FSDP2 handles sharding of parameters, gradients, and optimizer states, while expert parallelism (EP) distributes expert parameters across devices and aggregates them efficiently. The key to combining the two is ensuring that the sharding granularity aligns with the communication patterns of expert routing, so that All-to-All communication domains stay within high-bandwidth interconnects as much as possible.

Taking a 128-card supernode as an example, a combination of FSDP2 + EP64 can be used for LoRA fine-tuning validation: expert parameters are distributed across 64 expert parallel groups, while other dimensions are sharded by FSDP2. LoRA fine-tuning is advantageous here because it only trains low-rank adapter parameters, dramatically reducing memory usage for optimizer states and gradients, and making a trillion-parameter model practically trainable.

The following pseudocode illustrates the structural relationship of the parallel configuration:

# Pseudocode: structural representation of hybrid parallel configuration
parallel_config = {
    "fsdp_shard_dim": "param",          # shard params/gradients/optimizer states
    "expert_parallel_size": 64,          # expert parallel group size
    "sequence_parallel": True,           # long-sequence splitting
    "activation_checkpointing": "selective",  # selective recomputation
    "lora": {"rank": 16, "target": "attention_and_expert_proj"}
}
# Pseudocode: key stages of a single training step
def train_step(batch):
    with parallel_context(parallel_config):
        logits = model_forward(batch)     # fused operators for forward
        loss = compute_loss(logits, batch.labels)
        loss.backward()                   # fused operators for backward
        all_reduce_experts()              # aggregate expert parameters
        optimizer_step()                  # update only LoRA parameters

Operator Acceleration: Fusion Design for Heterogeneous Attention

For linear attention, large fused forward and backward operators consolidate state recurrence, chunked reduction, and output computation into a single kernel, reducing intermediate tensor spills. For sparse attention, fused AscendC operators—including index loading, sparse focused attention, and their variants—handle model integration by merging distributed index computation, sparse memory access, and attention computation, thereby lowering scheduling overhead. Together, these two fusion families form the computational foundation for optimization.

Communication and Memory: Co-Design Rather Than Silos

The focus of communication optimization is mitigating imbalance in expert routing: token reordering and capacity factor tuning bring the batch sizes of long-tail experts back into an acceptable range, while overlapping communication with computation hides All-to-All behind expert computation. Memory management depends on fine-grained control of sharding, recomputation, and the timing of parameter aggregation—aggregate when needed, shard whenever possible, and avoid keeping full weights resident.

The supernode architecture’s high bandwidth, low latency, and unified memory addressing provide the physical foundation for this co-design: unified addressing simplifies cross-device parameter access, high bandwidth keeps expert aggregation communication costs manageable, and low latency enables communication-computation overlap.

Practical Recommendations and Tuning Advice

1. Set the parallel strategy first, then tune operators: The parallel strategy determines communication patterns and memory ceilings. Establish the FSDP2 and EP ratio first, then perform operator fusion along bottleneck paths.
2. Use LoRA to validate trainability: On a trillion-parameter model, first run LoRA fine-tuning end-to-end, then gradually move to full-parameter training. This significantly reduces trial-and-error cost.
3. Focus on the long tail, not the mean: Performance bottlenecks in sparse routing often appear on the most imbalanced experts. Monitor quantiles of per-expert batch sizes, not just averages.
4. Validate fused operators thoroughly: After fusion, compare numerical precision and end-to-end throughput to avoid “faster single operators, slower overall network” false optimization.
5. Automate model onboarding: An agent-driven automated workflow that standardizes model structure parsing, operator mapping, and parallel configuration generation enables day-level onboarding and rapid iteration for new models.

Conclusion

Adapting trillion-parameter sparse MoE models for training is essentially about finding the optimal engineering trade-off among computational heterogeneity, communication imbalance, and rigid memory constraints. Fused operator design for linear and sparse attention addresses “computing fast,” the hybrid architecture of FSDP2 and expert parallelism addresses “fitting into memory,” and communication reordering and overlap address “transferring efficiently.” Only when all three work in concert can stable training and rapid iteration be achieved for multimodal models with million-token contexts. For teams looking to deploy such models on their own compute platforms, this “full-stack + automated” methodology offers strong reusable value.