Optimizing Long Context Training and Inference for Trillion-Scale Sparse Models

1. Background: Why Trillion-Parameter Models Struggle with Long Context and Sparse Experts

In recent years, the capability improvements of large language models have relied not only on scaling up parameters but also on two critical dimensions: longer context windows and stronger multimodal input capabilities. When a model’s total parameters reach the trillion scale, adopting a fully activated dense structure makes both training and inference costs unacceptable. Consequently, sparse expert models have become a common choice: the total parameter count is massive, but only a fraction of these parameters is activated during each forward pass.

However, a sparse expert model is not simply a standard Transformer with its feed-forward network replaced by expert networks. It introduces three prominent challenges:

1. Long context attention efficiency: When the context reaches 100K or 1M tokens, the memory footprint and computational load of traditional attention mechanisms increase drastically. It becomes necessary to introduce linear attention, delta rules, or incremental attention mechanisms. However, these mechanisms often suffer from numerical instability during low-precision training and tensor core parallelization.
2. Massive expert load balancing: As the number of experts grows from dozens to hundreds or even nearly a thousand, if some experts are frequently selected while others remain idle, model capacity is wasted, and training may become unstable. Traditional auxiliary losses can improve balance but may interfere with the optimization of the main task.
3. Representation bottlenecks in deep Transformers: As the number of layers increases, residual connections, while mitigating vanishing gradients, can restrict information flow, cause deep layer representations to converge, and increase memory and communication overhead.

To address these issues, a reusable set of engineering solutions has emerged: stabilizing linear attention with bounded decay parameterization, replacing auxiliary loss with a quantile balancing strategy, and enhancing deep information sharing via block-level attention residuals. Let’s explore each of these in detail.

2. Attention Mechanisms: From Full Attention to Incremental Attention

2.1 The Core Conflict of Long Context

In a standard Transformer, the computational complexity of self-attention scales quadratically with the sequence length. Taking a 1M token sequence as an example, directly computing the full attention matrix would result in an unacceptable memory footprint. While practical training often employs chunking strategies like FlashAttention, long sequences still impose immense computational pressure.

To support 1M-level contexts, a common approach is to introduce linear attention, Delta rules, or incremental attention. These methods decompose attention into a “state update + output read” format, allowing computation to proceed in blocks and avoiding the explicit storage of the complete attention matrix.

2.2 Numerical Issues in Incremental Attention

Although linear attention reduces complexity, it encounters numerical underflow during low-precision training. The reason is that state updates typically contain a cumulative decay factor. As the sequence progresses, the exponential term of this factor can become extremely small or large. Under BF16 precision, if the reciprocal of the decay factor is computed, it easily overflows to infinity or underflows to zero.

For instance, a common implementation might map the decay parameter to an unbounded value using a Softplus function. During computation within a long block, the numerical range of certain diagonal blocks may exceed the representable range of BF16. As a result, it becomes impossible to directly leverage Tensor Cores for matrix multiplication acceleration, forcing a fallback to a position-by-position computation path and significantly degrading computational efficiency.

2.3 Bounded Decay Parameterization

The solution is to set a lower bound for the decay parameter, ensuring it remains within a safe range across any block. One specific approach is to use a scaled Sigmoid function to replace the unbounded mapping. Setting a fixed lower bound of -5, we define:

g = g_min * sigmoid(raw_param)

Here, g_min is -5. This ensures that regardless of how the raw parameter varies, the decay value is strictly confined to the range [-5, 0]. Consequently, the reciprocal scaling factor within any 16-token block is also restricted to a safe interval.

The benefit of this approach is that the engineering changes are minimal: simply adjusting the parameterization format allows the linear attention operator to run stably on Tensor Cores, reducing the need for special handling of diagonal blocks. For teams utilizing low-level libraries like CUDA and CUTLASS, this is a highly reusable improvement.

3. Sparse Expert Routing: From Auxiliary Loss to Quantile Balancing

3.1 Why Auxiliary Loss Falls Short

In Mixture-of-Experts (MoE) models, the routing network determines which experts each token activates. When the number of experts is large, the routing network tends to assign most tokens to a few experts, creating a “winner-takes-all” scenario. To balance expert load, the traditional approach adds an auxiliary loss to the main loss function, forcing the routing distribution to be more uniform.

However, auxiliary loss presents several issues:

• The auxiliary loss coefficient requires manual tuning; if it is too large, it can harm the model’s final performance, and if too small, the load balancing effect is insufficient.
• Auxiliary loss is a global statistic, making it difficult to reflect instantaneous imbalances within local batches.
• Under extreme scale and sparsity conditions, the gradient of the auxiliary loss may conflict with the gradients of the main task.

3.2 Quantile Balancing Mechanism

The concept behind quantile balancing is to stop guiding the router through an additional loss and instead directly sort and group the routing scores. At each training step, based on the preference scores of tokens for different experts, tokens are divided among experts by quantiles, ensuring that each expert receives a roughly equal number of tokens.

The specific process can be understood as follows:

1. Compute the routing scores for all experts for the tokens in the current batch.
2. Sort the scores and determine the acceptance threshold for each expert.
3. Assign tokens to experts based on the thresholds, ensuring each expert processes a nearly equal amount of tokens.

The advantage of this mechanism is that it eliminates the need for an auxiliary loss coefficient, reducing the hyperparameter tuning burden. Furthermore, because load balancing is handled directly within each training step, the response is faster. For ultra-large scale models with hundreds of routing experts, this strategy can significantly improve training stability.

4. Cross-Layer Information Sharing: Block-Level Attention Residuals

4.1 Representation Issues in Deep Transformers

As model depth increases, the changes in each layer are stacked onto the residual stream. While residual connections facilitate gradient propagation, they can also lead to two problems: first, the magnitude of changes in deep representations becomes increasingly small, leading to diminishing returns from model depth; second, each layer must maintain a complete attention state, causing memory and communication overhead to grow linearly with the number of layers.

4.2 Design of Block-Level Attention Residuals

The core idea of block-level attention residuals is to treat several consecutive Transformer layers as a single block and introduce an additional attention residual connection within the block. Instead of going through a full layer-by-layer transformation, this residual fuses the attention state from the block’s entrance with the outputs of the layers within the block.

This approach provides at least three benefits:

• Reduced memory overhead: Attention states within the block can be shared, eliminating the need to store a complete copy for every single layer.
• Reduced communication overhead: In distributed training, cross-layer sharing can reduce certain tensor communications.
• Enhanced information flow: Deep networks no longer rely solely on layer-by-layer residuals; they can access a more direct information pathway via block-level residuals, mitigating deep representation degradation.

5. Multi-Stage Pretraining and Ultra-Long Context Post-Training

5.1 Phased Context Length Extension

Expanding the context from 8K to 1M in a single step introduces two problems: insufficient data and computational resources, and the model’s inability to adapt smoothly. A more common practice is phased expansion:

8K -> 64K -> 256K -> 1M

At each stage, the model is first allowed to train stably at the current length before entering the next. This approach controls training costs while enabling the model to gradually learn long-range dependencies. The data mixture can also be adjusted at each stage, progressively increasing the proportion of long-context data to prevent the forgetting of short-context capabilities.

5.2 Agent-Oriented Long-Range Inference Post-Training

After pretraining, post-training oriented towards agent tasks is required. Unlike standard conversational tasks, agent tasks often involve multi-step planning, tool calls, and long-term memory. Therefore, post-training requires constructing trajectory data at the million-token level, allowing the model to maintain goal consistency within a long context and avoid repetition or forgetting critical information.

The difficulty of such post-training lies in reward signal design. Simple outcome rewards may be too sparse, necessitating the use of process rewards, stage rewards, or automated scoring methods to help the model learn segmented decision-making within long trajectories.

6. Comprehensive Evaluation and Practical Value

Judging from technical reports, this architecture performs well on open-source benchmarks, indicating that the aforementioned designs are not just theoretically viable but have been validated through actual training and evaluation. For R&D teams, the most noteworthy aspect is not a single metric, but a reusable set of methodologies:

• Bounded decay parameterization: Minimal changes, clear benefits, suitable for teams optimizing linear attention operators to test rapidly.
• Quantile balancing mechanism: Replaces auxiliary loss, reduces tuning costs, and is ideal for sparse models with a large number of experts.
• Block-level attention residuals: Suitable for memory, communication, and representation optimization in deep models, though corresponding adjustments to model parallelism strategies are required.

When deploying such models to cloud inference services, further inference-side optimizations must be considered. For example, different experts may be distributed across different compute nodes, requiring specialized optimization for routing latency and communication costs. Additionally, ultra-long contexts during inference can utilize hierarchical caching, state compression, or segmented summarization to reduce initial response time and continuous memory footprint.

7. Conclusion

The long-context capability of trillion-parameter sparse models is not achieved by single-point algorithms but relies on collaborative design across multiple levels. The numerical stabilization of attention mechanisms determines whether the model can be trained efficiently; the load balancing strategy of sparse expert routing determines whether massive parameters can be effectively utilized; cross-layer information sharing mechanisms impact the quality of deep representations; and multi-stage context extension with post-training strategies determines whether the model possesses stable long-range reasoning capabilities.

For practitioners involved in large-scale model training and inference optimization, these designs provide a clear technical roadmap. Rather than simply stacking parameters, they emphasize striking a balance between hardware efficiency, training stability, and model capability. As context lengths continue to grow and expert counts expand further, the importance of these foundational mechanisms will only increase.