Optimizing Trillion-Scale MoE Models: Long Context and Sparse Training Best Practices

As multimodal large models scale into the trillion-parameter range, sustaining inference efficiency while supporting ultra-long contexts of up to a million tokens has become a core challenge for both architecture design and training engineering. The current industry mainstream utilizes a combination of Sparse Mixture-of-Experts (MoE), hybrid linear attention, and multi-stage context extension. This article delves into three highly reusable engineering problems within this technical route: the numerical stability of linear attention in matrix multiplication acceleration scenarios, load balancing for ultra-large-scale expert routing, and the efficiency of cross-layer information transfer in deep networks.

1. Architectural Trade-offs in Trillion-Scale MoE

Before discussing specific optimization techniques, it is essential to understand the fundamental form of this generation of architectures. A typical trillion-parameter MoE model employs a sparse activation strategy: “massive total parameters, smaller activated parameters.” The total parameter count can reach trillions, while the parameters actually activated during a single forward pass are merely in the hundreds of billions. This design decouples model capacity from inference costs but introduces three structural challenges:

1. Computational complexity of attention modules: The quadratic complexity of standard full attention over sequence lengths cannot support million-level contexts. Therefore, linear attention or state space modules must be introduced as supplements.
2. Training stability of expert routing: When the number of routing experts reaches hundreds or even nearly a thousand, traditional auxiliary loss balancing strategies often suffer from slow convergence and load oscillation.
3. Representation bottlenecks in deep Transformers: As the number of layers increases, information dilution along the residual path intensifies, degrading both gradient flow and representation quality in deep networks.

The following sections detail the technical solutions corresponding to these three challenges.

2. Lower-Bound Decay Parameterization: Maximizing Matrix Multiplication Units for Linear Attention

2.1 Problem Background

Linear attention and Delta Net-like structures require an cumulative decay factor Γ during chunked computations to describe the exponential forgetting of historical information. Within a chunk, to transform the recursive formulation into a matrix multiplication format, the decay factor’s reciprocal must be used for scaling. The issue arises when the decay factor itself is very small: its reciprocal quickly overflows under limited precision formats like BF16. Conversely, in certain intervals, it may underflow to zero, causing numerical instability.

This numerical issue directly forces engineering compromises: the diagonal tile must rely on an explicit position-by-position computation path, preventing it from utilizing Tensor Core matrix multiplication acceleration. Consequently, the theoretical linear complexity advantage is compromised at the implementation level, and operator efficiency falls significantly short of expectations.

2.2 Mechanism Design

The solution is to change the decay parameter in the logarithmic domain from an unbounded form to a bounded one. Traditional approaches use a negative Softplus function to generate the decay rate, with a range from negative infinity to zero, causing the reciprocal scaling factor to be unbounded at the lower end. The improved approach adopts a scaled Sigmoid form with a lower bound:

g = g_min * Sigmoid(e)
where g_min is a fixed lower bound (e.g., -5)

The key to this modification is that within any fixed-length chunk (e.g., 16 tokens), the reciprocal of the decay factor is locked within a safe numerical range and will not exponentially explode for later sequence positions. As a result, diagonal and non-diagonal tiles can adopt a unified causal chunked computation path, executing entirely via matrix multiplication.

2.3 Engineering Benefits

From an operator implementation perspective, the engineering effort for this change is minimal—essentially just replacing the parameterization of an activation function. However, it eliminates special branches in chunked computations, resulting in:

• All causal chunks (diagonal and non-diagonal) sharing the same compute kernel, reducing code paths and register pressure.
• Diagonal tiles fully utilizing matrix multiplication units, improving overall throughput.
• Enhanced numerical stability allowing for more aggressive use of low-precision formats during training and inference.

For teams developing their own linear attention operators, this is an optimization point with an exceptionally high return on investment.

3. Auxiliary-Loss-Free Expert Load Balancing

3.1 Limitations of Traditional Approaches

The most classic problem in MoE training is “routing collapse”: a few experts are excessively selected while most experts remain untrained, meaning the effective parameter count is far below the nominal count. The common industry countermeasure is introducing an Auxiliary Loss, adding a regularization term to the total loss that encourages load uniformity.

However, auxiliary loss presents an inherent contradiction: if the regularization weight is too high, it interferes with the main task’s optimization objective; if it is too low, it fails to balance the load. When the number of experts increases to hundreds, this contradiction is further amplified, often requiring a lengthy hyperparameter tuning process during early training phases.

3.2 Quantile Balancing Approach

An alternative approach does not rely on gradient-level regularization but instead makes dynamic adjustments directly at the routing decision level. The core idea is to maintain a dynamic bias term for each expert, adjusting the bias based on the recent actual number of tokens received—decreasing the bias for experts receiving too many, and increasing it for those receiving too few.

From a statistical perspective, this is equivalent to solving a quantile matching problem during each routing step: aligning the selection probability quantiles for each expert. Because the adjustment occurs at the routing scoring stage rather than within the loss function, it does not conflict with the main task’s gradients and requires no additional hyperparameter search.

3.3 Practical Considerations

In actual implementation, several points require attention for this method:

• Bias update frequency: Updating too frequently leads to routing oscillation, while updating too sparsely causes lag. Updating per training step or per micro-batch is generally reasonable.
• Bias amplitude constraints: Upper and lower limits must be set to prevent individual expert biases from being pushed to extreme values.
• Coordination with capacity factors: Even with good load balancing, capacity factors are still needed to handle instantaneous imbalances within a single batch.

Compared to auxiliary loss solutions, this mechanism converges faster in ultra-large-scale expert scenarios and involves fewer hyperparameters, making it suitable as a default configuration.

4. Cross-Layer Attention Residual: Mitigating Deep Representation Degradation

4.1 Representation Dilution in Deep Networks

In standard Transformers, the output of each layer is accumulated into the main stream via residual connections. As the number of layers increases, the proportion of early layer information in the main stream is continuously diluted. Deep networks essentially have to “re-learn” information that has already been lost. While this is tolerable in models with dozens of layers, further increasing the depth directly degrades representation quality in long-range reasoning tasks.

4.2 Blockwise Attention Residual Concept

The core improvement is to allow the attention module to directly access information from earlier layers, rather than relying solely on main stream accumulation. Specifically, this can be understood as dividing layers into blocks. Within a block, attention computation uses not only the current main stream but also introduces the representation from the beginning of the block as an additional reference.

This design offers two primary benefits:

• Shortened information path: Deep attention does not need to go through dozens of residual accumulations to “see” shallow information, reducing information loss.
• Controllable memory and communication overhead: Because it is organized by blocks, there is no need to store the complete activations of all intermediate layers. Retaining only the representations at block boundaries results in a limited memory footprint increase.

4.3 Synergy with Other Technologies

Cross-layer residuals and the aforementioned linear attention optimizations are not isolated. In linear attention scenarios, the state matrix itself handles cross-position information transfer, while blockwise residuals supplement information flow across the layer dimension. Combining the two improves both long-range dependency modeling and deep representation quality without significantly increasing computational overhead.

5. Multi-Stage Context Extension and Post-Training

Beyond architectural optimizations, extending context length typically employs a staged progression strategy: starting pre-training with a shorter context window (e.g., 8K) and gradually expanding to 64K, 256K, and up to 1M. The reasons for this progressive strategy include:

• The short-context phase allows for faster convergence of foundational language capabilities and multimodal alignment.
• The long-context phase primarily adjusts position encodings and attention distributions, usually requiring learning rate rescheduling.
• Directly training ultra-long contexts from scratch introduces massive computational waste and instability.

After pre-training, the post-training phase tailored for agentic long-range reasoning is equally critical. The goal here is no longer simple language modeling, but enabling the model to maintain stable reasoning chains across scenarios like multi-turn tool calling, long document retrieval, and cross-step planning. Common techniques include reinforcement learning based on verifiable rewards and credit assignment optimization for long trajectories.

6. Practical Recommendations for Engineering Teams

Synthesizing the analysis above, teams currently building or optimizing large-scale MoE models should prioritize the following initiatives:

Optimization Direction	Implementation Difficulty	Expected Benefit	Suggested Priority
Bounded decay parameterization	Low	Improved operator efficiency and numerical stability	High
Dynamic bias load balancing	Medium	Improved training stability and parameter utilization	High
Blockwise attention residual	Medium-High	Improved deep representation and long-range task performance	Medium
Multi-stage context extension	High	Long-context capability	As needed by business requirements

It is crucial to emphasize that these techniques are not mutually exclusive but complementary. Bounded decay parameterization resolves efficiency at the single-layer operator level; load balancing addresses utilization at the parameter level; cross-layer residuals tackle information flow at the depth level; and multi-stage training resolves capability construction along the temporal dimension. Together, they form the engineering foundation of ultra-long context MoE models.

For resource-constrained teams, it is advisable to start with the first item, which requires minimal changes and offers direct benefits. After validating the numerical stability improvements, progressively advance toward load balancing and cross-layer structural modifications.