Training Trillion-Scale MoE Models in Practice: From Sparse Architecture Adaptation to Efficient Training

From Model Scale to Engineering Challenges: What Sparse Architectures Change

In recent years, the parameter count of large language models has grown exponentially. The rise of Mixture of Experts (MoE) architectures has made it possible to decouple model capability from active parameter count: total parameters can be enormous, while only a fraction are activated for each inference. This sparse activation dramatically increases model capacity, but it also shifts the engineering challenge from “can we compute it at all” to “can we schedule it well.”

A new generation of native multimodal MoE models carries nearly 3 trillion total parameters, with tens of billions activated at any time. The model follows the Stable LatentMoE design, consisting of 896 routed experts, with each token activating only 16 of them. This means that during both training and inference, the system must constantly exchange token-to-expert mappings across compute devices.

From an architectural perspective, the key innovation is a hybrid attention mechanism: roughly three-quarters of the Attention Blocks use linear attention (KDA), while the remainder use traditional sparse attention (DSA). The former reduces complexity from quadratic to linear through state recursion and chunked computation; the latter relies on Top-2048 sparse indexing to preserve precise long-context modeling. Interleaving the two not only complicates operator fusion but also imposes sharply different requirements on inter-device communication patterns.

Supernode Architecture: The Foundation of High Bandwidth and Unified Memory Addressing

At this scale, single-node multi-GPU setups fall far short, and a supernode-level compute cluster becomes essential. The supernode architecture offers two core advantages. First, high-bandwidth interconnects between nodes keep the communication overhead of tensor parallelism and expert parallelism within acceptable limits. Second, unified memory addressing allows GPU memory distributed across multiple devices to be treated logically as a single pool, greatly simplifying parameter sharding and data management for massive models.

On top of this architecture, the training system needs end-to-end capabilities spanning five layers: model splitting, parallelism strategy, operator optimization, communication tuning, and memory management. Notably, after introducing an Agent-driven automated development workflow, a new model can be onboarded and iterated within days rather than the months previously required.

Three Core Challenges: Co-Optimizing Compute, Communication, and Memory

Compute: Operator-Level Challenges of Heterogeneous Attention

KDA linear attention computes very differently from standard attention. Standard attention maps naturally onto batched matrix multiplications, while KDA involves state recursion and chunked incremental computation, where each step depends on the state produced by the previous step. As a result, the operator fusion strategies commonly used on GPUs cannot be applied directly.

For KDA, an effective approach in practice is to build “large-forward, large-backward” fused operators — consolidating multiple computation stages, such as state initialization, recursive updates, and output projection, into a single kernel. This not only reduces kernel launch overhead but, more importantly, keeps intermediate states in L2 cache and avoids frequent HBM accesses.

DSA sparse attention presents the opposite problem: each token dynamically selects the Top-K historical positions by confidence score for attention computation. These positions are not stored contiguously in memory, generating large volumes of non-contiguous memory accesses. Executing this with traditional small operators is extremely inefficient; in practice, switching to a custom operator that fuses sparse indexing with Gather operations improved performance by more than 100%.

Notably, the model also introduces a deep residual mechanism at the attention layers (AttnRes), which passes raw residuals directly between attention layers instead of the values after LayerNorm. This novel structure can be implemented quickly through operator customization in the compilation framework, shortening the validation cycle for new architectures.

Communication: The Routing Dilemma of 896 Experts

MoE layers are where communication bottlenecks concentrate. When each token must select 16 experts out of 896, token distribution across experts naturally exhibits a long-tail effect — a few hot experts handle far more than the average load while cold experts sit idle. This imbalance directly causes two problems:

1. Inefficient token dispatch: All-to-All communication occurs frequently, and each round carries an imbalanced load, easily congesting certain links.
2. Undersized per-expert batches: When an expert receives too few tokens, Grouped GEMM struggles to reach the peak performance of tensor cores.

In practice, optimization typically proceeds along three directions simultaneously: first, optimizing the memory layout of Permute/Unpermute operators to reduce data movement; second, overlapping communication with computation, i.e., executing locally available compute tasks while waiting for tokens to arrive at their experts; and third, introducing an expert load-balancing loss that dynamically adjusts token assignment during training.

Combining these three dimensions effectively unblocks the full Dispatch-Compute-Combine pipeline — from token dispatch through computation to aggregation — yielding significant gains in overall training throughput.

Memory: The Storage Economics of 2.8T Parameters

Even storing weights in BF16 alone, 2.8 trillion parameters require roughly 5.6 TB of GPU memory. This means that regardless of how the parallelism strategy is designed, the combined footprint of parameters, gradients, and optimizer states approaches the physical limit of a single AI server.

For models at this scale, a hybrid strategy combining FSDP2 and expert parallelism (EP) is the mainstream solution. FSDP2 shards non-expert parameters (such as attention and dense layers), while EP distributes the 896 experts evenly across compute devices. The elegance of this combination lies in the natural fit: attention layers depend on the full sequence context and suit data parallelism, whereas expert layers are inherently parallel compute units and suit expert parallelism to spread both storage and compute pressure.

Experimental results show that on a 128-GPU supernode, an EP64 configuration can stably support LoRA fine-tuning. As a parameter-efficient fine-tuning method, LoRA freezes most of the weights and trains only a small set of low-rank increment matrices, drastically reducing memory and compute requirements and making scenario-specific adaptation of ultra-large models feasible.

Deep Memory Optimization: Swap Core and Chunk Loss

In ultra-large model training, activation memory is typically the second-largest consumer after parameters. In particular, intermediate results cached in GMM (grouped matrix multiplication) operators can occupy substantial HBM space.

Swap Core: Operator-Level Selective Caching

The core idea is a two-level caching mechanism: during computation, designated MatMul or GMM operators are given an optional output cache layer that stores intermediate results in HBM first and automatically spills them to DDR memory when HBM space runs out. This resembles virtual memory paging in an operating system, but implemented at the operator level.

The key to this “selective caching” strategy is that users can precisely control, based on model structure and training phase, which operator outputs should be cached and which can be discarded outright to save space. By dynamically scheduling between HBM and DDR, training throughput improves by roughly 15% without a significant latency penalty.

Chunk Loss and CCE Loss: From Sequence to Vocabulary Dimensions

Ultra-large models typically have vocabularies in the hundreds of thousands or even millions, making the final softmax layer extremely expensive. Standard cross-entropy requires computing the probability distribution over the entire vocabulary, incurring substantial memory and compute costs.

Chunk Loss splits the input sequence along its length into multiple chunks, computes the loss for each chunk separately, and then merges the gradient updates. This reduces the peak memory required for a single computation and adds flexibility for gradient accumulation along the sequence dimension.

CCE Loss approaches the problem from the vocabulary side: it partitions the vocabulary into smaller shards, computes a local softmax for each shard independently, and approximates the global softmax through cross-shard gradient synchronization. The two strategies can be stacked, reducing peak memory along both the sequence length and vocabulary size dimensions and making ultra-large batch training possible.

Practical Lessons and Takeaways

Having completed LoRA fine-tuning validation on a 128-GPU supernode with a hybrid parallel architecture, we can distill several lessons worth reusing:

1. Parallelism strategy must be chosen based on model architecture, not convention. For MoE models, the interplay between parameter sharding (FSDP) and expert parallelism (EP) requires repeated tuning against expert count, number of activated experts, per-expert batch size, and similar parameters.
2. Operator fusion is the most direct lever for performance. In large model training, the bottleneck is often memory access rather than compute. A well-designed large-kernel fusion plan merges many small kernels, effectively improving L2 cache hit rates and reducing API launch overhead.
3. Memory optimization requires a global view. Swap Core, Chunk Loss, and CCE Loss act on compute intermediates, the sequence dimension, and the vocabulary dimension respectively. When combined, their interactions must be carefully evaluated to avoid gradient truncation or excessive growth in communication volume.
4. Automated model onboarding is indispensable. When model iterations move at a weekly or even daily cadence, manually writing operator integration code is no longer realistic. Agent-driven automated compilation and tuning tools can make the mapping from model structure to training scripts nearly automatic, substantially boosting development efficiency.

Looking ahead, as parameter counts continue toward the tens of trillions, techniques such as sparsification, tiered storage, and heterogeneous computing will no longer be optional optimizations but the necessary infrastructure supporting continued model evolution. The work described here provides a reusable engineering methodology for that path.