Full-Stack Adaptation and Training Optimization for Ultra-Large-Scale Sparse MoE Models

Full-Stack Adaptation and Training Optimization for Ultra-Large-Scale Sparse MoE Models

As large model parameter scales push toward the trillion-level mark, native multimodal capabilities and ultra-long context lengths have become an industry consensus. Recently, a team open-sourced a new generation 3-trillion-level open-weight native multimodal model with a total parameter scale of 2.8 trillion. This model utilizes a Stable Latent MoE sparse expert architecture, featuring 896 routed experts where each token activates 16 experts. It supports a context length of up to 1 million tokens and native image understanding. Faced with the complex architecture of such ultra-large-scale models, how underlying compute platforms achieve efficient adaptation and rapid iteration has become a key technical challenge in the AI infrastructure field.

1. Analysis of Ultra-Large-Scale MoE Model Architecture Characteristics

This open-source flagship model deeply integrates sparsity and hybrid attention mechanisms in its architectural design to balance large parameter capacity with low inference and training costs.

1. Stable Latent MoE Sparse Expert Architecture: The model has a total of 2.8T parameters, but through a Top-K routing mechanism, only 16 experts are activated per token. This design significantly reduces computational complexity but simultaneously requires the system to possess extremely high memory bandwidth and communication capabilities to support dynamic addressing and weight loading for 896 experts.
2. Hybrid Attention Mechanism: The model adopts a hybrid architecture combining KDA (linear attention) and Gated DSA (gated sparse attention). Approximately three-quarters of the attention blocks use KDA linear attention to reduce the computational overhead of long sequences; meanwhile, it integrates a Top-2048 sparse attention mechanism to effectively capture ultra-long contexts (1 million tokens) and native image features.

Relying on large-scale joint reinforcement learning and system-level optimization, the model demonstrates strong capabilities in long-range autonomous programming, complex Agent tasks, 3D creation, and financial research. However, mapping these capabilities to underlying hardware for efficient training presents severe challenges across three dimensions: computation, communication, and memory.

2. Compute Infrastructure Advantages and Full-Stack Adaptation Challenges

To meet the training demands of ultra-large-scale MoE models, the underlying compute platform relies on core architectural advantages of “supernode high bandwidth, low latency, and unified memory addressing” to build a full-stack training capability system. In this process, the compute platform leveraged an Agent-driven automated development workflow to achieve day-level integration and rapid iteration for new models. This workflow covers five core stages: model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management.

However, during implementation, deep-seated issues on the computation, communication, and memory sides gradually emerged:

• Computation: The state recursion and chunked computation of KDA interleave with the sparse indexing and non-contiguous memory access of DSA. This heterogeneous computing mode significantly increases the difficulty of operator fusion, placing extremely high demands on the dynamic planning capabilities of the execution scheduling engine.
• Communication: The sparse routing of 896 experts easily triggers “expert load imbalance” issues. Popular experts receive too many tokens, while cold experts have extremely small computation batches, leading to severe All-to-All communication bottlenecks across the cluster.
• Memory: The 2.8T parameters alone require about 5.6TB of memory in BF16 format. Combined with optimizer states, gradients, and activations, a single machine’s memory is entirely insufficient, necessitating highly efficient distributed memory management strategies.

3. Computation Optimization: Heterogeneous Attention Operator Fusion and Scheduling

Under the hybrid attention architecture, traditional static computation graph optimization can no longer meet requirements. The system needs to differentiate processing based on the characteristics of KDA and DSA.

1. KDA Linear Attention Operator Fusion: KDA avoids explicitly computing the attention matrix through a state recursion mechanism. The system deeply fuses its chunked computation process, merging fine-grained operations such as matrix multiply-accumulate and state updates into a single kernel. This reduces the number of intermediate results written back to memory, thereby lowering memory access latency.
2. DSA Sparse Attention Non-Contiguous Memory Access Optimization: Addressing the non-contiguous memory access caused by the Top-2048 sparse indexing mechanism, the system introduces a Software Prefetch mechanism. This initiates memory access requests for the next block in advance while computing the current block. Simultaneously, leveraging hardware unified addressing capabilities, it bypasses traditional PCIe addressing overhead and directly maps remote memory data.

# Pseudocode: Dynamic operator scheduling logic for hybrid attention
class HybridAttentionScheduler:
    def __init__(self, layer_id, total_layers):
        # Determine whether to use KDA or DSA based on model structure
        self.use_kda = (layer_id < total_layers * 0.75)
    def forward(self, hidden_states, topk_indices=None):
        if self.use_kda:
            # KDA: Execute fused state recursion and chunked computation
            return fused_kda_linear_attention(hidden_states)
        else:
            # DSA: Execute sparse indexing and non-contiguous memory access optimization
            if topk_indices is None:
                raise ValueError("DSA requires topk_indices")
            return optimized_dsa_sparse_attention(hidden_states, topk_indices)

4. Communication Optimization: Expert Load Balancing and All-to-All Optimization

The sparse routing of 896 experts is the communication bottleneck in MoE training. When expert token distribution is uneven in All-to-All communication, communication latency is infinitely amplified by the barrel effect.

1. Dynamic Load Redistribution: The system introduces a dynamic fine-tuning mechanism at the routing layer. Before token dispatch, it quickly estimates the load of each expert and reroutes overflowing tokens for experts exceeding the computation threshold, ensuring the computation batch of each expert node remains within a reasonable range.
2. Communication and Computation Overlap: Relying on supernode high bandwidth and low latency characteristics, the system decomposes All-to-All communication into micro-batch pipelines. While experts receive the first batch of tokens to begin computation, subsequent batches of tokens continue to transfer, achieving extreme overlap between communication and computation.
3. Small Batch Operator Merging: Addressing the issue of insufficient GPU Tensor Core utilization caused by small computation batches for single experts, the scheduler pads and merges GEMM operations from multiple small batches to increase computation density.

5. Memory Optimization: Unified Addressing and Multi-Dimensional Hybrid Parallelism

The 5.6TB pure weight memory requirement means that even a single node with 8x128GB memory cannot contain it. The system must rely on multi-dimensional hybrid parallelism and unified addressing capabilities.

1. Application of Unified Memory Addressing: Relying on the underlying hardware’s unified addressing architecture, CPU memory and cross-node GPU memory form a single logical address space. The system offloads infrequently accessed expert weights or optimizer states to CPU memory or SSDs. When needed, they are fetched directly via a zero-copy mechanism, avoiding the cumbersome memory copy overhead found in traditional solutions.
2. Coupling of Expert and Tensor Parallelism: For the non-MoE parts (such as shared attention layers), Tensor Parallelism (TP) is used for splitting; for the 896 experts, Expert Parallelism (EP) is used to distribute them across different nodes. Concurrently, combined with ZeRO-3 optimization, gradients and optimizer states are sharded across the entire cluster.
3. Memory Fragmentation Governance: During dynamic routing, frequent memory allocation and release can cause severe fragmentation. The system manages the lifecycle of expert weights uniformly through pre-allocated memory pools and the Buddy System algorithm, ensuring memory stability during long-term training.

6. Automated Development Workflow: Agent-Driven Day-Level Integration

Faced with the pressure of rapid new model iteration, traditional manual adaptation can no longer meet business timeliness requirements. The system introduced an Agent-driven automated development workflow.

1. Graph Analysis and Operator Matching: The Agent automatically parses the open-source model’s computation graph, identifies structures like KDA, DSA, and MoE Router, and performs semantic matching within a pre-configured operator library.
2. Code Generation and Verification: For missing or poorly performing operators, the Agent invokes code generation tools to produce an initial version of CUDA/C++ code and verifies accuracy and performance on a small-scale cluster.
3. Full-Stack Automatic Compilation: Using JIT (Just-In-Time) compilation technology, hybrid parallelism strategies and communication optimization operations are automatically inserted into the computation graph, ultimately generating an executable training script.

Through this automated closed loop, model adaptation work that originally took weeks is compressed to the day level, significantly improving the compute platform’s response speed to new models.

7. Enhancing the Training Guarantee System via General Cloud Services

Beyond underlying computing power and adaptation frameworks, complete model training cannot be separated from general cloud service support. When constructing ultra-large-scale training tasks, it is typically necessary to combine object storage for low-cost hosting of massive multimodal data, and utilize container platforms for elastic scheduling and environment isolation of training tasks. Meanwhile, message queues facilitate asynchronous state synchronization between training nodes, while logging services and monitoring alert platforms collect core metrics such as cluster compute utilization, communication latency, and memory water levels in real time. Upon detecting expert load skew or node failure, the monitoring system triggers alerts within seconds and automatically initiates recovery processes, ensuring the continuity and stability of trillion-level model training.

Conclusion

From heterogeneous operator fusion on the computation side, to dynamic load balancing on the communication side, and unified addressing with multi-dimensional parallelism on the memory side, the training adaptation of ultra-large-scale sparse MoE models is an extremely complex systems engineering endeavor. By building full-stack optimization capabilities and introducing an Agent-driven automated development workflow, the underlying compute platform successfully achieved day-level integration and efficient training for the 2.8T parameter flagship model. This not only validates the advancement of current high-bandwidth, low-latency compute architectures but also provides a highly valuable practical paradigm for the infrastructure construction of future 10-trillion-level or even larger scale models.