Engineering Practices for Training and Adapting Trillion-Parameter MoE Multimodal Models

1. Background: New Demands on Training Systems from Trillion-Parameter Open-Weight Models

In recent years, leading teams have released open-weight models at the trillion-parameter scale, natively multimodal, and supporting million-token context. While model capability ceilings keep rising, these releases impose requirements on enterprise-side training and fine-tuning systems that far exceed previous norms—parallelism strategies, operator implementations, and memory management schemes originally designed for dense models largely need to be rebuilt from scratch when facing ultra-large sparse MoE (Mixture of Experts) architectures.

Taking a next-generation open-weight native multimodal model as an example, its key specifications are:

• Total parameter count of roughly 2.8 trillion, using a stable sparse expert architecture with 896 routed experts, activating only 16 experts per token;
• Attention layers use a hybrid design of “linear attention + gated sparse attention”: about three-quarters of the Attention Blocks are linear attention, and the rest are Top-K sparse attention (each Query interacts with only on the order of a thousand Keys);
• Supports up to 1 million tokens of context and has native image understanding capability.

Translating these specifications into system-side pressure gives an intuitive sense of the difficulty:

• BF16 weights alone are about 5.6 TB, far exceeding any single machine’s memory capacity, so parameters must be distributed;
• The 16/896 activation ratio means tokens must be dynamically grouped by routing results and scheduled across devices;
• The state recurrence of linear attention and the non-contiguous memory access of sparse attention interleave, dramatically increasing the difficulty of operator fusion and execution scheduling.

Therefore, adapting such models is far from just “getting it to run”—it is a systems engineering effort coordinating three fronts: compute, communication, and memory. The following sections break down the bottleneck mechanisms one by one and provide reusable engineering solutions.

2. Compute Side: Hybrid Attention Brings Operator Fragmentation and Scheduling Challenges

First, linear attention. By compressing Key/Value information into a fixed-size recurrent state, it reduces the O(N²) interaction complexity of standard attention to O(N), making it a key technology for supporting million-token long context. But its computation flow naturally contains three stages: intra-chunk interaction computation, cross-chunk state recurrence, and state readout. If split into small operators executed serially in the traditional way, intermediate states are repeatedly read from and written to HBM; and the sequential dependency of state recurrence further complicates backpropagation scheduling.

Next, sparse attention. Top-K sparse indexing means each Query interacts with only the most relevant small subset of Keys: first similarity filtering (essentially sorting/selection), then gathering data by index to complete scoring and aggregation. The entire flow is a highly non-contiguous memory access pattern; if pieced together from small operators like “filter—gather—score—aggregate,” kernel launch overhead and random memory access will drown out the actual computation.

The forward logic of chunked linear attention can be summarized in the following pseudocode:

# Chunked linear attention forward (pseudocode)
state = init_state(batch, heads, d_k, d_v)
for chunk in split(x, chunk_size):         # Split along sequence dim to break long-sequence memory bottlenecks
    q, k, v = input_proj(chunk)
    intra = intra_interaction(q, k, v)     # Intra-chunk interaction term
    state = state_update(state, k, v)      # State recurrence (sequential dependency)
    out.append(intra + readout(q, state))  # State readout

More troublesome is that the hybrid architecture makes a single computation graph contain two types of subgraphs with very different characteristics: linear attention is compute-density-oriented, while sparse attention is memory-bandwidth-oriented. Execution scheduling must identify the resource preference of each Block type and orchestrate them differently, placing high demands on the graph optimizer.

3. Communication Side: Load Imbalance and Scheduling Overhead in Sparse Routing

Scaling MoE cannot avoid All-to-All communication: tokens need to be sent from the data-parallel domain to the device hosting the expert for computation, then returned along the same path. The routing results of 896 experts naturally follow a skewed distribution—hot experts accumulate tokens while cold experts have very small batches. This has two consequences:

1. The effective batch per expert is small, reducing the computational efficiency of Grouped GEMM, and compute utilization is diluted by communication and waiting;
2. Communication volume varies across devices, creating long-tail latency that slows the overall synchronization rhythm.

Combined with million-token long-context scenarios, the per-step token volume is enormous, and All-to-All message sizes swell accordingly, making communication overhead a significantly larger share of end-to-end time. Common industry countermeasures include: expert load balancing strategies (auxiliary loss or loss-free balancing algorithms), capacity factors and overflow control, micro-batch splitting to overlap communication and computation, and at the hardware layer, adopting high-bandwidth, low-latency supernode interconnect architectures with unified memory addressing—aggregating multiple AI chips into a single unified compute unit so that cross-card access is closer to “local large memory,” reducing scheduling penalties at the physical layer.

4. Memory Side: Storing Multi-TB Weights and Managing Intermediate Results

About 5.6 TB of BF16 weights is just the “entry ticket.” Full training also requires reserving space for activations, gradients, and optimizer states. Even with LoRA fine-tuning—freezing the backbone and training only injected low-rank bypasses, which can compress trainable parameters and optimizer states to a very small size—the 2.8T backbone weights still must be fully loaded into the cluster.

Therefore, memory management must answer two questions simultaneously: first, how to shard and deploy the weights (corresponding to the hybrid parallelism in Section 5); second, how to “store less and store smartly” for forward intermediate results (corresponding to the selective caching in Section 8). If these two things are not done well, even the strongest compute will be stuck at the memory wall.

5. Solution 1: Hybrid Architecture of Fully Sharded Data Parallel + Expert Parallel

For ultra-large MoE, a single parallelism strategy struggles to balance memory and efficiency. The mature combined approach is:

• Fully Sharded Data Parallel (FSDP): shard parameters, gradients, and optimizer states along the data-parallel dimension, aggregating them on demand during use, significantly reducing per-card memory usage;
• Expert Parallel (EP): distribute 896 experts across multiple device groups, with each expert storing only one complete set of parameters within its group, and tokens dispatched via All-to-All to the device hosting the expert for computation.

Combined, the dense parts go through FSDP and the expert parts go through EP, both keeping memory in check and preserving effective batch sizes for expert computation. A representative validation practice is: on a 128-card supernode cluster, using a “Fully Sharded Data Parallel + EP64” configuration (896 experts distributed across 64 expert parallel groups), completing LoRA fine-tuning training validation for a 2.8T-parameter model. The strategy can be illustrated as follows:

# Hybrid parallelism strategy schematic
parallel:
  data_parallel: 2        # Fully sharded data parallel dimension
  expert_parallel: 64      # 896 experts distributed across 64 EP groups
  devices: 128
expert:
  placement: distributed  # Expert parameters deployed in distributed fashion
dispatch:
  transport: all_to_all    # Token Dispatch / Combine via All-to-All
finetune:
  method: lora            # Backbone frozen, only low-rank bypasses trained

6. Solution 2: Large-Granularity Fused Operators for Hybrid Attention

For linear attention, the core idea is to design “large forward, large backward” fused operators: integrating input projection, intra-chunk interaction, state recurrence, state readout, and other stages into a single kernel, with the backward pass similarly fused as a whole to avoid intermediate states being repeatedly written to and read from memory. For sparse attention, index construction, Top-K filtering, scoring, and weighted aggregation are fused into a small number of large operators, with memory access patterns manually optimized for hardware characteristics. Engineering practice shows that such large-granularity fusion typically achieves more than 2x performance improvement over serial concatenation of small operators.

For emerging structural innovations (such as attention variants with deep residual connections), a more efficient path is to rely on a compilation stack of “frontend DSL + compiler auto-optimization”: developers describe forward logic in a way close to its mathematical definition, and the compilation framework handles automatic differentiation, operator generation, and performance tuning, compressing the cycle from paper to trainable implementation from weeks to days. The competitiveness of operator development is shifting from “hand-written labor” to “description + automatic generation.”

7. Solution 3: Full-Pipeline Optimization for MoE Training

The complete execution pipeline of MoE training is: Token Dispatch (Permute, rearranging tokens by target expert) → Expert Compute (Grouped GEMM grouped matrix multiplication) → Combine (Unpermute, restoring the original order). From a full-pipeline perspective, the optimization points include:

• Permute/Unpermute fusion: both are memory-intensive operations, and executing them separately is costly; deeply fusing them with Grouped GEMM, completing token aggregation and scattering inside the GEMM, can eliminate redundant data movement;
• High-performance All-to-All communication: message bucketing, multi-stream concurrency, topology-aware routing, reducing long-tail latency;
• Communication-computation overlap: at the micro-batch level, letting “sending the next batch” and “computing the current batch” proceed in parallel;
• Expert load balancing: mitigating hotspots through capacity factors, redundant experts, and other means, trading off between communication volume and compute balance.

Single-point optimization easily trades one problem for another; only by measuring and tuning Dispatch, Compute, and Combine as a whole can end-to-end gains be achieved.

8. Solution 4: Operator-Level Selective Caching and Memory Tiering

Backpropagation needs forward intermediate results: keeping all of them causes memory explosion, while recomputing all of them wastes compute. Selective caching is an engineering compromise between the two. One mechanism validated in practice is: applying two-level caching to the outputs of specified operators such as MatMul/Grouped GEMM—hot data resides in HBM, cold data sinks to larger-capacity, lower-bandwidth DDR, swapped in and out on demand.

Its value lies in trading cheaper capacity for scarce HBM space: what is gained is larger batch sizes and longer sequences, which is especially critical for million-token long-context training. When designing such mechanisms, two points need careful evaluation: whether cache hit rates and swap-in/swap-out bandwidth match the computation rhythm; and whether the impact of cached data on numerical precision is controllable.

9. Workflow Engineering: Enabling “Day-Level Onboarding” for New Models

For engineering teams, the biggest hidden challenge is: the speed of model architecture innovation far outpaces the speed of system adaptation. Manually adapting a new architecture takes weeks, clearly unable to keep up with the iteration rhythm. Feasible engineering paths include:

• Establishing an automated adaptation loop: model architecture parsing → automatic operator mapping → automatic parallel partitioning → performance profiling and localization → numerical alignment regression;
• Introducing Agent-driven R&D assistance workflows, delegating repetitive labor such as looking up interfaces, writing operator prototypes, and running alignment tests to automated toolchains, so engineers focus on high-value areas like fused operators and scheduling strategies;
• Establishing a baseline evaluation system based on four metrics: end-to-end throughput, Model FLOPs Utilization (MFU), latency, and peak memory, with any optimization passing the baseline before being merged.

Once this workflow is in place, a new model can go from weight release to trainable validation within a cycle measured in “days.”

10. General Path for Cloud Deployment

If building a self-managed cluster is too high a barrier, a cloud platform is a more realistic starting point. The general deployment combination is:

• Compute layer: GPU/heterogeneous compute clusters, paired with high-bandwidth, low-latency networking, meeting the physical link requirements of All-to-All communication;
• Platform layer: container platforms to orchestrate training jobs, object storage to host weights and datasets, and logging and monitoring alerts to ensure full-process observability;
• For example, leveraging Baidu AI Cloud’s general compute clusters, container platform, object storage, and logging/monitoring capabilities, a training and fine-tuning environment can be set up fairly quickly, letting the team focus on the model and algorithms themselves.

Practical recommendations for deployment: first run at small scale and complete numerical alignment, then gradually scale to the target size; run baseline evaluations for any change; establish fine-grained observability for communication and memory so bottlenecks are “visible.”

11. Summary

Adapting ultra-large MoE multimodal models for training is essentially fighting three walls: the “compute wall” of operator fragmentation and complex scheduling, the “communication wall” of uneven expert routing, and the “memory wall” of storing multi-TB weights. The four corresponding levers are: a hybrid architecture of Fully Sharded Data Parallel + Expert Parallel, large-granularity fused operators, full-pipeline Dispatch-Compute-Combine optimization, and selective caching with HBM/DDR tiering. Supplemented by automated, intelligent adaptation pipelines, this allows the system to keep pace with the speed of model architecture innovation. Mastering this methodology means that no matter which generation of new model arrives, engineering teams can compress the adaptation cycle from “months” to “days,” turning uncertainty into replicable engineering capability.