Scaling MoE Architectures to Trillion Parameters: Long-Context and Sparse Training Stability Practices

Scaling MoE Architectures to Trillion Parameters: Long-Context and Sparse Training Stability Practices

As large models push toward trillion-parameter scales, Mixture-of-Experts (MoE) architectures have become the mainstream choice for balancing compute cost and model performance. However, maintaining training stability under extreme sparsity, breaking through the representation bottlenecks of deep Transformers, and efficiently handling million-token contexts remain common industry challenges. Recently, the industry proposed an MoE multimodal architecture with 2.8 trillion total parameters and approximately 104 billion activated parameters. This architecture natively integrates visual multimodal capabilities and achieves key breakthroughs in compute efficiency and reasoning through multi-stage pretraining and million-token-level reinforcement learning post-training. This article distills its technical paradigms and engineering reuse value across three dimensions: core architecture design, long-context computation optimization, and MoE load balancing.

1. Attention Mechanism Redesign: Breaking the Bottleneck of Long-Context Computation

When processing ultra-long contexts, the self-attention mechanism of traditional Transformers introduces quadratic growth in memory and compute overhead. Although linear attention mechanisms can alleviate this to some extent, in practical engineering deployments—especially when optimized for mainstream underlying compute architectures such as CUDA/Tensor Cores—they often face severe numerical overflow and poor parallel efficiency.

1.1 Numerical Underflow and Diagonal Block Computation Challenges in Linear Attention
In conventional linear attention or Delta Net-like computation mechanisms, intra-chunk computation requires frequent handling of the reciprocal of cumulative decay factors. Under limited precision (e.g., BF16), as sequence length increases, the reciprocal of the decay factor can easily exceed the representable numerical range, causing severe overflow or underflow. To avoid this, traditional implementations often force the diagonal portion of data chunks to go through an explicit “position-pair path” for special computation. This branching logic breaks the continuity of large block matrix multiplication, preventing the underlying hardware from fully utilizing Tensor Core dense matrix acceleration units, which significantly drags down overall throughput.

1.2 Lower-Bounded Decay Parameterization
To completely eliminate the aforementioned special computation branches, this architecture introduces a lower-bounded decay parameterization mechanism. The core idea is to reconstruct the originally unbounded decay function (such as Negative-Softplus) into a Scaled Sigmoid function with a strict lower bound.
Specifically, by fixing a minimum decay threshold (e.g., set to -5), the reciprocal scaling factors within any 16-token tile are locked within an absolutely safe numerical range. This minor engineering change brings enormous benefits: it completely eliminates the special handling branch for diagonal blocks, allowing all causal computation blocks (including both diagonal and non-diagonal parts) to directly invoke dense GEMM (general matrix multiplication) operators. This not only pushes the utilization of underlying hardware matrix multiply-accumulate units to their limit, greatly improving throughput, but also naturally supports direct extrapolation to million-level contexts in NoPE (No Positional Encoding) scenarios.

2. MoE Architecture Innovation: Training Stability and Load Balancing Under Extreme Sparsity

In an extremely sparse MoE architecture containing nearly a thousand routed experts, ensuring dynamic load balance across experts is the core challenge determining model convergence quality. Traditional approaches typically rely on auxiliary loss functions to constrain routing strategies, but this not only increases optimization complexity but also easily introduces gradient interference.

2.1 Auxiliary-Loss-Free Fast Dynamic Balancing Mechanism
This architecture adopts an auxiliary-loss-free load balancing strategy for ultra-large-scale MoE—a quantile balancing mechanism. Traditional auxiliary loss methods typically apply soft constraints at a global granularity, converging slowly and making it difficult to precisely control individual expert loads. In contrast, the quantile balancing mechanism achieves hard constraints on expert loads with rapid response by computing and dynamically adjusting the distribution quantiles of routing weights.
At each training step, the system does not need to compute and backpropagate complex auxiliary gradients; instead, it directly performs quantile fine-tuning based on the current reception frequency of each expert. This mechanism maintains rapid dynamic balance even in extreme sparsity scenarios with thousands of experts, avoiding the “winner-takes-all” expert collapse phenomenon and greatly improving the stability of large-scale pretraining.

2.2 Stable Latent MoE and Activation Function Design
On the expert network side, the architecture employs a Stable Latent MoE framework combined with the SiTU-GLU activation function. Conventional GLU (Gated Linear Units) tend to suffer from uncontrollable activation value ranges leading to gradient explosion or vanishing when scaling to ultra-large parameters. SiTU-GLU introduces smoother and bounded nonlinear transformations, ensuring forward propagation numerical stability under deep network layers and massive expert stacking.

3. Cross-Layer Information Sharing: Breaking Deep Transformer Representation Bottlenecks

As model depth increases, traditional Transformer architectures tend to suffer from “representation over-smoothing,” where feature representations across different layers in deep networks become homogenized, leading to wasted model capacity and degraded generalization.

3.1 Block Attention Residuals
To address deep representation degradation, this architecture introduces a Block Attention Residual mechanism that enables direct cross-layer information sharing. Traditional residual connections only establish shortcuts between adjacent layers, while Block AttnRes builds information channels across multiple layers, allowing high-frequency detail features from lower layers to be directly injected into deep abstract representations.
This not only breaks the representation bottleneck of deep Transformers but also significantly reduces memory access overhead and inter-layer communication latency during backpropagation. Combined with Gated Multi-Head Latent Attention, the model can maintain ultra-long context memory depth while effectively filtering redundant information, ensuring logical coherence in Agent multi-step long-range reasoning.

4. Training Paradigm Evolution: From Multi-Stage Pretraining to Agentic RL

4.1 Progressive Context Window Extension
During pretraining, the model does not directly train with ultra-long contexts but adopts a progressive extension strategy of 8K -> 64K -> 256K -> 1M. This multi-stage strategy first solidifies the model’s foundational language knowledge and logical capabilities under short contexts, then gradually introduces long-sequence data, enabling the model to smoothly transition to million-token context extrapolation. This not only reduces initial training compute costs but also avoids interference from long-sequence noise on foundational capability training.

4.2 Agentic RL Post-Training for Long-Range Reasoning
After pretraining, the model incorporates million-token-level Agentic RL (agent reinforcement learning) post-training. Traditional RLHF mainly focuses on aligning human preferences and safety, while Agentic RL shifts the focus to tool invocation, multi-turn long-range planning, and complex task decomposition. By simulating realistic agent interaction trajectories in ultra-long context environments, the model learns to extract key clues from vast context memories and maintain task goal stability. This provides an effective path for addressing the long-range reasoning deficiencies of large models in complex enterprise applications such as codebase-level code generation and deep long-document analysis.

5. Engineering Reuse Guide and Cloud-Native Adaptation Outlook

The above technical innovations have high reuse value for industry large model development. Lower-bounded decay parameterization and quantile balancing mechanism both feature “minimal engineering changes, maximum benefits” and can be directly integrated into various mainstream training frameworks.

For enterprise-level large model deployment, training and inference of such ultra-large-scale architectures impose stringent requirements on underlying compute infrastructure. In cloud-native practice scenarios, high-performance container platforms and elastic compute resource pools are typically needed to carry distributed training of hundred-billion-parameter models. For example, when utilizing general object storage for high-speed distribution of massive multimodal datasets, high-bandwidth, low-latency RDMA network architectures are needed to ensure gradient synchronization under thousand-GPU training; meanwhile, mature monitoring, alerting, and logging services are required to track the load status of thousands of routed experts and the operator performance of underlying Tensor Cores in real time. Only by deeply integrating cutting-edge MoE architecture innovations with standardized cloud-native infrastructure can the closed loop from theoretical breakthroughs to industrial deployment of ultra-large-scale models truly be achieved.