Training Adaptations for Ultra-Large-Scale MoE Models: From Operator Fusion to Hierarchical Memory Management

1. Background: Engineering Challenges Posed by Trillion-Scale Sparse Models

Recent open-weight native multimodal models have reached total parameter counts of approximately 2.8 trillion. These models use a sparse Mixture-of-Experts (MoE) architecture with 896 routed experts, where each token activates only 16 experts. This design keeps per-token computation manageable while maintaining very large parameter capacity.

The attention mechanism combines linear attention and sparse attention: roughly three-quarters of the attention blocks use linear attention, while the remaining blocks combine Top-2048 sparse attention. The architecture supports context windows of up to 1 million tokens and includes native image understanding. With large-scale joint reinforcement learning and system-level optimization, these models show measurable gains in long-horizon autonomous programming, complex agent tasks, 3D content creation, financial research, and deep knowledge work.

From an infrastructure perspective, however, adapting and training such models goes far beyond “loading weights and launching a job.” At BF16 precision, the weights alone consume roughly 5.6 TB of memory. Sparse routing across 896 experts introduces inherent load imbalance. Linear attention requires recurrent state updates, while sparse attention involves non-contiguous memory access. Both place heavy demands on operator fusion and execution scheduling.

This article examines the training adaptation methodology for ultra-large-scale MoE models from three angles: computation, communication, and memory.

2. Overall Approach: End-to-End Training Capability Plus Automated Development Workflows

Adaptation work should leverage architectural advantages such as high bandwidth, low latency, and unified memory addressing in large-node systems to build end-to-end training capabilities spanning model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management.

At the same time, agent-driven automated development workflows enable rapid onboarding and iteration of new models. This matters in practice because model architecture innovation often outpaces manual adaptation speed.

For the parallelization strategy, the core approach is a hybrid FSDP2 + EP architecture:

• FSDP2 (Fully Sharded Data Parallel) shards parameters, gradients, and optimizer states.
• EP (Expert Parallelism) distributes expert parameters and enables efficient aggregation.

In a representative validation configuration, LoRA fine-tuning with FSDP2+EP64 has been completed on a 128-card node, providing a reference baseline for larger-scale training.

3. Computation: Operator Fusion Challenges for Hybrid Attention

3.1 Linear Attention: Large Forward and Backward Fusion

The core of linear attention lies in recurrent state updates and chunked computation. If recurrence, chunking, and normalization are split into many small operators in the conventional way, the result is frequent kernel launch overhead and excessive intermediate tensor reads and writes.

The optimization approach is to design large fused forward and backward operators that combine multiple computation stages into a single kernel, reducing memory traffic and scheduling overhead.

3.2 Sparse Attention: Fused Operators Replace Small Operator Combinations

Sparse attention involves sparse indexing and non-contiguous memory access, which interleave with the computation patterns of linear attention and further complicate operator fusion and scheduling. In practice, fused operators such as LI/SFA/SFAG can be used for model integration. Compared with the original combination of small operators, performance can improve by more than 2x.

3.3 Deep Residual Mechanisms: Rapid Adaptation Through Compiler-Based Frameworks

For novel structures such as attention residuals, compiler-based frameworks like Triton and compatible domestic compilers can accelerate forward and backward operator development and integration. The advantage of this compiler-based path is that when a model introduces new structural variants, usable operators can be generated quickly through a high-level DSL rather than writing low-level kernels from scratch, significantly shortening the adaptation cycle.

4. Communication: Load Balancing for Expert Routing

Sparse routing across 896 experts creates two common problems: uneven token distribution across experts (some experts overloaded while others idle) and small per-expert batch sizes (reducing matrix multiplication efficiency).

To address MoE training bottlenecks across the full pipeline, deep fusion and high-performance optimization should target core operators such as Permute, Unpermute, and Grouped GEMM. These efforts should be combined with high-performance All-to-All communication, communication-computation overlap, and expert load balancing to streamline the full Token Dispatch → Expert Compute → Combine execution path.

A simplified illustration of communication-computation overlap:

# Pseudocode: basic idea of overlapping communication and computation
def moe_forward(tokens, experts):
    # 1. Local routing computation
    routing = gate(tokens)
    # 2. All-to-All dispatch (asynchronous)
    handle = async_all_to_all(tokens, routing)
    # 3. While waiting for communication, process locally available tokens
    local_out = grouped_gemm(local_tokens, local_experts)
    # 4. Wait for communication to complete and merge results
    remote_out = handle.wait()
    return combine(local_out, remote_out)

5. Memory: Parameter Sharding and Hierarchical Caching

At BF16 precision, 2.8T parameters occupy roughly 5.6 TB for weights alone—far beyond the capacity of a single accelerator. Beyond FSDP2+EP parameter sharding, two additional directions can further optimize memory usage.

First, operator-level selective caching. A Swap Core mechanism introduces second-level caching for the outputs of designated MatMul/GMM operators, supporting hierarchical HBM and DDR storage. Under current test configurations, training performance improves by approximately 15% on average. The core idea is to offload certain intermediate results from high-bandwidth memory to a larger but lower-bandwidth storage tier, trading controllable bandwidth cost for more available memory space.

Second, chunked loss for large vocabularies. To address the memory bottleneck caused by large vocabularies, Chunk Loss and CCE Loss provide dual optimization along the sequence and vocabulary dimensions. Chunk Loss chunks computation along the sequence dimension, while CCE Loss chunks along the vocabulary dimension, avoiding materialization of the full logits tensor at once.

6. Summary and Engineering Recommendations

Training adaptation for ultra-large-scale MoE models is ultimately a systems engineering effort centered on compute efficiency, communication efficiency, and memory efficiency. Based on the practices described above, several reusable engineering principles emerge:

1. Define the parallelization strategy first. The FSDP2+EP combination achieves a good balance between parameter sharding and expert parallelism, making it suitable for scenarios with many experts and large per-expert parameter counts.
2. Prioritize operator fusion. The hybrid structure of linear attention and sparse attention requires a fused-operator approach. Otherwise, small-operator scheduling overhead will consume most of the potential gains.
3. Communication-computation overlap is mandatory. All-to-All communication accounts for a high proportion of MoE execution time. Asynchronous execution and overlap are essential to hide latency.
4. Memory optimization must be hierarchical. From parameter sharding to operator-level caching to large-vocabulary chunking, a multi-level memory management system is required.
5. Automate the adaptation workflow. Agent-driven development processes can compress the onboarding cycle for new models to days, which is necessary to keep pace with model iteration speed.

For teams working on production deployment of ultra-large models, this methodology provides a practical path from “able to run” to “running efficiently.” As model scale continues to grow, the competitiveness of training infrastructure will increasingly depend on these seemingly low-level engineering details.