Full-Stack Training Optimization Practices for Ultra-Large-Scale MoE Multimodal Models

Full-Stack Training Optimization Practices for Ultra-Large-Scale MoE Multimodal Models

As large model parameter scales approach the trillion-level, Mixture of Experts (MoE) architectures and native multimodal capabilities have become core directions for industry evolution. Recently, the industry has open-sourced several ultra-large native multimodal models with total parameters exceeding 2.8 trillion. These models typically employ sparse expert architectures, featuring nearly a thousand routing experts, activating dozens of experts per token, and supporting million-level context windows alongside native image understanding. However, fine-tuning such ultra-large models on limited computing clusters not only carries a high risk of interruption but also poses extreme challenges to distributed parallel strategies, operator fusion optimization, and memory scheduling. This article systematically outlines and analyzes the full-stack training optimization practices for ultra-large-scale MoE models in a cluster environment.

I. Core Challenges in Ultra-Large Model Training

Before initiating model fine-tuning, it is essential to analyze the underlying bottlenecks across three dimensions: computation, communication, and memory.

1. Computation: Interwoven Complexity of Heterogeneous Attention Mechanisms
New-generation models often adopt hybrid attention architectures (e.g., combining linear attention and sparse attention). The state recursion and chunked computation of linear attention interweave with the sparse indexing and non-contiguous memory access of sparse attention. This heterogeneous computing pattern significantly increases the difficulty of underlying operator fusion and execution scheduling, rendering traditional single-point optimizations ineffective.

2. Communication: Load Imbalance from Fine-Grained Expert Routing
In an MoE architecture with nearly a thousand routing experts, the sparse routing mechanism easily leads to uneven token distribution. Certain popular experts may be overly targeted, resulting in small, concentrated computation batches for a single expert, while less popular experts remain idle. This not only causes communication and computation bubbles but also subjects the All-to-All communication pattern to severe congestion and long-tail latency.

3. Memory: Storage Pressure from Massive Parameters
Taking a 2.8-trillion-parameter model in BF16 format as an example, the model weights alone require approximately 5.6 TB of storage space. Combined with optimizer states, gradients, activations, and long-context KVCache, the memory demand far exceeds the physical limits of a single machine or even a standard 8-GPU cluster. Furthermore, the logit computation at the final output layer for a massive vocabulary can trigger instantaneous memory peak overflows.

II. Hybrid Parallel Architecture Design and Full-Stack Integration

To overcome the aforementioned bottlenecks, the industry typically adopts a hybrid parallel architecture based on FSDP2 (Fully Sharded Data Parallel 2) and EP (Expert Parallel).

1. FSDP2 + EP Distributed Deployment Strategy
Compared to traditional FSDP, FSDP2 offers greater flexibility in sharding granularity and memory management. Combined with EP, dense parameters can be evenly sharded across all compute cards via FSDP2, while the massive expert parameters are distributed across different nodes using the EP strategy. On a 128-card high-performance cluster, the FSDP2 + EP64 configuration has been proven to effectively achieve distributed deployment and efficient aggregation of expert parameters.

2. Full-Stack MoE Performance Optimization
For the full stack of MoE training—from token dispatch to computation recovery—deep fusion must be performed around core operators such as Permute, Unpermute, and Grouped GEMM. The following pseudo-code illustrates the overlapping design of computation and communication:

# Pseudo-code: MoE computation and communication overlap strategy
def moe_forward(tokens, router_probs):
    # 1. Token Dispatch: All-to-All communication distribution
    dispatched_tokens = all_to_all_comm(tokens, router_probs)
    # 2. Asynchronous launch of expert computation (Grouped GEMM optimization)
    expert_output = grouped_gemm_kernel(dispatched_tokens, expert_weights)
    # 3. Combine: Overlapping computation and control signal communication
    combined_output = all_to_all_comm_reverse(expert_output)
    return combined_output

By integrating high-performance All-to-All communication libraries, computation-communication overlap mechanisms, and dynamic expert load balancing technologies, the execution loop from Token Dispatch to Expert Compute and finally to Combine can be effectively streamlined.

III. Operator-Level Deep Optimization for Heterogeneous Attention and Residual Mechanisms

To address the computational challenges posed by hybrid attention and deep residual mechanisms, optimizations must start from operator compilation and underlying fusion:

1. Large Forward/Backward Fusion for Linear Attention
For linear attention mechanisms, large forward and large backward fusion operators are designed. Multiple discrete computation stages—such as multi-head computation and state updates—are consolidated into a single GPU kernel. This reduces the number of HBM read/write operations for intermediate results, significantly lowering memory access overhead.

2. AscendC Fusion Operators for Sparse Attention
For sparse attention (e.g., DSA), fusion operators like LI/SFA/SFAG are used to replace original combinations of small operators. This deep fusion scheme reduces kernel launch overhead and achieves more than a 2x performance improvement compared to traditional small operators.

3. Innovative Integration of Deep Residual Mechanisms
For innovative structures like AttnRes deep residuals, frameworks like Triton Ascend and the BiSheng compiler can be leveraged to rapidly develop and integrate forward and backward operators. This compiler-level optimization not only shortens the adaptation cycle for new structures but also ensures operator execution efficiency.

IV. Extreme Memory Management and Massive Vocabulary Optimization Strategies

1. Operator-Level Selective Caching Mechanism
To alleviate HBM memory pressure, a Swap Core operator-level selective caching mechanism can be introduced. It performs secondary caching for the computation results of specified MatMul or GMM operators, supporting hierarchical storage scheduling between HBM and DDR. Under tested configurations, this mechanism improved training performance by an average of 15%, effectively avoiding frequent memory recomputation.

2. Dual Chunking Optimization for Massive Vocabulary Loss
To address the memory peak issue caused by a million-level vocabulary at the output layer, a dual optimization scheme combining Chunk Loss and CCE Loss must be adopted:

• Chunk Loss: Chunks along the sequence dimension to avoid loading the logits of an ultra-long context into memory all at once.
• CCE Loss: Chunks computations along the vocabulary dimension, decomposing the originally massive logit matrix multiplication into multiple narrow and tall matrix operations, thereby smoothing the memory usage curve.

V. Automated Development Workflow and Infrastructure Support

Faced with the rapid iteration of complex models, traditional manual adaptation can no longer meet time-to-market requirements. Relying on an Agent-driven automated development workflow enables the integration of new models within a single day. This workflow covers the full lifecycle, from model operator parsing and automatic parallel strategy generation to accuracy verification.

At the infrastructure level, general cloud services like Baidu AI Cloud provide container platforms, object storage, and high-performance network services, offering stable underlying support for large model training. Leveraging cloud-native elastic computing scheduling and monitoring alerts allows developers to focus on optimizing model algorithms and parallel strategies, rather than getting bogged down in underlying cluster operations and maintenance.

Conclusion

Training ultra-large-scale MoE multimodal models is not merely a computing power arms race, but an extreme test of the distributed system software stack and underlying operator optimization capabilities. From FSDP2 + EP hybrid parallelism to deep operator fusion, and down to memory chunking management for massive vocabularies, systematic full-stack optimization is the key to breaking through computing bottlenecks. As automated development workflows mature and cloud-native infrastructure improves, the fine-tuning and deployment of trillion-scale models will become more widespread and efficient in the future.