Adapting and Optimizing Ultra-Large-Scale MoE Model Training: A Full-Stack Engineering Approach

1. Background: Sparse MoE as the Mainstream for Ultra-Large Models

In recent years, the capabilities of open-weight large models have expanded rapidly, with parameter scales advancing from hundreds of billions to trillions. Take the recent open-source, native multimodal sparse MoE model as an example: it has a total parameter scale of approximately 2.8 trillion, utilizing a sparse expert architecture with 896 routed experts, of which only 16 are activated per token. Its attention mechanism employs a hybrid architecture combining linear and gated sparse attention, with about three-quarters of the attention modules using linear attention to reduce long-context computation and memory overhead. It supports a context window of up to one million tokens alongside native image understanding. Leveraging large-scale joint reinforcement learning and system-level optimization, such models excel in long-horizon autonomous programming, complex agent tasks, financial analysis, and deep knowledge work.

However, these capability leaps impose extreme demands on training and fine-tuning infrastructure. For training platforms, “adaptation” goes far beyond simply loading weights and running forward passes. It requires simultaneous engineering efforts across five dimensions: model architecture parsing, hybrid parallelism, operator acceleration, communication optimization, and memory management, forming a full-link engineering capability. This article follows this thread, first breaking down the challenges and then presenting engineering solutions validated by public practices.

2. Three Core Challenges: Computation, Communication, and Memory

2.1 Computation: Hybrid Attention Complicates Fusion and Scheduling

Traditional Transformer attention calculations are highly regular and easily covered by mature fusion schemes. The new generation of hybrid attention architectures, however, introduces two complex computation paradigms:

• Linear Attention: Relies on state recursion and chunked calculations with sequential dependencies across time steps, preventing straightforward parallelization.
• Sparse Attention: Relies on sparse indexing and non-contiguous memory access, resulting in poor memory locality and low bandwidth utilization.

Because these two paradigms interleave within the same network, execution scheduling must frequently switch between them, significantly increasing the difficulty of operator fusion and scheduling. Additionally, innovative structures like attention deep residuals often lack off-the-shelf implementations in existing operator libraries, requiring the platform to possess rapid engineering capabilities for developing forward and backward operators.

2.2 Communication: Sparse Routing Inherently Causes Load Imbalance

The routing network dynamically selects 16 experts out of 896 for each token. While sparse activation reduces computation, it introduces a classic load skew problem: tokens pile up at popular experts while unpopular ones idle, degrading large matrix multiplications into inefficient, small-batch GEMMs. When experts are distributed across multiple cards, tokens must be transferred between cards via All-to-All communication, and this communication overhead scales up with the parallelism size.

2.3 Memory: 5.6TB Weights are Just the Entry Ticket

Storing 2.8 trillion parameters in BF16 precision requires approximately 5.6TB just for weights, excluding gradients, optimizer states, activations, and intermediate tensors for ultra-long contexts. A more hidden bottleneck comes from the ultra-large vocabulary: the output layer logits tensor expands based on “sequence length × vocabulary size,” causing instantaneous inflation during long-sequence training that often generates the actual peak memory usage.

3. Parallelism Strategy: Combining FSDP and Expert Parallelism

The first step in addressing these challenges is designing a rational parallel layout. Public industry practices generally adopt a hybrid scheme of “Fully Sharded Data Parallel (FSDP) + Expert Parallelism (EP)”:

• Dense Parts (Attention layers, shared layers): Uses FSDP to shard parameters, gradients, and optimizer states along the data parallel dimension. Each card retains only its shards, and reduction communication is handled by efficient collective communication primitives.
• MoE Expert Layers: Uses EP to deploy expert groups across different accelerator cards. Tokens are dispatched via All-to-All to the nodes hosting the experts and merged back after computation.
• Fine-Tuning Scenarios: Integrates LoRA, freezing the backbone and training only low-rank adaptation matrices. This compresses trainable parameters by several orders of magnitude, making fine-tuning of trillion-level models feasible on clusters with around a hundred cards.

# Pseudocode: Hybrid parallel layout
parallel_plan = {
    "dense / attention layers": FSDP(shard_size=128, wrap="transformer_block"),
    "moe.experts layers":      ExpertParallel(ep_size=64, dispatch="all_to_all"),
    "fine_tuning mode":        LoRA(freeze_weights=True),
}

Industry platforms have already validated LoRA fine-tuning for models of this scale on 128-card clusters using the FSDP+EP64 configuration, utilizing high-bandwidth, low-latency intra-node interconnects and unified memory addressing capabilities. This proves the engineering feasibility of this combination.

4. Operator Optimization: Fusing Fragmented Small Operators into Large Kernels

1. Linear Attention: Large Forward/Backward Fused Operators. The linear attention computation chain includes state initialization, recursive updates, intra-block attention, and output projection. Dispatching small operators stage-by-stage allows kernel launch overhead and repeated intermediate memory writing to devour performance. An effective approach is to integrate the entire forward (or backward) chain into a single fused kernel, keeping intermediate states in on-chip storage as much as possible to minimize global memory reads and writes.

2. Sparse Attention: Integrated Fetch-Compute-Writeback. Construct fused operators for index loading, Key/Value non-contiguous gather, attention scoring, and result aggregation, executing scattered memory access and computation in a combined manner. Public practices show that such deep fusion can achieve over 2x performance improvement compared to assembling small operators.

3. Rapid Integration of Innovative Structures. For new structures like attention deep residuals, compiler-based operator development frameworks can be used to describe computation logic with high-level DSLs and automatically generate high-performance implementations. This accelerates forward and backward operator development and automatic differentiation integration, compressing the support cycle for new structures from months to days.

4. Core MoE Full-Link Operators. Permute (reordering tokens by expert), Unpermute (restoring original order), and Grouped GEMM (expert-grouped matrix multiplication) are the performance choke points of MoE training. Conducting deep fusion and high-performance implementation around these three operators, while streamlining the complete execution chain of “Token Dispatch → Expert Compute → Combine,” is key to improving end-to-end throughput.

5. Communication Optimization: Maximizing All-to-All Utilization and Preventing Idle Compute

• High-Speed Interconnect: Leveraging high-bandwidth, low-latency intra-node interconnects and unified memory addressing to optimize All-to-All dispatch and retrieval paths, reducing token transfer latency.
• Communication-Computation Overlap: Splitting communication and computation into parallel micro-tasks, using a pipeline approach to overlap communication with computation and improve hardware time utilization.
• Expert Load Balancing: Constraining routing distribution via auxiliary loss, combined with dynamic capacity factors and overflow control to mitigate popular expert skew. When necessary, expert redundancy or dynamic replacement strategies are used to bring single-expert batch sizes back into efficient ranges.

6. Memory Optimization: Operator-Level Selective Caching and Tiered Storage

Beyond conventional methods like parameter sharding and gradient recomputation, a notable practice is the “operator-level selective caching” mechanism:

• Identify key MatMul and grouped matrix multiplication operators in the training computation graph and establish L2 caches for their results.
• The caching medium supports tiered storage across HBM and DDR: high-frequency reused data resides in HBM, while low-frequency data offloads to DDR, paged in and out on demand.
• When recomputation is needed during backpropagation, it prioritizes cache hits to avoid executing the full computation repeatedly.

Public tests indicate that this mechanism can yield an average training performance improvement of around 15% under typical configurations. Its essence is “trading bandwidth for compute,” which is particularly beneficial in ultra-large-scale MoE scenarios with high recomputation overheads and tight memory constraints.

7. Large Vocabulary Loss Optimization: Dual Chunking to Reduce Peak Memory

Facing ultra-large vocabularies, public practices frequently employ a dual chunking scheme across “sequence dimension + vocabulary dimension”:

1. Chunk Loss (Sequence Dimension Chunking): Splits long sequences into several chunks, calculating logits and cross-entropy chunk by chunk, then aggregating them with weights.
2. Vocabulary Dimension Chunking (Cross-Chunked Cross-Entropy Schemes): Slices the vocabulary, calculating logits and local softmax slice by slice, utilizing streaming updates via online log-sum-exp to ensure numerical equivalence.

# Pseudocode: Sequence dimension chunked Loss
total = 0.0
for i, chunk in enumerate(hidden_states.split(seq_chunk_size)):
    logits = lm_head(chunk)                # [chunk_len, vocab]
    total += ce_loss(logits, labels[i*chunk:(i+1)*chunk])
loss = total / seq_len

By combining these two approaches, the output layer peak memory is reduced from a “full sequence × full vocabulary” scale down to a single chunk scale, making ultra-long sequence training possible.

8. Engineering Methodology: Pipelining New Model Integration

As new architectures emerge endlessly, the pace of manual adaptation cannot keep up with open-source model releases. A viable path is building an automated development pipeline: automated processes (or even agent-driven development paradigms) handle model structure parsing, operator mapping checks, forward/backward numerical alignment, and performance baseline regression. Engineers only intervene for exception handling and deep tuning. Industry practices show that with such workflows, new model integration can be compressed from weeks to days, supporting continuous and rapid iteration.

9. Verification Loop and Conclusion

A credible training adaptation must complete at least three types of verification: numerical correctness (aligning the loss curve with the reference implementation), convergence (downstream task metrics meeting standards), and performance (end-to-end throughput and hardware utilization). Building full-link capabilities centered on “model adaptation — hybrid parallelism — operator acceleration — communication optimization — memory management” is becoming the core competitiveness of ultra-large-scale MoE model training platforms.

For enterprise teams looking to quickly implement such practices, leveraging cloud platforms that provide general infrastructure like large-scale training clusters, high-speed interconnect networks, parallel file storage, and containerized scheduling (such as Baidu AI Cloud’s compute services) can significantly lower the barrier of building and tuning from scratch, allowing teams to focus on the model value itself.