Challenges and Optimization Practices for Fine-Tuning Ultra-Large-Scale MoE Models

1. Background: The Adaptation Problem Posed by Trillion-Parameter Models

In recent years, the parameter scale of open-weight large models has continued to climb. New-generation flagship models have entered the multi-trillion parameter range, commonly adopting sparse Mixture-of-Experts (MoE) architectures and native multimodal designs. For training infrastructure, these models are far from a simple matter of “scaling up the old workflow.” Sparse routing, linear attention, gated sparse attention, and ultra-long context windows introduce systematic challenges for kernel implementation, parallelism strategies, communication patterns, and memory management.

This article uses the adaptation and fine-tuning of a recently released 3-trillion-class multimodal MoE model (approximately 2.8 trillion total parameters) as a case study. It breaks down the engineering bottlenecks across three dimensions — computation, communication, and memory — and summarizes widely applicable optimization methodologies and implementation paths.

2. Architectural Characteristics of the Target Model

Before designing an optimization strategy, it is essential to understand the model’s key architectural features:

• Sparse expert architecture: The model has 896 routed experts, with only 16 experts activated per token. Sparse activation keeps per-token computation manageable while maintaining enormous total model capacity, but it inherently complicates routing decisions and load balancing.
• Hybrid attention design: Approximately three-quarters of the attention modules use linear attention with state recurrence and block-wise computation. The remaining modules use gated sparse attention, which performs precise computation only on the Top-2048 key positions. This combination of linear approximation and precise computation on critical positions supports a context window of up to 1 million tokens and native image understanding.
• Training paradigm and capability boundaries: Built on large-scale joint reinforcement learning and system-level optimization, models of this class excel at long-horizon autonomous programming, complex agent tasks, 3D creation, financial research, and deep knowledge work. This makes them popular base models for fine-tuning and secondary development.

3. Three Core Engineering Challenges

3.1 Computation: Interleaved Attention Types Increase Fusion and Scheduling Complexity

Linear attention requires kernels to maintain recurrent states that propagate across blocks, with computation organized block by block. Sparse attention, by contrast, relies on dynamic indexing for non-contiguous memory access. These two fundamentally different computation patterns alternate within the same network, leading to a large number of fragmented kernels and frequent materialization of intermediate results in memory. This significantly increases the difficulty of kernel fusion and execution scheduling. If a generic implementation composed of small per-segment kernels is used, memory access overhead becomes the throughput bottleneck.

3.2 Communication: Uneven Expert Token Distribution and Small Per-Expert Batches

Sparse routing across 896 experts naturally creates a “hot expert” problem: different experts receive vastly different numbers of tokens, and some experts have input batches too small for grouped matrix multiplication to reach efficient compute granularity. At the same time, All-to-All communication is highly sensitive to network latency and bandwidth. Uneven load distribution directly causes some compute cards to idle while waiting, slowing down the overall iteration speed.

3.3 Memory: Weight Volume Approaches Cluster Limits

Storing 2.8 trillion parameters in BF16 alone requires approximately 5.6 TB of memory, far exceeding the capacity of a single GPU or even a single node. When optimizer states, gradients, and activations are added, memory pressure persists throughout the entire training process. Memory is therefore not a single-point optimization target but the primary constraint that determines the parallelism strategy design.

4. Parallelism Strategy: Hybrid Design of Fully Sharded Data Parallelism and Expert Parallelism

To address the dual challenges of memory and communication, the industry-standard approach is a hybrid architecture combining data parallelism sharding with expert parallelism:

• Fully sharded data parallelism (FSDP): Dense parameters, gradients, and optimizer states are sharded along the data-parallel dimension. Each GPU holds only its shard, with communication primarily consisting of shard aggregation.
• Expert parallelism (EP): The 896 experts are distributed across different compute cards, with All-to-All operations handling token dispatch and result collection, ensuring that “each expert resides in exactly one place.”

Illustrative pseudocode:

for layer in model.layers:
    if layer.type == "dense":
        wrap_sharded(layer)                       # Dense layers use fully sharded data parallelism
    else:                                         # MoE layers
        experts = shard(layer.experts, ep_size=64)  # Shard experts across EP group
        wrap_sharded(experts, dp_group)           # Expert parameters further sharded by data parallelism
        replicate(layer.router)                   # Router parameters are replicated

Under this architecture, LoRA fine-tuning is a key cost-reduction technique: all base model weights are frozen, and only injected low-rank matrices are trained. This can reduce memory requirements for optimizer states and gradients by an order of magnitude, enabling trillion-parameter model fine-tuning on clusters of roughly a hundred GPUs. Public cases exist of completing LoRA fine-tuning of a 2.8-trillion-parameter model on a 128-GPU supernode cluster using fully sharded data parallelism with 64-way expert parallelism.

5. Kernel-Level Optimization: Fusion Is the First Principle

5.1 Linear Attention: Large Forward/Backward Fused Kernels

The forward pass of linear attention includes multiple stages such as normalization, state updates, and block-wise matrix operations. The backward pass additionally requires recomputing intermediate states. Integrating these stages into a single large kernel — the “large forward, large backward” fusion approach — eliminates memory reads and writes of intermediate tensors and reduces kernel launch overhead. This is the primary path to improving linear attention efficiency.

5.2 Sparse Attention: Multi-Stage Fused Kernels

Token selection, index construction, and weighted computation in sparse attention can also be fused into a small number of large kernels. Public industry practices show that designing separate fused kernels for the two attention types yields more than 2x performance improvement compared to executing small per-segment kernels.

5.3 New Structures: Rapid Integration via Compiler Frameworks

Ultra-large-scale models often introduce novel structures such as deep residuals. Hand-writing all forward and backward kernels is prohibitively expensive. Using general-purpose compiler frameworks and low-level kernel programming interfaces, and following a compilation pipeline of “operator expression → automatic differentiation → code generation,” can reduce the development cycle for forward and backward passes of new structures from weeks to days, significantly improving model iteration speed.

6. End-to-End MoE Optimization: From Dispatch to Combine

The MoE forward pass can be abstracted into three stages: Token Dispatch → Expert Compute → Combine. Performance optimization must address the entire pipeline:

1. Deep fusion of core kernels: Permute (reordering tokens by expert), Unpermute (restoring original order), and Grouped GEMM (grouped matrix multiplication) are the three hotspots. They require deep fusion and high-performance implementations to eliminate redundant copies and repeated indexing.
2. All-to-All communication optimization: Use high-performance collective communication implementations combined with communication-computation overlap. While expert computation is in progress, asynchronously send and receive the next batch of tokens to hide communication latency.
3. Expert load balancing: Mitigate the hot expert problem through capacity factor control and auxiliary load-balancing losses to prevent some compute cards from idling.

def moe_forward(tokens):
    scores, topk = router(tokens)                 # Router scoring
    perm, seg = permute(tokens, topk)             # Reorder tokens by expert
    async_send_recv(perm)                         # Asynchronous All-to-All dispatch
    out = grouped_gemm(perm, seg, expert_weights) # Grouped computation
    out = unpermute(out, perm)                    # Restore token order
    return combine(out, scores, topk)             # Weighted merge output

7. Memory Optimization: Tiered Caching and Loss Sharding

7.1 Operator-Level Selective Secondary Caching

To address the memory-bandwidth conflict caused by weight residency, an operator-level selective caching mechanism can be introduced: cache the computation results of specified operators such as MatMul and grouped GEMM, and support tiered storage across HBM (high-bandwidth memory) and DDR (host memory). Hot data stays in HBM, while cold data spills to DDR. Industry practices show that this mechanism delivers an average training performance improvement of approximately 15% under typical training configurations.

7.2 Large Vocabulary Loss Sharding: Chunk Loss and CCE Loss

Vocabulary sizes for ultra-large-scale models often reach hundreds of thousands of tokens, making the logits tensor enormous in both sequence and vocabulary dimensions. The standard solution is dual-dimension sharding:

• Chunk Loss (sequence-dimension sharding): Split long sequences into multiple chunks, compute loss chunk by chunk, and accumulate the results. This avoids materializing the full logits tensor at once.
• CCE Loss (vocabulary-dimension sharding): Split the vocabulary into column-wise blocks and compute cross-entropy in batches to control peak activation memory per step.

def chunked_ce_loss(logits, labels, seq_chunk, vocab_chunk):
    total = 0
    for s in split(logits, seq_chunk):               # Sequence-dimension chunking
        for v in split(s, vocab_chunk, axis=-1):     # Vocabulary-dimension chunking
            total += cross_entropy(v, labels_slice)  # Incremental loss accumulation
    return total / num_tokens

When combined, these two techniques can keep output-layer peak memory within a controllable range, making training with million-token contexts engineeringly feasible.

8. Engineering Implementation Recommendations

• Pipeline automation: To achieve the goal of “new model onboarding within days,” the industry widely uses agent-driven automated pipelines that connect model structure parsing, kernel generation, accuracy alignment, and performance regression. This reduces manual intervention to key decision points and significantly shortens iteration cycles.
• Cloud-based practice: Teams without access to on-premises clusters can leverage general-purpose infrastructure from providers such as Baidu AI Cloud — using container platforms for training job orchestration, object storage and file storage for managing massive weights, message queues for data pipelines, and monitoring and alerting for cluster health tracking — allowing teams to focus on model architecture and kernel-level optimization.

9. Summary

Training adaptation for ultra-large-scale sparse MoE models is fundamentally an “architecture-driven systems engineering” endeavor: fusion on the compute side, balancing and overlap on the communication side, and sharding and tiering on the memory side. By mastering five key capabilities — the hybrid strategy of fully sharded data parallelism plus expert parallelism, large-granularity fused kernels, communication-computation overlap, operator-level selective caching, and dual-dimension loss sharding — and supplementing them with automated engineering workflows and general-purpose cloud infrastructure, teams can maintain low-cost, fast-paced adaptability as trillion-parameter models continue to evolve.