Full-Link Optimization Practices for Training and Fine-Tuning Ultra-Large-Scale MoE Multimodal Models

As large model parameter scales approach the trillion-level, native multimodal capabilities and long context lengths have become industry standards. However, training sparse Mixture-of-Experts (MoE) models with 2.8 trillion total parameters and hundreds of routed exerts introduces unprecedented pressure on underlying training systems across compute, communication, and memory dimensions. This article strips away specific product attributes to deeply explore full-link optimization practices for distributed training and fine-tuning of ultra-large-scale MoE models from a general engineering perspective.

I. Underlying System Challenges of Ultra-Large-Scale MoE Models

A typical 3-trillion-level open-weight native multimodal model usually adopts a Stable Latent MoE sparse expert architecture, such as containing 896 routed experts where only 16 experts are activated per token. Additionally, the model introduces a hybrid architecture of KDA (linear attention) and Gated DSA (gated sparse attention) to support million-token contexts. While this innovative structure enhances model capabilities, it brings three major challenges to the underlying infrastructure:

1. Compute: Operator Fusion Difficulties Caused by Heterogeneous Attention
   The state recursion and block computation of KDA are deeply intertwined with the sparse indexing and non-contiguous memory access of DSA. In traditional execution flows, numerous fragmented small operators lead to frequent kernel launch overheads and memory access bottlenecks, significantly increasing the difficulty of operator fusion and execution scheduling.

2. Communication: Load Imbalance Triggered by Sparse Routing
   The sparse routing mechanism of 896 experts is prone to the “hot expert” phenomenon during runtime, leading to severe token distribution imbalances. This not only causes small compute batches for single experts and low GEMM compute efficiency but also triggers data skew across nodes during All-to-All communication.

3. Memory: Memory Walls from Massive Parameters and Long Contexts
   With a parameter scale of 2.8T, the BF16 weights alone consume about 5.6TB of memory. Coupled with the massive KV Cache and optimizer states brought by million-token contexts, traditional single-node or simple data parallelism can no longer meet deployment requirements.

II. Hybrid Parallel Architecture Design: Breaking Memory and Communication Limits

To address these challenges, the mainstream industry approach is to build full-link training capabilities, encompassing model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management. At the parallel architecture level, a hybrid parallel solution based on FSDP2 (Fully Sharded Data Parallel 2) and EP (Expert Parallel) has become critical.

1. FSDP2+EP Distributed Deployment and Aggregation

Compared to the first-generation solution, FSDP2 offers more flexible sharding strategies and communication scheduling, enabling fine-grained control over non-parameter components. Combined with EP (Expert Parallel), the system can distribute 896 experts across large-scale multi-node clusters.
For example, LoRA fine-tuning validation was conducted on a 128-GPU super-node cluster using the FSDP2+EP64 architecture. This architecture not only achieves efficient distributed deployment of expert parameters but also ensures the efficiency of cross-node parameter aggregation through unified memory addressing and a network topology with high bandwidth and low latency.

2. Agent-Driven Automated Integration

Faced with the rapid iteration of new model structures, traditional manual adaptation is time-consuming and labor-intensive. Modern training frameworks often introduce Agent-driven automated development workflows. By parsing the model computation graph, the Agent can automatically identify fusable operators, recommend parallel strategies, and generate distributed deployment scripts, making day-level integration and rapid iteration of new models possible.

III. Deep Operator Fusion and High-Performance Scheduling

To unleash the computing power of underlying hardware, deep customization targeting the model’s unique computation graph is required.

1. Fused Operator Design for KDA and DSA

For KDA linear attention, the system needs to design large forward and large backward fused operators, integrating originally scattered computation stages (such as projection, state update, and output projection) into a single kernel. This significantly reduces the offloading of intermediate results and kernel launch overhead.
For DSA sparse attention, composite fused operators based on LI/SFA/SFAG (implemented via underlying programming paradigms like AscendC) can be integrated into the model. By deeply fusing sparse index extraction with attention computation, performance can be improved by more than 2x compared to the original small operator mode.

2. Compiler-Level Support for Innovative Structures

For innovative mechanisms like AttnRes (Attention Deep Residual), the underlying system can rely on general compilation frameworks like Triton Ascend and the BiSheng compiler. By utilizing compiler Pass mechanisms, the development of forward and backward operators and automatic differentiation integration can be rapidly completed, lowering the engineering threshold for implementing innovative structures.

3. Full-Link Operator Optimization for MoE

The core execution flow of MoE includes Token Dispatch, Expert Compute, and Combine. To address performance bottlenecks in this link, deep fusion must be performed around core operators such as Permute, Unpermute, and Grouped GEMM. Combined with high-performance All-to-All communication and communication-computation overlap technology, local expert computation can be initiated while waiting for cross-node data transmission. Paired with dynamic load balancing strategies, this completely clears full-link execution bottlenecks.

IV. Hierarchical Storage and Large Vocabulary Memory Optimization

After resolving compute and communication, memory management is the final threshold determining whether a model can actually run.

1. Swap Core: Operator-Level Selective Caching

To alleviate pressure on High Bandwidth Memory (HBM), the system can introduce a Swap Core mechanism. This is an operator-level selective caching strategy that performs secondary caching on the computation results of specified MatMul or GMM operators, supporting hierarchical storage between HBM and DDR (Double Data Rate synchronous dynamic random-access memory). By intelligently predicting data access frequency, cold data is offloaded to DDR while hot data remains in HBM. Tests show that this mechanism can improve training performance by an average of 15% under the current configuration.

2. Dual Optimization of Chunk Loss and CCE Loss

A large vocabulary is another memory killer for ultra-large-scale models. When computing cross-entropy loss, the massive Logits matrix easily leads to Out of Memory (OOM) errors.
A general solution to this problem is dual chunking optimization along the sequence and vocabulary dimensions:

• Chunk Loss: Chunking along the sequence dimension, splitting long sequences into multiple sub-sequence fragments to compute loss separately, thereby reducing the instantaneous memory peak.
• CCE Loss: Chunked computation along the vocabulary dimension to avoid loading the entire vocabulary’s weight matrix at once.
  The following is a simplified logic illustration of Chunk Loss:

def compute_chunk_loss(logits, labels, chunk_size=1024):
    total_loss = 0.0
    seq_len = logits.shape[0]
    for start_idx in range(0, seq_len, chunk_size):
        end_idx = min(start_idx + chunk_size, seq_len)
        # Extract logits and labels for the current chunk
        chunk_logits = logits[start_idx:end_idx, :]
        chunk_labels = labels[start_idx:end_idx]
        # Compute local cross-entropy and accumulate
        local_loss = cross_entropy(chunk_logits, chunk_labels)
        total_loss += local_loss * (end_idx - start_idx)
    return total_loss / seq_len

Through this dual chunking strategy, the model effectively circumvents the memory surge caused by the combination of large vocabularies and long contexts.

V. Fine-Tuning Practices and Convergence Validation

Theoretical optimizations must ultimately be validated in engineering practice. Based on the aforementioned full-link optimization technologies, LoRA fine-tuning practices were conducted on a 2.8T parameter model using a SuperPoD cluster composed of 128 compute nodes.

During fine-tuning, parameter sharding was executed via the FSDP2+EP64 strategy. Combined with operator fusion and Swap Core memory management, the system successfully carried out parallel training at the 100-GPU scale. Regarding convergence metrics, the initial loss of randomly initialized training was approximately 13.5. After over 400 steps of iterative training, the loss smoothly converged to around 1.7. This data not only validates the feasibility of training large-scale sparse MoE models on distributed clusters but also proves the stability and high efficiency of the underlying hybrid parallelism and memory optimization strategies.

Conclusion

Training and fine-tuning ultra-large-scale MoE models is essentially a stress test on the limits of underlying compute, storage, and communication architectures. By building a full-link optimization system covering hybrid parallelism, deep operator fusion, communication-computation overlap, and hierarchical memory management, developers can effectively break through the engineering barriers of trillion-parameter models. As automated integration workflows and compiler technologies continue to evolve, the iteration cycle for ultra-large models will be significantly shortened in the future, providing a more solid computing foundation for frontier applications such as complex Agent tasks and deep knowledge work.