Breaking Through Trillion-Parameter Multimodal MoE: Key Innovations from Ultra-Long Context to Extreme Sparsity
This article explores the critical engineering and algorithmic breakthroughs in next-generation trillion-parameter Mixture-of-Experts (MoE) multimodal large models. It addresses three primary bottlenecks: numerical instability in linear attention, load balancing in ultra-large MoE routing without auxiliary loss, and communication overhead in deep Transformer architectures. Key innovations discussed include bounded decay parameterization for Tensor Core optimization, quantile balancing for efficient routing, and block-level attention residuals. The article also covers multi-stage training strategies for 1M context lengths, Agentic RL post-training for long-horizon reasoning, and reusable engineering patterns for the broader AI development community.
1. Technical Background and Core Challenges
As large model parameter scales surpass the trillion-level mark, the sparse Mixture-of-Experts (MoE) architecture has become the mainstream solution for balancing computational costs with model capacity. However, when simultaneously pushing model scale and context length to their absolute limits, both engineering and algorithmic layers encounter three interconnected bottlenecks.
Bottleneck 1: Numerical Stability in Linear Attention. While traditional Softmax attention offers high precision, its computational complexity grows quadratically with sequence length, making it incapable of supporting million-token processing demands. Consequently, the industry has largely shifted towards Linear Attention or its linearized variants to reduce complexity to linear. However, linear attention introduces a decay factor during chunked parallel computation. The cumulative effect of this factor in low-precision formats like BF16 easily triggers numerical underflow or overflow. This causes certain computation paths to fail to leverage Tensor Core matrix multiplication units for acceleration, degrading into slow, position-by-position loops.
Bottleneck 2: Load Balancing Contradictions in Ultra-Large MoE. In MoE layers containing nearly a thousand routing experts, routing strategies that rely on Auxiliary Loss to enforce balance often damage the specialization capabilities of the experts. Furthermore, this introduces additional synchronization overhead during large-scale distributed training. Maintaining near-equal computational load across every expert within dynamic batches, without relying on auxiliary loss, is critical to ensuring end-to-end training throughput.
Bottleneck 3: Representation Degradation and Communication Overhead in Deep Transformers. When a model reaches hundreds of layers, the deep network faces the risk of Representation Collapse, causing deep modules to degenerate into “identity mappings.” Simultaneously, the cross-device communication cost brought by Expert Parallelism expands rapidly as the number of experts increases. Traditional cross-layer attention connection schemes further exacerbate memory and bandwidth pressures.
2. Core Innovative Technologies
To address these bottlenecks, the model’s technical report proposes three highly engineering-driven innovations, targeting attention computation, MoE routing, and hierarchical connections, respectively.
2.1 Bounded Decay Parameterization: Unlocking Tensor Core Performance for Linear Attention
In standard linear attention (such as Delta Net), the cumulative decay factor Γ within a chunk determines the decay rate of historical information. Its logarithmic form typically uses an unbounded negative Softplus function (i.e., Log(σ(x)), with a value range of (-∞, 0)). The underlying issue is that when input feature values are too large, the reciprocal of Γ (i.e., 1/Γ) can inflate to hundreds or even thousands under BF16. This exceeds the representable range of 16-bit floating-point numbers, leading to numerical instability.
The technical solution applies a hard lower bound to the decay parameter, rewriting it as a scaled Sigmoid function. The expression can be simplified as:
g_t = g_min · σ(w · x_t + b)
Here, g_min is set to a negative constant (e.g., -5). This means the value range of the decay factor at any given time is compressed to [g_min, 0), and its reciprocal is strictly controlled within a safe floating-point representation range.
The engineering benefits are profound: Because the numerical scaling within any 16-token tile is within a safe range, the computation kernel no longer requires separate CUDA Kernels for diagonal tiles and non-diagonal tiles, nor does it need explicit position indexing paths. All tiles can uniformly invoke Dense GEMM (General Matrix Multiply), achieving near-theoretical peak FLOPs utilization on GPUs like the A100/H100. This modification is entirely implemented at the operator level and is zero-intrusion to upper-layer training scripts and model structures.
2.2 Quantile Balancing: Auxiliary-Loss-Free Routing for Ultra-Large MoE
In giant MoE layers containing 896 routing experts, traditional Auxiliary Loss load balancing methods (like the load balancing loss introduced by Switch Transformer) require calculating the KL divergence between the expert distribution and a uniform distribution for each token. Its gradients often conflict with the main task loss. Moreover, under dynamically shaped inputs, auxiliary loss struggles to respond timely to unbalanced fluctuations in token flow.
The model’s innovation does not compute auxiliary loss during training. Instead, it directly performs real-time statistics on routing weights (i.e., the logits generated for each token input). Specifically, each expert maintains a quantile estimator for its historical logits. When an expert’s cumulative load exceeds a preset threshold (e.g., 105% of the mean), the routing algorithm dynamically introduces a random rejection sampling step, causing the expert to temporarily “skip” the tokens it is least confident about.
Its core value lies in: requiring no learnable auxiliary parameters and being entirely driven by deterministic statistics. In distributed scenarios, it only requires lightweight AllReduce communication within the expert parallel group to synchronize quantile information. The communication volume is significantly smaller than the complete forward-backward computation of traditional auxiliary loss. This ensures that even at a scale of 896 experts, the training process maintains near-zero idle waiting.
2.3 Block Attention Residual (Block AttnRes): The Bandwidth Killer via Cross-Layer Sharing
Representation degradation in deep Transformers is typically mitigated by Residual Connections, but standard residuals only occur between adjacent layers. To promote deeper information flow, the industry has explored dense connections or cross-layer connections. However, these schemes require passing shallow activation tensors to deep layers, introducing extremely high memory usage in MoE models.
The report proposes a “intra-block residual sharing” strategy: treating several consecutive Transformer layers as a single computation Block. Within the block, the intermediate computation results of the attention mechanism (i.e., the hidden state before the attention output projection) are shared as residuals with all subsequent layers in the block. Due to the pre-normalization architecture, this residual connection does not alter the computational manifold of each layer but effectively shortens the information propagation path of the deep network.
The key advantage is memory optimization: This residual connection is compiled as an in-place operation at the computation graph level and does not require passing intermediate tensors between blocks, resulting in almost zero marginal memory overhead. Concurrently, because the intra-block attention output sharing ratio is high, partial KV cache update operations can be merged during the inference phase, reducing memory bandwidth pressure. This is crucial for online serving of million-token-level contexts.
3. Multi-Stage Training Strategy and Engineering Implementation
The model adopts a multi-stage pre-training scheme that gradually expands context from 8K to 64K, 256K, and ultimately 1M. This avoids the slow convergence issues associated with single-stage training under extremely long contexts.
The training phases are as follows:
- Phase 1 (8K Context): Focuses on dense semantic learning, utilizing a very large batch size (e.g., 4096). The optimizer used is AdamW, with a peak learning rate set to 1.5e-4, supplemented by Warmup and Cosine decay.
- Phase 2 (64K Context): Freezes the parameters of several bottom layers, only unfreezing the relative positional encoding module of the attention layers and the top-level classifier, to accelerate the acquisition of long-range dependencies.
- Phase 3 (256K to 1M Context): Introduces a Spiral Packing strategy, interleaving and packing multiple short documents into a single long sequence to improve training efficiency and avoid perplexity sudden shifts caused by positional encoding extrapolation.
In terms of engineering implementation, the model utilizes Sequence Parallelism to split long sequences across different GPUs, combining it with Expert Parallelism and Tensor Parallelism to form a 3D parallel strategy. For 1M context, merely computing the causal mask and storing the KV Cache requires massive memory. The solution adopts a combination of Chunked Prefill and Paged KV Cache.
4. Agentic RL Post-Training and Long-Horizon Reasoning
The model innovatively expands the post-training phase from traditional SFT and RLHF to Agentic RL. Specifically, it defines a long-horizon task within the environment that includes tool usage, code execution, and multi-step search. The reward model is no longer based on a single scalar of the final answer, but rather on the legality of each step’s behavior, the effectiveness of the intermediate reasoning path, and the attainment of the final goal.
Technical Key Points: In RL training, the model needs to process complete trajectories exceeding 1 million tokens. Gradient calculation requires backpropagation across hundreds of time steps, placing extremely high demands on optimizer memory usage and numerical stability. The solution is to perform gradient accumulation at the trajectory level, paired with a “relaxed truncation” strategy—for overly long trajectories, the gradients of early time steps are compressed into low-rank representations to save memory consumption.
In practice, following Agentic RL post-training, the model demonstrates stronger long-horizon planning capabilities. For instance, in multi-hop knowledge QA and repository-level bug fixing tasks, its success rate is significantly higher than baseline models that only underwent SFT.
5. Reusable Key Technology Checklist
Although the technologies mentioned in this article are based on large MoE models, the underlying operators and algorithmic concepts are highly universal and suitable for migration in the following scenarios:
- Linear Attention Operator Optimization: For any model utilizing linear attention or linearized Transformers (such as speech recognition or long document classification), introducing bounded decay factor parameterization can directly improve operator efficiency on Tensor Cores without altering the model structure.
- Large-Scale Recommendation Systems / Risk Control Models: The Quantile Balancing (QB) mechanism can be used independently of Transformers. In scenarios with thousands of feature bins or streaming model shards, QB can serve as a lightweight dynamic load balancer, ensuring the utilization of downstream computational resources.
- Distributed Training Framework Optimization: The design pattern of Block Attention Residual can be used to guide hierarchical cascading training strategies, reducing the frequency of cross-node activation synchronization.
6. Conclusion and Future Outlook
As seen from this technical report, the core competitiveness of trillion-parameter MoE models no longer rests solely on parameter count, but rather on extreme resource utilization efficiency and numerical stability design. For the developer community, the parameterization approach of bounded decay and the auxiliary-loss-free balancing strategy of QB provide targeted inspiration—they prove that algorithmic innovations tailored to underlying hardware characteristics can yield performance gains far surpassing the mere stacking of computational resources. In the future, as context lengths approach millions of tokens, the joint optimization of attention operators and communication topologies will become the main battlefield for the engineering deployment of large models.
