Large-Scale Sparse MoE: Key Technologies and Distributed Training Practices

Engineering Challenges and Technological Evolution in Ultra-Large MoE Models

As Large Language Model (LLM) parameter counts surpass the trillion-level mark, traditional dense Transformer architectures face significant bottlenecks in both training costs and inference latency. The sparse Mixture of Experts (MoE) architecture, through conditional computation mechanisms, has become the mainstream technical solution for building ultra-large-scale models in the industry, as it maintains model capacity while controlling actual computational overhead. However, when a model’s total parameters reach 2.8 trillion with activated parameters exceeding the 100 billion mark, issues such as training stability, long-context modeling efficiency, and distributed system communication overhead become exceptionally prominent.

This article systematically outlines the architectural design philosophies and engineering practices of a recently released multimodal MoE model. The model features 2.8 trillion total parameters and 104 billion activated parameters, natively supporting visual multimodal inputs and an ultra-long context of 1 million tokens. Its core innovations cover attention mechanism modifications, MoE sparse routing optimizations, multi-stage pre-training strategies, and reinforcement learning post-training methods tailored for agentic tasks. The following sections will provide a detailed analysis across multiple dimensions and discuss the reusable value of these general technical principles.

1. Low-Cost, High-Yield Modifications to Attention Mechanisms

The computational complexity of traditional causal attention scales quadratically with sequence length. When processing a long text of 1 million tokens, both computational volume and memory consumption become unbearable. The industry typically adopts linear attention or sparse attention patterns to reduce complexity, but these solutions often face a trade-off between numerical stability and hardware utilization during practical engineering implementation.

1.1 Design Logic of KDA’s Lower-Bounded Decay Parameterization

In linear attention mechanisms, the decay factor controls the forgetting rate of historical information. The standard approach uses a lower-bounded decay parameterization to map the decay factor within a safe numerical range. Specifically, traditional implementations often use an unbounded Negative-Softplus function to map raw scores into the log space, which performs well on small-scale models. However, when sequence chunking is applied to physical computing units (such as Tensor Cores), the reciprocal of the cumulative decay factor within each 16-token chunk may exceed the numerical representation range under BF16 precision. This leads to floating-point underflow or overflow, forcing the computing units to fall back to explicit position-pair traversal paths, thereby failing to leverage the hardware acceleration capabilities of matrix multiplication.

KDA (Kimi Decay Attention) replaces the Log-decay with a lower-bounded scaled Sigmoid function, fixing the minimum value:

decay = g_min * Sigmoid(activation)

Where g_min is fixed at -5. This parameterization limits the reciprocal scaling factor within any 16-token chunk to a safe numerical interval, fundamentally eliminating the risk of numerical overflow. This allows all computation paths to be mapped to high-throughput matrix multiplication instructions. The brilliance of this design lies in its minimal engineering change—merely modifying the activation function form—yet it drastically improves the implementation efficiency of linear attention operators in underlying compute libraries like CUDA and CUTLASS.

From a reuse perspective, this technical solution is highly universal. Any distributed training project involving linear attention or State Space Models (SSMs) can adopt this approach of “redesigning parameterization for hardware precision constraints.” It reminds developers that during operator optimization, one must not only focus on reducing algorithmic FLOPs (floating-point operations) but also pay close attention to the matching relationship between numerical ranges and hardware instruction sets.

1.2 Synergistic Effects of Gated MLA and Attention Residuals

Gated Multi-head Latent Attention (Gated MLA) is another critical innovation. Standard MLA reduces KV cache size through low-rank compression, while the gating mechanism introduces learnable gating scalars for each attention head, allowing it to adaptively adjust the contribution weights of different representation subspaces. This design reduces memory pressure while endowing the model with more flexible representation capabilities.

Attention Residuals break the inherent bottleneck of information flow attenuation across deep Transformer layers. In specific implementation, Block AttnRes allows attention outputs to be passed across layers, enabling local patterns extracted in lower layers to directly supplement high-level semantic representations. This approach not only improves the representation capability of deep networks but also reduces unnecessary intermediate layer computations—some layers can directly reuse the attention outputs of preceding layers, thereby lowering memory footprint and communication overhead.

It is worth noting the natural synergy between attention residuals and Gated MLA: the gating mechanism compresses the information redundancy of each layer’s attention output, while the residual connections ensure that the compressed critical information can propagate over long distances. Together, they form a closed-loop optimization path of “compress-transmit-enhance.”

2. Training Stability Solutions for Extremely Sparse MoE Architectures

When an MoE layer contains 896 routed experts, each token activates only a minuscule proportion (about 0.05%) of the expert parameters. While this extreme sparsity brings significant computational efficiency gains, it introduces two tricky engineering problems: first, severe load imbalance among experts, where some experts become “hot experts” carrying too many tokens, causing computational and communication hotspots to overlap; second, under large-batch distributed training, the communication volume of Expert Parallelism can account for over 50% of the overall training latency.

2.1 Stable Latent MoE Architecture

In traditional MoE, expert networks directly process the hidden states of tokens. Latent MoE, however, introduces a layer of low-dimensional latent variables for intermediate mapping, projecting the token’s representation into a shared low-dimensional subspace before distributing it to different experts. The advantage of this approach is that the latent variable space acts as a compressed representation, making the input distribution received by each expert more stable and reducing training oscillations caused by drastic input distribution fluctuations. Simultaneously, Stable Latent MoE incorporates the SiTU-GLU activation function, which is suited for sparse gradients, further enhancing gradient flow.

From an engineering perspective, the latent variable design brings an additional benefit—it inherently supports more flexible load balancing. Since all tokens pass through the same latent variable mapping before entering the experts, routing decisions can be made quickly based on low-dimensional features, without calculating similarity for each high-dimensional hidden state individually. This significantly reduces the computational latency of the routing module itself.

2.2 Quantile Balancing Mechanism

In ultra-large-scale MoE, Auxiliary Loss is a common load balancing method, but it has two drawbacks: first, it interferes with the gradient direction of the main loss function; second, its coefficient requires fine-tuning, and improper scaling can trigger training instability. Stable Latent MoE introduces the Quantile Balancing mechanism into the dynamic routing process, completely eliminating auxiliary loss. Instead, it calculates quantile thresholds and adjusts routing probabilities in real-time by tracking the distribution of token reception counts for each expert over a continuous training window.

Specifically, Quantile Balancing treats load balancing as an online optimization problem. The system maintains real-time load statistics for each expert and uses fixed quantiles (e.g., the top 25% and bottom 25% boundaries) as dynamic adjustment signals. When an expert’s load exceeds the high quantile, its routing probability is penalized; conversely, experts with loads below the low quantile receive a boost. This mechanism can quickly respond to data stream distribution shifts without introducing extra loss terms during forward computation, greatly simplifying the training process.

The universality of this strategy lies in the fact that it is not limited to MoE layers. Any system involving multiple parallel computing units with dynamically changing loads (such as multi-node, multi-GPU pipeline parallelism) can adopt this statistical distribution-based quantile adjustment philosophy.

3. Multi-Stage Pre-training and Agentic RL Post-training

The model adopts a multi-stage pre-training strategy scaling context from 8K to 1M: advancing step-by-step through 8K → 64K → 256K → 1M. Each stage continues training with a portion of the weights from the previous stage and proportionally increases the ratio of long-sequence data. Compared to directly jumping from short to long sequence training, this approach significantly reduces the risk of loss divergence. The core bottleneck of ultra-long context training is that while attention computation grows linearly with length (after adopting linear attention), the extrapolation capability of position information like RoPE (Rotary Positional Embedding) at 1M length remains limited. Multi-stage training allows the positional encoding to gradually adapt to longer relative distance ranges, effectively mitigating the degradation of extrapolation performance.

The post-training stage introduces Agentic RL (Agentic Reinforcement Learning). Agents need to execute complex long-horizon tasks such as multi-turn tool calls, code generation and execution, and information retrieval. Reinforcement learning utilizes a Rule-based Reward Model and a Process Reward Model to provide feedback signals on intermediate steps rather than just the final result. This fine-grained reward mechanism effectively guides the model to maintain coherence during inference journeys far longer than the training sequence length.

It is crucial to emphasize that the success of Agentic RL highly depends on the synergy among “long context + tool invocation + code execution.” Pure text generation models often suffer from “forgetting early instructions” in long-horizon tasks. However, a 1-million-token context window combined with RL feedback allows the model to repeatedly review and correct its own actions within the same conversation, laying the technical foundation for building genuinely usable intelligent assistant products.

4. Key Insights from Engineering Practice

From the technical analysis above, we can extract the following engineering experiences that offer direct guidance to developers:

First, operator limits determine algorithmic limits. The KDA case demonstrates that under FP16/BF16 precision, the theoretical advantages of linear attention can be largely offset by numerical range requirements. When designing new operators, one should start from the instruction cycle perspective of the target hardware (e.g., the Tensor Core pipeline) and reverse-engineer the algorithm’s parallelizable granularity and numerical constraints.

Second, communication optimization in sparse patterns is a top engineering priority. The most easily overlooked aspect in MoE parallelism is the All-to-All communication pattern. At the scale of 896 experts, each token needs to be distributed across machines, and communication volume increases linearly with the number of experts. A sound engineering strategy is to perform two-stage routing based on “expert distribution and data locality”—local aggregation first, followed by global exchange—transforming communication hotspots into compute-intensive operations.

Third, load balancing requires online adaptive mechanisms. Fixed-threshold load balancing strategies often fall short when facing data drift. The Quantile Balancing approach reminds us that treating load balancing as a statistical process control problem—adjusting allocation weights in real-time through distribution drift detection—is the reliable path to supporting stable training over tens of thousands of steps.

Fourth, the context window is a capability boundary, not a decorative metric. The training path of gradually expanding from 8K to 1M is essentially a progressive shaping of the model’s “memory.” In distributed environments, long-context training must also be paired with a 3D hybrid parallel strategy combining Sequence Parallelism, Tensor Parallelism, and Expert Parallelism, requiring dynamic tuning across dimensions based on actual memory budgets and communication bandwidth.

Conclusion

Ultra-large-scale sparse MoE models are not merely about stacking parameters; they represent a systems engineering challenge deeply intersecting numerical computation, network communication, distributed scheduling, and reinforcement learning. The four core technologies—KDA parameterization, stable training with “latent experts,” quantile-based load balancing, and multi-stage context expansion—provide a high-quality engineering paradigm for subsequent large-scale model development. For teams exploring trillion-parameter models or ultra-long context applications, these methods offer directly reusable guidance. The core lies not in copying the architecture, but in understanding the hardware constraints and training stability demands behind each design decision, thereby establishing a proprietary engineering methodology.