Architecture Evolution and Training Engineering for Trillion-Scale Sparse Multimodal Models

    Scaling model architectures to trillion-parameter totals with billion-parameter activations and million-token contexts introduces severe bottlenecks in traditional dense Transformers. This article explores a comprehensive architecture combining a sparse Mixture-of-Experts (MoE) backbone with native multimodal processing and ultra-long context capabilities. It details three key engineering solutions: lower-bounded decay parameterization for stable linear attention, quantile balancing for auxiliary-loss-free expert routing, and block attention residuals to mitigate deep layer representation saturation. Additionally, it outlines a multi-stage pre-training pipeline and agent reinforcement learning post-training paradigm for million-token scenarios.

    _

    Architecture Evolution and Training Engineering for Trillion-Scale Sparse Multimodal Models

    1. Background: Crossing the Trillion-Parameter Threshold

    When a model’s total parameters cross the trillion-scale mark, its activated parameters remain at the hundred-billion scale, and its context window expands to one million tokens, the scaling path of traditional dense Transformers simultaneously encounters hard constraints in three areas:

    • Attention Computation: The computational and memory overhead of standard full attention grows quadratically under a million-token context. While linear or hybrid attention mechanisms are necessary, the numerical instability of linear attention under finite precision has long constrained its implementation efficiency on tensor cores.
    • MoE Training Stability: As the number of experts grows from dozens to nearly a thousand, routing load imbalance leads to some experts “starving” while others overload. This drastically amplifies gradient variance and makes training highly prone to divergence.
    • Deep Representation Degradation: As network depth increases, the differences in hidden states between adjacent layers tend to saturate. Deep stacking brings diminishing marginal returns, while cross-layer communication and memory overhead rise.

    This article uses an architectural approach designed for trillion-scale sparse multimodal models as a thread to systematically outline the causes and engineering solutions for these three types of problems. It also presents the overall pipeline design for long-context pre-training and agent post-training. The technical points described are general methodologies and can serve as a reference for designing similar systems.

    2. Overall Architecture: Sparse Backbone + Native Multimodal + Ultra-Long Context

    The overall form of this architecture can be summarized by three pillars:

    1. Sparse MoE Backbone: With approximately 2.8 trillion total parameters and around 104 billion activated parameters per forward pass, the routing expert scale reaches 896, representing an extreme sparse configuration. Higher sparsity means larger model capacity per unit of compute, but it also imposes much stricter requirements on routing stability.
    2. Native Multimodal Input: Vision and text are jointly modeled from the very first stage of pre-training, rather than performing late-stage adaptation on a language model. The key to native multimodality is that visual tokens and text tokens share the exact same attention and expert routing paths, thereby avoiding representation fragmentation between modalities.
    3. Million-Token Context: The attention module adopts a fused structure of incremental linear attention and gated multi-head latent attention. Cross-layer information sharing is achieved through attention residuals, allowing the million-token context to maintain acceptable computational efficiency during inference.

    Within this framework, what truly determines whether the model “can be trained, runs fast, and scales long” are four engineering details: the numerical lower bound of linear attention, the balancing strategy for expert routing, the block-level organization of cross-layer residuals, and the coordination between activation functions and normalization.

    3. Key Technology 1: Lower-Bounded Decay Parameterization for Linear Attention

    3.1 The Problem: Diagonal Tile Dilemma Under Finite Precision

    Linear attention (including incremental forms like Delta Net) is typically computed in chunks. Intra-tile computation requires introducing the reciprocal of a cumulative decay factor Γ for normalization, but this reciprocal easily overflows or underflows under finite precision like BF16. Once the numerical value goes out of bounds, the diagonal tile (the tile containing the current position itself) cannot follow the standard dense matrix multiplication path. It must degenerate into explicit position-by-position computation—a special branch that not only drags down throughput but also makes it difficult for the compiler to perform unified tile scheduling.

    3.2 Mechanism: Changing Unbounded Decay to Bounded Decay

    The solution is to change the log-domain decay parameter from an unbounded negative softplus form to a scaled sigmoid form with a lower bound. By fixing the lower bound (e.g., -5), the decay coefficient is constrained within a safe interval:

    1. # Conceptual illustration: Bounded decay parameterization
    2. g_min = -5.0
    3. g = g_min * sigmoid(e + b) # Value is always in the (g_min, 0) interval
    4. Gamma = exp(cumsum(g)) # Cumulative decay, reciprocal scaling factor is bounded

    Because the reciprocal scaling factor of the cumulative decay is locked within a safe numerical range for any fixed-size tile (e.g., 16 tokens), the diagonal and non-diagonal tiles can adopt the exact same computation path:

    1. for tile in causal_tiles(Q, K, V):
    2. # Diagonal and non-diagonal tiles uniformly use dense GEMM
    3. O_tile = dense_gemm(Q_tile, K_tile, V_tile, scale=inv_gamma(tile))

    3.3 Benefits: Unified Path + Natural Extrapolation

    • Throughput Improvement: After eliminating the special branch for diagonal tiles, all causal tiles can directly invoke dense matrix multiplication on tensor cores. The complexity of operator implementation drops significantly, requiring minimal engineering changes.
    • Numerical Stability: The reciprocal scaling factor is bounded, avoiding overflow/underflow under BF16, keeping cumulative errors for long sequences controllable.
    • Positional Encoding Decoupling: This parameterization naturally supports No Positional Encoding (NoPE) forms, giving the million-token context direct extrapolation capabilities without relying on post-processing methods like positional interpolation.

    The reuse value here is that it is not an exclusive trick for a specific model, but rather a numerical engineering method that any system employing chunked linear attention and hoping to achieve high utilization on tensor cores can adopt.

    4. Key Technology 2: Auxiliary-Loss-Free Quantile Balancing

    4.1 The Problem: The Dilemma of Auxiliary Loss

    Traditional MoE load balancing relies on Auxiliary Loss to penalize skew in the routing distribution. However, in scenarios with nearly a thousand experts and extreme sparsity, auxiliary loss presents two issues: First, the loss weight is difficult to tune—too strong and it harms the representation quality of the main task; too weak and it fails to suppress load collapse. Second, the auxiliary loss and the main loss gradient directions may conflict, causing instability in the early stages of training.

    4.2 Mechanism: Dynamic Routing Thresholds via Quantiles

    The idea behind Quantile Balancing is to move “balancing” from the loss function level down to the routing decision level. It performs online statistics on the routing score distribution of each expert and takes the quantile based on the target load ratio as a dynamic threshold. Tokens exceeding the capacity are reallocated or discarded, achieving load balancing without auxiliary loss.

    1. # Conceptual illustration: Quantile threshold routing
    2. scores = router(x) # [num_tokens, num_experts]
    3. threshold = quantile(scores, 1 - capacity_ratio) # Online quantile statistics
    4. mask = scores >= threshold # Dynamic capacity constraint
    5. dispatch(x, mask)

    4.3 Benefits: Stable, Fast, Tuning-Free

    • Training Stability: The balancing constraint is embedded in the routing, no longer competing with the main loss for gradients. Expert utilization remains uniform under extreme sparsity.
    • Faster Convergence: Avoids the manual tuning cycle for auxiliary loss weights; the routing strategy adapts dynamically with training.
    • Engineering Friendly: Quantile statistics can be maintained incrementally, with controllable additional communication and computation overhead, suitable for large-scale distributed training.

    5. Key Technology 3: Block Attention Residuals

    5.1 The Problem: Representation Saturation in Deep Stacking

    In deep Transformers, the cosine similarity of hidden states between adjacent layers increases with depth, meaning the effective “new information” added by deep layers is decreasing. Simply increasing the number of layers yields diminishing returns, while the overhead of layer-by-layer activation storage and cross-device communication grows linearly.

    5.2 Mechanism: Cross-Layer Residual Reuse at the Block Level

    Block Attention Residuals organize several consecutive layers into a block. Fine-grained residuals are retained within the block, while attention-style aggregation between blocks directly introduces representations from early blocks into deep blocks:

    1. for block in blocks:
    2. h = block_layers(h) # Intra-block standard residual
    3. h = attn_res(h, block_cache) # Cross-block attention-style residual aggregation
    4. block_cache.append(h)

    5.3 Benefits: Mitigating Depth Bottlenecks, Reducing Memory and Communication

    • Representation Reuse: Deep layers can directly read early block representations, mitigating information saturation in the depth direction.
    • Controllable Overhead: Saving and transmitting at the block granularity significantly reduces memory footprint and cross-device communication volume compared to layer-by-layer full residuals.
    • Training Friendly: The block-level structure facilitates pipeline parallel splitting, reducing the impact of pipeline bubbles on the residual path.

    6. Coordination Between Activation Functions and Sparse Conditions

    In extremely sparse MoEs, the choice of activation function inside experts affects gradient scale and the numerical behavior of sparse activations. Adopting a gated SiTU-GLU-like structure can improve deep gradient propagation while maintaining the expressive power of GLU gating. Combined with the aforementioned quantile balancing, this allows the combination of “sparse routing + gated activation” to remain numerically controllable at the trillion-parameter scale. Engineering-wise, it is important to note that the implementation of the activation function should be aligned with the expert parallelism partitioning method to avoid extra synchronization caused by gating branches crossing devices.

    7. Multi-Stage Long-Context Pre-Training and Agent Post-Training

    7.1 Context Progression: 8K → 64K → 256K → 1M

    Long-context capability should not be achieved in a single step, but rather progressed in stages. The short-context stage uses massive general corpora to establish foundational language and multimodal alignment. Subsequently, the window is expanded level by level. Each stage is proportioned with long documents, repository-level code inputs, and multi-image multi-page visual samples of corresponding lengths, while position-related parameters and chunking strategies are adjusted synchronously. The core benefit of staged progression is that each level expands upon the already converged representations of the previous level, avoiding taking on the full computational and stability risks of a million-length context from the very beginning.

    7.2 Agent Reinforcement Learning Post-Training

    Agent reinforcement learning post-training at the million-token level aims to transform long-context capabilities into long-horizon task execution capabilities. Key engineering points include:

    • Trajectory Organization: Organize multi-step tool calls, environment feedback, and intermediate reasoning into long trajectories, and perform advantage estimation based on episode segmentation.
    • Reward Design: Use task completion as the main signal, supplemented by process constraints to avoid reward sparsity over long trajectories.
    • Sampling Efficiency: Sampling long trajectories is costly; it requires coordination with parallel environments and trajectory reuse strategies to control overall overhead.
    • Stability: Backpropagation over long sequences is sensitive to memory and numerical precision, and must be used in conjunction with the aforementioned linear attention numerical scheme and block-level residuals.

    8. Summary of Engineering Implementation

    Technical Point Problem Solved Reuse Value
    Lower-bounded decay parameterization Linear attention diagonal tile numerical overflow, operator branch redundancy High, minimal changes, direct throughput gains
    Quantile balancing Load collapse and auxiliary loss tuning under extreme sparsity with ~1,000 experts High, auxiliary-loss-free, adaptive
    Block attention residuals Deep representation saturation, layer-by-layer residual memory and communication overhead Medium-High, requires coordination with parallelism strategies
    Multi-stage context progression Training stability and cost for million-token contexts General methodology
    Long-horizon agent post-training Transforming long-context capability into task capability Depends on task and environment design

    9. Conclusion

    The difficulty of trillion-scale sparse multimodal systems often lies not in “drawing the model architecture diagram,” but in “running numerically stable, compute-efficient, and scaling to long contexts.” Lower-bounded decay parameterization resolves the conflict between finite precision and tensor core utilization; quantile balancing resolves the conflict between sparse routing and training stability; block attention residuals resolve the conflict between diminishing depth returns and communication overhead. Together, these three form a transferable engineering paradigm: pushing mathematical stability constraints forward into operator design, embedding balancing constraints into routing decisions, and sinking representation reuse down to the block level. For teams planning similar systems, these three paths offer low trial-and-error costs and high deterministic returns.