Adapting Domestic SuperNodes for Trillion-Scale Sparse MoE Models: From Operator Fusion to Full-Stack Parallelism

1. Background: Crossing the “Trillion” Threshold

As native multimodality and ultra-long contexts become the mainstream evolution path for large models, parameter scales are rapidly advancing from hundreds of billions to the trillion-level. A recent flagship open-weight model adopted a sparse Mixture-of-Experts (MoE) architecture with approximately 2.8 trillion total parameters. It features nearly 900 routing experts, with each token activating only about a dozen of them. Additionally, it introduces a hybrid attention mechanism combining linear and sparse attention. Approximately three-quarters of the attention blocks utilize linear attention, paired with Top-2048 sparse attention indexing, enabling a context window of up to a million tokens and native image understanding capabilities.

These models demonstrate significant capability leaps in scenarios like long-range autonomous programming, complex agent tasks, 3D content creation, financial research, and deep knowledge work. However, this evolution places immense pressure on the training system across three dimensions: computation, communication, and memory. From a general technical perspective, this article outlines the key engineering methods for adapting and fine-tuning such trillion-scale sparse MoE models on domestic SuperNode architectures.

2. Analysis of Three Core Bottlenecks

2.1 Computation: Scheduling Challenges from Heterogeneous Attention

In hybrid attention architectures, the linear attention component relies on state recursion and chunked computation, creating a computation graph with strict sequential dependencies. Conversely, the sparse attention component relies on sparse indexing and non-contiguous memory access, resulting in highly irregular memory access patterns. When intertwined, the complexity of operator fusion and execution scheduling increases significantly. A large number of fine-grained, small operators cannot be efficiently pipelined, making kernel launch overhead and memory access redundancy the primary sources of performance degradation.

2.2 Communication: Load Imbalance in Sparse Routing

Sparse routing across nearly 900 experts often leads to uneven token distribution during actual training. Popular experts may process tokens far exceeding the average, while less popular ones remain nearly idle. This results in two major consequences: first, the computation batch size per expert becomes too small, reducing GPU/NPU utilization; second, pronounced long-tail effects emerge in All-to-All communication, dragging down overall throughput as the system waits for the slowest links.

2.3 Memory: Hard Constraints of Parameter Volume

Storing 2.8 trillion parameters in BF16 precision alone requires approximately 5.6TB of memory. This far exceeds the capacity of a single card. It is mandatory to distribute weights across the entire SuperNode cluster using expert parallelism and parameter sharding, while simultaneously reserving space for activations, optimizer states, and communication buffers.

3. SuperNode Architecture Solutions

Domestic SuperNodes (such as the Atlas 950 SuperPod cluster architecture) possess architectural advantages in high bandwidth, low latency, and unified memory addressing, providing a systemic foundation to resolve the aforementioned bottlenecks. For trillion-scale MoE training, a full-stack capability covering model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management can be constructed. Supported by an agent-driven automated development flow, this enables day-level integration and rapid iteration for new models.

3.1 Hybrid Parallelism: The FSDP2 + EP Combination

In practice, a hybrid parallel architecture utilizing FSDP2 (Fully Sharded Data Parallel 2) + EP (Expert Parallelism) is employed:

• FSDP2 handles the sharding of non-expert parameters (attention, embeddings, normalization, etc.), reducing the parameter residency pressure on individual cards.
• EP distributes hundreds of experts across different devices, achieving distributed deployment and efficient aggregation of expert parameters.

On a 128-card SuperNode, a LoRA fine-tuning training path based on FSDP2 + EP64 has been successfully verified, demonstrating the feasibility of this parallel combination for ultra-large-scale MoE models.

3.2 Attention Operators: Fusion as the Main Theme

For the state recursion and chunked computation in linear attention, large forward and large backward fused operators are designed. These integrate multiple previously scattered computation stages into a single kernel, reducing the repeated read/write operations of intermediate results.

For sparse attention, a series of fused operators such as LI, SFA, and SFAG are utilized for model integration. Compared to implementations using fragmented small operators, the fused approach achieves over 2x performance improvement in testing.

For innovative structures like deep residuals, leveraging Triton-like compilation frameworks and BiSheng-like compilers enables rapid development and integration of forward and backward operators, significantly shortening the deployment cycle for new structures.

3.3 Full-Stack MoE: From Dispatch to Combine

The performance bottleneck in MoE training spans the entire chain. It requires deep fusion and high-performance optimization around core operators like Permute, Unpermute, and Grouped GEMM, coupled with:

• High-performance All-to-All communication;
• Communication and computation overlap;
• Expert load balancing strategies.

The synergy of these three elements is essential to streamline the full execution path from Token Dispatch to Expert Compute to Combine, ensuring no single环节 becomes a bottleneck.

4. Two Key Memory Optimization Strategies

4.1 Operator-Level Selective Caching

A Swap Core mechanism is introduced to apply secondary caching to the computation results of specified MatMul / GMM operators, supporting hierarchical storage between HBM and DDR. Under test configurations, this mechanism yielded an average training performance improvement of approximately 15%. Fundamentally, it trades a larger-capacity, lower-bandwidth storage layer to free up high-bandwidth memory, thereby allowing for larger batch sizes or longer sequences.

4.2 Chunked Loss for Large Vocabularies

In scenarios with large vocabularies, the memory footprint of the output layer’s logits is substantial. Double optimization is achieved through Chunk Loss and CCE Loss:

• Chunk Loss partitions along the sequence dimension, avoiding the materialization of intermediate loss variables for the entire sequence at once.
• CCE Loss computes in chunks along the vocabulary dimension, preventing the full vocabulary logits from residing in memory.

Combined, these methods drastically compress the memory peak of the output layer with virtually no loss in numerical precision.

5. A Reproducible Fine-Tuning Validation

A LoRA fine-tuning practice was conducted on a 128-card SuperNode, yielding the following results:

• The initial loss from random initialization training was approximately 13.5;
• After over 400 steps of training, the loss converged to around 1.7.

This curve demonstrates that the combination of parallel strategies, operator fusion, and memory optimizations can support the stable convergence of trillion-scale sparse MoE models, proving they are not just “runnable” but practically trainable.

6. Engineering Best Practices

1. Fix Communication Before Computation: MoE training bottlenecks often surface in All-to-All communication first. Prioritizing communication-computation overlap and load balancing typically yields higher returns than solely optimizing GEMM.
2. Prioritize High-Impact Operator Fusion: Focus on fusing attention and core MoE operators that are called frequently and generate large intermediate results, as these offer the highest ROI.
3. Layered Memory Optimization: Start with parameter sharding, followed by activation recomputation, and finally utilize operator-level caching and chunked loss to extract remaining memory capacity.
4. Leverage Automated Integration to Reduce Iteration Costs: Utilize an agent-driven development workflow to standardize operator development, graph rewriting, and precision alignment for new model structures, enabling day-level integration.

7. Conclusion

Training trillion-scale sparse MoE models is fundamentally a systems engineering challenge characterized by heterogeneous computation, intensive communication, and strict memory constraints. Single-dimension optimizations are insufficient; collaborative design across parallelism strategies, operator fusion, communication scheduling, and memory management is mandatory. Domestic SuperNodes, with their architectural features of high bandwidth, low latency, and unified memory addressing, provide a viable deployment path for such models. Ultimately, the engineering methodologies centered around hybrid parallelism and full-stack operator optimization determine whether a system can transition from merely being able to train, to training fast and training stably.