Training and Optimization Practices for Ultra-Large-Scale Sparse MoE Models

Introduction: How Flagship Sparse MoE Models Reshape Training System Design

In recent years, large model architecture evolution has followed two clear directions: sparse Mixture-of-Experts (MoE) for “large capacity, low activation” parameter organization, and linear or sparse attention for “long context, low cost” sequence modeling. When these two techniques are deeply integrated in flagship models, the design logic of the training infrastructure changes accordingly.

A flagship sparse MoE model exhibits three typical architectural characteristics:

1. Ultra-large-scale sparse experts: The model employs a sparse MoE architecture that constructs stable expert representations in a latent space, with 896 routing experts in total, where each token activates only 16 of them. The model has approximately 2.8T total parameters with an extremely low activated parameter ratio, effectively controlling per-token computation while maintaining large knowledge capacity.
2. Hybrid attention mechanism: Approximately three-quarters of the attention blocks use gated incremental linear attention (modeling sequence dependencies through state recurrence, with computational complexity growing nearly linearly with sequence length), while the remaining blocks incorporate Top-K sparse attention (e.g., Top-2048). This enables the model to support a maximum context window of 1 million tokens and native image understanding capabilities.
3. Reinforcement learning-driven capability leaps: Through large-scale joint reinforcement learning and system-level collaborative optimization, the model achieves enhanced capabilities in long-horizon autonomous programming, complex agent tasks, 3D creation, financial research, and deep knowledge work.

These architectural advantages do not come for free. Sparse routing, state recurrence, ultra-long sequences, and ultra-large vocabulary tables simultaneously impact three critical dimensions of the training system: computation, communication, and memory. This article uses the adaptation process of such models as a thread to systematically review the challenges and full-chain optimization solutions for ultra-large-scale MoE training.

Three Core Challenges from an Engineering Perspective

Computation: Interleaved Attention Mechanisms Complicate Operator Fusion and Scheduling

Linear attention is essentially a gated recurrent state update process: the model maintains a fixed-shape hidden state along the sequence and applies a recurrent transformation to the state for each token processed. This recurrent form is inherently serial, and a direct implementation would be severely limited by kernel launch overhead and memory access bottlenecks. The common industry practice is to introduce chunked computation: the sequence is divided into blocks, computations within each block run in parallel, and state transitions between blocks follow a recurrent pattern.

Meanwhile, sparse attention requires Top-K index selection first, followed by non-contiguous memory access based on the selected indices, resulting in a large number of fine-grained data rearrangement operations. When linear attention blocks and sparse attention blocks are interleaved within the same model, the computation graph exhibits “multi-stage, multi-shape, heterogeneous memory access patterns”:

• A single attention block contains multiple computational stages including state initialization, chunked recurrence, gated fusion, and output transformation.
• The index computation and data rearrangement on the sparse path generate numerous small operators.
• Different stages have significantly different memory access densities, making it difficult to cover them with a unified scheduling strategy.

If each stage is implemented as an independent kernel, intermediate results need to be repeatedly transferred between memory and compute units, significantly degrading bandwidth utilization and computational density. Therefore, the core challenge on the computation side is: while preserving numerical equivalence, fuse as many computational stages as possible into a few large kernels, and enable the scheduler to reliably feed them.

Communication: Load Imbalance from 896-Expert Sparse Routing

Expert parallelism distributes different experts across different devices. Before tokens enter the expert layer, they require two rounds of All-to-All communication: dispatch (distributing tokens according to routing results) and combine (aggregating outputs back). When the number of experts reaches 896 and each token activates only 16 of them, communication and load issues are significantly amplified:

• Imbalanced token distribution across experts: Skewed routing distributions are common in MoE. Devices hosting hot experts may accumulate several times the average number of tokens, creating “long-tail waits.”
• Small per-expert batch sizes: On average, each expert receives a limited number of tokens per training step, leading to fragmented computation granularity that fails to saturate compute units.
• Difficulty in overlapping communication and computation: Load imbalance causes All-to-All completion times to vary widely, and the overlap pipeline can easily be blocked by the slowest shard.

Under long-context and multimodal inputs, token distribution differences across experts are further amplified. Without systematic design on the communication side, communication can easily become the primary bottleneck for overall training throughput.

Memory: Multiple Pressures from Parameters to Activations

Memory is the most rigid constraint in ultra-large-scale model training. Taking the 2.8T parameter scale as an example, BF16 weights alone occupy approximately 5.6TB. On top of this, we must add:

• Gradients (same magnitude as weights).
• Optimizer states (commonly maintained in FP32 for first-order and second-order moments, potentially several times the weight size).
• Long-context activations (with million-token sequences, intermediate states remain substantial even with activation recomputation enabled).
• Logits and loss computation buffers for the ultra-large vocabulary output layer.

The gap between single-GPU memory and the above requirements spans more than two orders of magnitude. Sharding is the only way forward, but sharding strategies in turn affect communication volume and computational efficiency. All three dimensions must be designed jointly.

Full-Chain Training Optimization Solutions

To address the above challenges, a mature training capability system for ultra-large-scale MoE models typically covers five areas: model adaptation, hybrid parallelization, operator acceleration, communication optimization, and memory management, supported by supernode-level hardware and automated tooling.

Hybrid Parallelization: Orthogonal Combination of Fully Sharded Data Parallelism and Expert Parallelism

Facing 5.6TB of weights, no single parallelization strategy can balance sharding granularity and communication efficiency. The common engineering practice is an orthogonal combination:

• Fully Sharded Data Parallelism (FSDP): Shards parameters, gradients, and optimizer states across the machine group. Parameters for a given layer are temporarily gathered during forward/backward passes and released immediately after use. This trades controlled communication overhead for extreme memory compression, making it suitable for dense components such as attention and embedding layers.
• Expert Parallelism (EP): Shards the 896 experts across different devices along the expert dimension. Expert parameters reside on their assigned devices, and tokens traverse devices via All-to-All communication for cross-device routing. This is suitable for the sparse expert components.

When combined, expert parameters achieve distributed deployment and efficient aggregation: dense parts rely on FSDP’s “gather on demand,” while sparse parts rely on EP’s “each expert on its own device.” On a 128-GPU supernode cluster, this FSDP + EP64 (expert parallelism degree of 64) combination has been validated for LoRA fine-tuning of such models, demonstrating the correctness and usability of this paradigm at the hundred-GPU scale. An illustrative configuration approach is shown below (pseudo-code, not tied to a specific framework interface):

# Illustrative: FSDP + EP orthogonal combination configuration approach
parallel_config = {
    "data_parallel": 2,           # data parallel group size
    "expert_parallel": 64,        # expert parallelism degree, 896 experts sharded across 64 GPUs
    "sharding": "full",           # fully shard parameters / gradients / optimizer states
    "reshard_after_use": True,    # release parameters after use, re-gather on demand
    "expert_placement": "uniform" # uniformly distribute experts to mitigate static skew
}

Operator Acceleration: Large Fused Kernels for Forward and Backward Passes

For the “multi-stage, serial recurrence” nature of linear attention, a direct and effective method is to design large fused forward and backward kernels: integrate state initialization, intra-chunk parallel computation, inter-chunk state propagation, and gated fusion into a single kernel. Inter-chunk states ideally remain in on-chip high-speed memory to avoid repeated read/write of intermediate results to HBM. The backward pass is handled similarly, combining backward recurrence and gradient restoration of intermediate quantities into one large kernel. This provides three benefits:

1. The number of kernels is significantly reduced, lowering launch and scheduling overhead.
2. Intermediate tensors are no longer materialized to global memory, substantially reducing memory traffic.
3. Fused kernels exhibit stable computational density, enabling compilers and schedulers to perform deeper optimization.

Memory Tiering: Operator-Level Selective Caching

Within operators, the property of “reusable computation results” can be leveraged for memory tiering: for designated MatMul/GMM (grouped matrix multiplication) operators, computation results are written to a secondary cache. High-frequency accessed portions stay in HBM, while low-frequency but large-volume portions are demoted to DDR memory. When subsequent computations hit the cache, results are reused directly, avoiding repeated computation or repeated data movement.

The key is “selectivity”: a larger cache is not always better. Only results with high hit rates and high recomputation costs are worth caching; otherwise, additional management overhead is introduced. In the relevant test configuration, this mechanism delivered an average 15% training performance improvement, confirming that in HBM-constrained scenarios, the approach of “using tiers to expand capacity and caching to trade for bandwidth” yields practical benefits.

Large-Vocabulary Loss Computation: Dual Chunking Across Sequence and Vocabulary Dimensions

The output layer of large-vocabulary models is a common source of memory peaks: the logits matrix of sequence length × vocabulary size, plus the normalization statistics required for cross-entropy, can easily create transient memory spikes during the loss stage. The solution is to chunk in two orthogonal dimensions simultaneously:

• Chunk Loss (sequence-dimension chunking): Split the sequence into multiple chunks, compute logits and loss chunk by chunk, and release intermediate results immediately after computation. This flattens memory usage along the sequence direction.
• CCE Loss (vocabulary-dimension chunking): Cross-entropy only requires global logsumexp statistics. The vocabulary dimension can be chunked: compute local maxima and exponential sums per chunk, then merge to obtain global statistics, followed by a backward pass to compute per-token losses. This avoids keeping full-vocabulary logits in memory.

These two dimensions are orthogonal and can be stacked for dual optimization in both sequence and vocabulary. An illustrative pseudo-code example:

# Illustrative: chunked cross-entropy loss with sequence + vocabulary chunking (pseudo-code)
def chunked_ce_loss(hidden, weight, labels, seq_chunk, vocab_chunk):
    total_loss, total_tokens = 0, 0
    for h_blk, y_blk in split_sequence(hidden, labels, seq_chunk):
        m, s = -inf, 0                    # global max and exponential sum
        for w_blk in split_vocab(weight, vocab_chunk):
            logits_blk = h_blk @ w_blk.T  # only small logits blocks land in memory
            m_blk = logits_blk.max(-1)
            m, s = merge_max_sum(m, s, m_blk,
                     (logits_blk - m_blk).exp().sum(-1))
        lse = m + log(s)                  # global logsumexp
        total_loss += (lse - gather_logits(h_blk, weight, y_blk)).sum()
        total_tokens += y_blk.numel()
    return total_loss / total_tokens

Supernode Architecture and Automated Integration: Platformizing Optimization

For the above software optimizations to deliver their full value, a solid hardware foundation is essential. Supernode-style compute clusters provide three key characteristics:

• High bandwidth, low latency: Significantly reduces the cost of All-to-All communication and parameter gathering, allowing finer EP sharding granularity.
• Unified memory addressing: Makes HBM+DDR tiered caching transition from “usable” to “convenient,” enabling software to transparently migrate data across a unified address space.

Additionally, given the constant emergence of new models, adaptation efficiency itself is a competitive advantage. Through agent-driven automated development workflows, stages such as operator adaptation, parallel strategy generation, and precision alignment can be templated and automated, compressing the integration cycle for new models to the day level. This enables rapid iteration—which is particularly important for engineering teams that need to keep pace with rapid model architecture evolution.

Training Validation: LoRA Fine-Tuning Convergence

The ultimate proof of any solution lies in the training curves. On a 128-GPU supernode cluster, LoRA fine-tuning of such a model was validated using the FSDP + EP64 parallel combination. Results showed: under random initialization, the initial training loss was approximately 13.5; after more than 400 training steps, the loss converged to around 1.7.

These numbers carry two layers of engineering significance:

First, the numerical pathway is correct. Under random initialization, model output approximates a uniform distribution over the vocabulary, and the initial loss should be on the order of “the logarithm of the vocabulary size.” The initial value of 13.5 matches this theoretical expectation, confirming that forward computation, parallel sharding, and loss computation are numerically self-consistent—any sharding error or precision anomaly would typically surface directly in the initial loss.

Second, the convergence path is healthy. The loss dropped smoothly and significantly from 13.5 to 1.7 within 400+ steps, indicating that backward gradients remain correct after FSDP shard aggregation and EP routing aggregation. The learning rate and precision configuration are sound, with no signs of gradient explosion, vanishing gradients, or silent numerical degradation. LoRA itself only trains low-rank incremental parameters injected into the backbone, further compressing trainable parameters and optimizer state memory usage—complementing the full-sharding strategy.

The significance of this validation is that it proves the training pipeline for wide-gate sparse MoE models can be fully operationalized at the hundred-GPU supernode scale, laying the foundation for heavier training workloads such as full-parameter training and large-scale reinforcement learning.

Methodology: An Engineering Checklist for New Model Adaptation

The above practices can be abstracted into a reusable methodology, summarized as an adaptation checklist:

1. Perform a static memory calculation first: Before starting, account for weights, gradients, optimizer states, activations, and communication buffers to determine the minimum sharding granularity required.
2. Combine parallel strategies orthogonally: Use fully sharded data parallelism for dense parts, expert parallelism for sparse experts, and sequence-dimension chunking for ultra-long sequences. Avoid over-sharding along any single dimension.
3. Fuse operators where the impact is largest: Prioritize fusing computation paths that are memory-intensive, multi-stage, and chunkable (such as linear attention recurrence) for the most direct gains.
4. Tier memory rather than brute-force it: Leverage HBM+DDR tiering and selective caching to trade controlled access latency for capacity and bandwidth.
5. Chunk loss computation in both dimensions: Chunk both sequence and vocabulary dimensions to eliminate memory spikes during the loss stage.
6. Manage communication and load jointly: Design for uniform expert placement, dynamic load balancing, and overlapping computation with communication.
7. Use convergence curves as the final validation: Initial loss aligned with theoretical values and a smooth, anomaly-free convergence path are the ultimate arbiters of parallelization and numerical correctness.
8. Automate the process: Template and agentify the adaptation pipeline so that “day-level integration for new models” becomes the norm rather than the goal.

Conclusion

Training adaptation for ultra-large-scale sparse MoE models is fundamentally a systematic trade-off among computation, communication, and memory. Each step forward in model architecture demands corresponding evolution in parallel strategies, operator shapes, memory hierarchies, and scheduling approaches. From hybrid parallelization and fused operators to tiered caching and dual-chunked loss computation, this full-chain optimization path has been proven viable in practice. For teams planning to train or fine-tune similar models, what matters more than individual techniques is establishing a framework for mapping architectural characteristics to system bottlenecks and optimization measures—and leveraging automated tooling to minimize the ongoing cost of continuous adaptation. At the engineering deployment level, teams can either deep-tune on self-built clusters or leverage general-purpose cloud platform capabilities (such as AI compute instances, container scheduling, object storage, and monitoring/alerting) to quickly set up experimental environments and focus their efforts on the models and algorithms themselves.