Distributed Training Architecture Evolution: Full-Link Optimization for Ultra-Large-Scale MoE Models

Introduction: Systemic Engineering Challenges of Ultra-Large-Scale MoE Models

As foundation model parameter scales leap into the trillions, architectures represented by Mixture-of-Experts (MoE) have emerged as the mainstream choice for balancing model capability with training costs. These models typically encompass hundreds of billions to trillions of parameters, yet their forward computation activates only a subset of experts. Theoretically, this allows them to achieve greater expressive power with significantly lower computational overhead than fully dense models. However, this sparsity introduces entirely new engineering complexities—every component, from underlying operator implementation to distributed communication topologies and memory lifecycle management, requires redesign and validation.

Consider a typical 2.8-trillion-parameter native multimodal model with open weights. It employs a stable latent MoE architecture configured with 896 routed experts, activating only 16 experts per token during inference. Simultaneously, the model innovatively fuses linear attention and sparse attention mechanisms. Approximately three-quarters of the attention layers utilize linear attention to reduce the quadratic complexity associated with sequence length, while integrating Top-2048 sparse attention to support a context window of up to 1 million tokens and native image understanding capabilities. This combination of “multimodal + massive context + ultra-large-scale MoE” transforms model training and fine-tuning into a systemic engineering challenge rather than a mere accumulation of compute power.

This article analyzes the core pain points of distributed training for such models across three key dimensions: compute, communication, and memory. Drawing on mainstream technological paradigms, it proposes a full-link optimization approach covering model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management, offering a practical reference for R&D teams.

I. Deep Dive into Architectural Features and Training Bottlenecks

1.1 Computational Complexity from Heterogeneous Attention Mechanisms

In traditional Transformers, the computational complexity of dense self-attention grows quadratically with sequence length. To break this limitation, the model introduces a hybrid attention architecture. Some attention layers employ KDA linear attention, reducing complexity to linear through state recursion. Others retain DSA-based sparse attention, utilizing Top-2048 sparse indexing to select key tokens for computation. While this heterogeneous design enhances long-sequence processing, it presents significant challenges for underlying operator implementation.

Specifically, KDA linear attention relies on chunked state recursion, requiring the sequence to be divided into multiple blocks with state information passed between them. This imposes special requirements on operator data flow organization and kernel scheduling. Conversely, DSA sparse attention exhibits highly non-contiguous memory access patterns. It requires dynamically querying KV vectors based on sparse index tables, leading to poor memory access locality and low cache hit rates. More critically, when these two computation modes are interleaved within a single training step, traditional fusion strategies struggle to accommodate both: excessive fusion may restrict the flexibility of the sparse path, while insufficient fusion leads to massive kernel launch overheads and high memory read/write bandwidth consumption.

1.2 Communication and Load Imbalance Triggered by Sparse Routing

The core mechanism of MoE models lies in token-level routing—each token is dispatched to only a few expert networks. With a configuration of 896 experts, the routing distribution should theoretically be relatively uniform. However, in actual training, due to natural skew in data distribution and the dynamics of routing network learning, a “hot expert” phenomenon frequently occurs: some experts receive a massive influx of tokens, while others remain idle. This uneven token distribution directly leads to two consequences. First, the computation batch sizes for individual experts vary wildly, making it difficult to saturate GPU utilization. Second, cross-device communication volume surges, as tokens must be sent from their original device to the expert’s device, and the computation results must be returned.

This is especially true when employing Expert Parallelism (EP). The communication pattern shifts from conventional All-Reduce to massive All-to-All exchanges. In a thousand-GPU cluster, All-to-All communication easily becomes the system bottleneck, and its performance is heavily constrained by network topology bandwidth and latency characteristics.

1.3 Memory Pressure from Trillion-Level Parameters

Parameter scale is the most direct constraint on memory planning. For a 2.8-trillion-parameter model, storing just the BF16 weights requires approximately 5.6TB of memory. If optimizer states (such as the second-moment estimates in AdamW) are considered, even under mixed-precision training, an FP32 master weight and optimizer state usually need to be stored additionally, expanding the memory requirement several times over.

Clearly, the memory capacity of a single GPU or even a single node is far from sufficient to host such a massive parameter set. Distributed training strategies must be able to shard parameters, gradients, and optimizer states across multiple devices while ensuring data availability and synchronization efficiency during forward and backward computations.

II. Hybrid Parallelism Strategy: Collaborative Design of FSDP2 and EP

2.1 FSDP2: Evolution from Parameter Sharding to State Sharding

Fully Sharded Data Parallel (FSDP) is the mainstream solution for addressing the memory bottlenecks of large models. Traditional FSDP horizontally shards model parameters, gradients, and optimizer states simultaneously, with each device holding only a fraction of the complete model. FSDP2 (the second generation of FSDP, typically referring to an upgraded version supporting mixed sharding strategies) introduces finer-grained control logic on top of this, allowing different sharding dimensions and communication timings to be applied to different types of states.

In the scenario of a 2.8-trillion-parameter MoE model, FSDP2 is particularly valuable. We can treat non-expert parameters (such as attention layers, embedding layers, and routing networks) as the dense part, managed via FSDP2’s preset sharding methods. Meanwhile, the 896 expert parameters are treated as the sparse part and handed over to Expert Parallelism. The key value of FSDP2 lies in providing a unified state management interface for both, avoiding state redundancy and maintenance complexity when switching between multiple parallel strategies (like Tensor Parallelism and Pipeline Parallelism).

2.2 Expert Parallelism (EP): Trading Communication for Memory and Compute

The core idea of Expert Parallelism is to assign different groups of experts to different devices. Each device stores and computes only a subset of experts. During forward propagation, tokens are dispatched to the corresponding devices based on routing results, and returned after computation. The primary advantage of EP is that it does not require replicating all expert parameters, thus significantly saving memory. At the same time, the computation for each expert is dense and independent, making it suitable for improving computational efficiency through large batch parallelism.

However, as a “communication-intensive” strategy, EP’s performance heavily depends on the efficiency of communication primitives. In a typical configuration of 896 experts across 128 GPUs, EP64 means each device hosts 14 experts, and tokens must be massively exchanged across 64 device groups. Here, the implementation of All-to-All communication (such as whether hierarchical switching is used, or whether a hybrid network of NVLink and InfiniBand is utilized) directly dictates the iteration time of the training.

2.3 Collaborative Practice of FSDP2 + EP64

The FSDP2 + EP64 hybrid parallel architecture provided an excellent demonstration when validating LoRA fine-tuning training on a 128-GPU Atlas 950 SuperPod. LoRA (Low-Rank Adaptation) inherently trains only a small fraction of additional parameters while freezing the vast majority of original weights. This requires the distributed framework to precisely distinguish between “gradients needing synchronization” and “frozen weights only participating in forward passes”—only initiating optimizer synchronization for LoRA parameters and maintaining them in full precision. This is exactly what FSDP2’s state-based fine-grained management achieves.

In this solution, non-expert layer weights and optimizer states are sharded via FSDP2, while expert layers are distributed via EP64. During forward and backward passes, each device only processes its local subset of expert tokens, and the output is aggregated via All-to-All communication. Validation results showed that this architecture maintains stable and controllable memory usage in long-sequence scenarios. Because the LoRA parameter scale is extremely small, overall communication volume is significantly reduced compared to full fine-tuning.

III. Compute-Side Optimization: Operator Fusion and Scheduling Strategies

3.1 Execution Optimization for KDA Linear Attention

The state recursion and chunked computation of KDA linear attention naturally align with a “segmented scan” strategy. At the operator fusion level, operations such as Query-Key products within a chunk, state updates, and Value weighting can be fused into a custom kernel, eliminating memory round-trips for intermediate tensors. Furthermore, state transfer between chunks can be implemented using Persistent Kernels. By retaining the state in shared memory or registers at the end of each block and passing it directly to the next block, memory read/write overhead is eliminated.

However, it is crucial to note that this persistent scheduling requires strict serial execution of computation blocks—subsequent blocks must wait for the state output of preceding blocks. In multi-GPU scenarios, if a KDA layer is distributed across devices, this serial dependency severely limits parallelism. Therefore, when partitioning model parallelism, efforts should be made to ensure that all sequence chunks of a KDA layer reside within the same device, avoiding cross-device state recursion.

3.2 Memory Access Reordering for DSA Sparse Attention

The non-contiguous memory access issue faced by DSA sparse attention can be mitigated via a “pre-fetch + reorder” strategy. In specific implementation, the KV addresses pointed to by sparse indices can be collected first, and the scattered KV vectors can be batch-prefetched into shared memory or registers before performing dense computation. Concurrently, tokens can be reordered based on index distribution patterns, aggregating tokens with similar sparse patterns within a batch to improve cache line utilization.

Additionally, the sparse indices themselves consume memory bandwidth. A common optimization is to use bitmaps instead of index tables, paired with compressed storage formats to reduce the transmission volume of index data. This is particularly effective when the Top-2048 sparsity is high.

3.3 Orchestration Strategy for Heterogeneous Layer Scheduling

Because KDA and DSA differ significantly in data flow patterns, forcibly fusing them into the same kernel often does more harm than good. A more reasonable approach is to orchestrate at the graph compilation level: marking KDA-type layers as “sequence-dependent” and DSA-type layers as “sparse memory-access.” In the training scheduler, sparse layers can be prioritized to fill the pipeline bubbles caused by sequence-dependent layers, thereby achieving interleaved utilization of computational resources.

IV. Engineering Details of Communication and Memory Management

4.1 Hierarchical Routing for Uneven Communication

To address the uneven token distribution caused by 896 experts, a “hierarchical routing” mechanism can be added on top of global routing. The first layer coarsely assigns tokens to EP groups, while the second layer allows the routing network within the group to finely select specific experts. This approach effectively reduces cross-group communication volume—the vast majority of tokens only need to be exchanged within the group, triggering cross-group communication only when intra-group load is saturated. Coupled with a dynamic load balancing strategy, where expert load statistics are collected periodically and routing biases are adjusted accordingly, the hot expert effect can be further smoothed out.

4.2 Three-Tier Memory Planning: Weights, Activations, and Communication Buffers

Memory management requires comprehensive planning before the training process begins. It can generally be divided into three tiers:

• Static Memory: The largest footprint consists of weights and optimizer states, which are fixed on respective devices via FSDP2’s sharding mechanism and do not participate in dynamic allocation.
• Dynamic Memory: Activations during forward and backward passes. In ultra-long sequence scenarios, activations can dominate memory. Activation Checkpointing should be prioritized to trade time for space, with CPU Offloading as an emergency contingency strategy.
• Communication Buffers: All-to-All communication requires temporary buffers to store token tensors pending transmission. This memory footprint is often overlooked, but insufficient buffering can cause serial waiting between communication and computation. It is recommended to reserve an independent memory pool for the communication engine and employ a ring buffer design to support asynchronous communication.

4.3 Coordination of Asynchronous Pipelines and Gradient Accumulation

At a scale of a thousand GPUs, communication latency is now on the same order of magnitude as computation latency. To hide communication overhead, backward gradient computation can be split into an asynchronous pipeline triggered layer-by-layer, allowing the gradient computation of the Nth layer to overlap with the All-Reduce/All-to-All communication of the (N-1)th layer. Gradient Accumulation can also coordinate with this strategy: advancing the pipeline in Micro-batches and executing a complete weight synchronization only after accumulating several micro-batches, thereby trading fewer communication instances for higher compute density.

V. Automated Adaptation: From Manual Tuning to Agent-Driven Processes

Faced with increasingly complex combinations of parallel strategies, relying solely on senior engineers to manually adjust communication and computation mappings can no longer meet the demands of rapid iteration. The industry is currently exploring Agent-driven automated development pipelines. Processes such as model structure parsing, operator mapping, parallel strategy searching, and performance bottleneck profiling are abstracted into a series of programmable tasks executed by automated intelligent agents.

A typical path for this pipeline is as follows: First, a model structure parser identifies the distribution of MoE and attention layers. Then, based on hardware topology (device interconnect bandwidth matrix), it automatically searches for the optimal expert splitting dimensions and communication groupings. Next, combined with a performance model, it estimates the latency of each strategy and automatically generates kernel fusion schemes. Finally, small-scale trial runs are conducted on the actual cluster to correct the performance model with empirical data, forming a closed-loop optimization.

The superiority of this paradigm lies in compressing the onboarding cycle for new models from days to hours, significantly reducing reliance on manual experience. This is particularly suited for ultra-large-scale models like the 2.8-trillion-parameter architecture, where structural changes are frequent.

VI. Practical Validation and Key Takeaways

The effectiveness of the aforementioned solutions was fully validated during LoRA fine-tuning training on a 128-GPU Atlas 950 SuperPod. The high bandwidth, low latency, and unified memory addressing capabilities provided by the SuperPod architecture offer hardware-level support for FSDP2’s shard synchronization and EP’s All-to-All communication.

The following key takeaways can be extracted:

1. No One-Size-Fits-All Parallel Strategy: For hybrid architecture models mixing dense and sparse modules, the combination of FSDP2 and EP is vastly superior to either strategy used independently.
2. Compute Optimization Must Embed Architectural Details: The scheduling strategies for linear and sparse attention should be customized. Blindly applying fusion templates from dense Transformers is counterproductive.
3. The Core of Communication Optimization is Topology Awareness: Prioritizing the high-bandwidth domain (e.g., NVLink domain) within a SuperPod and only utilizing higher-level networks for cross-domain traffic effectively reduces All-to-All latency.
4. Memory Planning Must Intervene Early: Before writing training scripts, capacity budgeting and reservation strategies must be planned across the three tiers of weights, activations, and communication buffers.
5. Automated Tools are Inevitable for Scaling: For ultra-large-scale MoE models, manual tuning has hit an efficiency ceiling. Agent-driven automated adaptation pipelines will be the core direction of future evolution.

Conclusion

Distributed training for 2.8-trillion-level MoE multimodal models is a multi-dimensional systemic engineering challenge involving parallel strategies, operator implementation, network communication, and memory engineering. The solutions outlined in this article—a hybrid FSDP2 + EP parallelism, customized operator optimization for heterogeneous attention, hierarchical routing and asynchronous communication pipelines, and an Agent-driven automated adaptation pipeline—provide a valuable engineering blueprint for this domain.

It is worth noting that the support of advanced hardware is equally critical. The high bandwidth and unified addressing capabilities of SuperPods provide an ideal foundation for training optimization. As model scales continue to expand, innovations in distributed training technology will remain the key engine driving the continuous expansion of large model capabilities.