Adapting and LoRA Fine-Tuning Ultra-Large-Scale MoE Flagship Models: Architecture, Bottlenecks, and End-to-End Optimization

1. Background: When Model Scale Enters the Multi-Trillion Range

As foundation model capabilities continue to evolve, native multimodality, extremely long context, and complex reasoning have become defining features of the latest flagship models. At the same time, model architectures are undergoing a significant shift: traditional dense models are being replaced by sparse Mixture-of-Experts (MoE) models. Total parameter counts are rising rapidly, but only a subset of parameters is activated per token during inference, striking a balance between capability and compute cost.

This article focuses on a representative ultra-large-scale sparse MoE model with the following characteristics:

• Native multimodal design
• 2.8 trillion total parameters
• 896 routed experts in the sparse MoE architecture
• 16 experts activated per token
• Hybrid attention: roughly three-quarters of attention modules use linear attention, while the remainder use sparse attention
• Top-2048 sparse attention mechanism
• Support for up to 1 million tokens of context
• Native image understanding

Models of this class show significant capability gains in long-horizon autonomous programming, complex agentic tasks, 3D content generation, financial research, and deep knowledge work. However, these improvements come with more complex computation graphs, extreme memory pressure, and highly sensitive communication bottlenecks for training and fine-tuning systems.

2. Engineering Challenges Posed by the Model Architecture

2.1 Hybrid Attention: Increased Complexity in Fused Kernel Development

In traditional Transformer models, attention structures are relatively uniform, allowing engineering teams to leverage mature operator libraries for rapid integration. In a hybrid attention architecture, however, different layers use different attention implementation paths:

• Some layers use linear attention, whose core involves state recurrence and block-wise computation
• Other layers use sparse attention, characterized by sparse indexing and non-contiguous memory access
• The model may also introduce deep residual structures that deviate from the classic Transformer forward and backward computation graph

This multi-operator configuration within a single model significantly reduces the value of simply reusing existing operator libraries. If implemented naively as a sequence of small operators, execution scheduling becomes fragmented, kernel launch overhead increases, and memory access efficiency degrades.

2.2 Sparse Expert Routing: Communication and Load Balancing Challenges

With 896 experts, input tokens must be dynamically routed to different experts during training. Ideally, each expert should receive a roughly balanced number of tokens. In practice, however, token traffic across experts often shows significant skew.

This creates two direct consequences:

• Some experts process very small batches, leaving GPU/NPU compute resources underutilized
• All-to-All communication becomes unbalanced, with certain links becoming bottlenecks

For models with hundreds of billions to trillions of parameters, experts are distributed across multiple accelerators and even multiple nodes. Communication overhead is no longer a negligible constant but a core variable that directly affects end-to-end training efficiency.

2.3 Memory Pressure: Weights Alone Exceed Single-Accelerator Capacity

For a 2.8-trillion-parameter model, BF16 weights alone require approximately 5.6 TB of memory. Adding gradients, optimizer states, and activations multiplies the memory requirement. A single accelerator clearly cannot hold the full model, so parameter sharding is mandatory.

Even with sharding, several issues remain:

• Trade-offs between sharding granularity and communication volume
• Balancing activation recomputation against memory usage
• Memory spikes in the output layer caused by a large vocabulary
• Coexistence of LoRA adapter parameters and frozen base model weights in large-scale fine-tuning

3. Overall Solution Design: Hybrid Parallelism, Fused Kernels, and Hierarchical Storage

Addressing these challenges requires coordinated design across four dimensions: parallelism strategy, kernel optimization, communication scheduling, and memory management. These cannot be optimized in isolation.

3.1 Hybrid Parallelism: Combining FSDP2 with Expert Parallelism

In dense models, Fully Sharded Data Parallelism (FSDP) is a common strategy. It shards model weights, gradients, and optimizer states across data-parallel ranks, reconstructing full parameters on demand during forward and backward passes. In sparse MoE models, however, FSDP alone cannot address the communication and load imbalance introduced by expert routing.

A typical solution is to introduce expert parallelism: different experts are placed on different devices, and each device is responsible for computing only its assigned experts. This means a single expert’s parameters no longer need to be fully replicated across all data-parallel ranks. Instead, they are concentrated on a smaller set of devices, reducing communication volume from parameter sharding and aggregation.

In practice, FSDP2 can be combined with expert parallelism. For example, on a 128-accelerator environment, training can be configured with an expert parallelism degree of 64. Key benefits include:

• Expert parameters are distributed along the expert-parallel dimension, reducing per-accelerator memory pressure
• Data parallelism and expert parallelism are decoupled, enabling independent tuning
• Routing, dispatching, and aggregation are completed within expert-parallel groups, reducing cross-group communication

In LoRA fine-tuning scenarios, base model parameters remain frozen, and only low-rank incremental parameters are trained. This further reduces optimizer state and gradient overhead. However, the forward pass still needs to load base parameters and perform sparse activation, so communication and memory management remain critical.

3.2 Compute-Side Optimization: From Fragmented Small Kernels to Fused Kernels

The forward pass of a hybrid attention model includes multiple stages, such as state recurrence, block-wise matrix multiplication, sparse indexing, RoPE embedding, and residual connections. Implementing each stage as a separate kernel leads to frequent memory reads and writes plus high launch overhead.

For linear attention, a large fused forward and backward kernel approach can be adopted: merging recurrence computation, state updates, and output projection into a single kernel. This enables cache reuse within the kernel, reduces HBM accesses, and lowers scheduler pressure.

For sparse attention, the core challenge is non-contiguous memory access. The following fusion strategies can be applied:

• Merge index generation, sparse attention score computation, and weighted aggregation
• Generate contiguous blocks during indexing to reduce memory access randomness
• Use custom operators to achieve performance several times higher than a naive small-kernel implementation

For novel structures such as deep residual connections, compilation frameworks and domain-specific compilers can be used to develop high-performance kernels and enable efficient integration of new architectural components.

3.3 Communication-Side Optimization: All-to-All and Load Balancing

Sparse MoE models require token dispatching, expert computation, and result aggregation in every forward pass. This process typically involves multiple All-to-All communication operations. Key optimization areas include:

• Reducing the number of communication operations by fusing compatible ones
• Overlapping communication with computation by starting the next batch of token transfers before expert computation finishes
• Balancing expert load through dynamic routing adjustments or auxiliary constraints to avoid long-tail experts
• Using custom communication kernels tuned for different expert parallelism degrees and message sizes

At high expert parallelism degrees, All-to-All message volume grows significantly. Optimizing individual communication latency alone yields limited gains, so the communication pattern itself must be restructured. For example, fusing Permute and Unpermute operations with communication kernels reduces data movement. Batch aggregation during token dispatch also avoids generating many small messages.

3.4 Memory-Side Optimization: Hierarchical Caching and Chunked Loss

Even with hybrid parallelism, large vocabularies and activations can still cause excessive memory spikes. Two strategies can address this.

The first is hierarchical caching for selected matrix multiplication or grouped matrix multiplication results, storing them in secondary storage such as a DDR tier alongside HBM. By offloading results not immediately needed to DDR, HBM peak pressure is reduced. In the tested configuration, this improved training throughput by approximately 15% on average.

The second targets chunked loss computation for large vocabularies. In large-vocabulary models, the output layer typically computes logits proportional to the vocabulary size, incurring significant memory overhead. Chunking can be applied along both sequence and vocabulary dimensions:

• Sequence-dimension chunking avoids generating the loss for an entire sequence at once
• Vocabulary-dimension chunking splits the large vocabulary into sub-blocks computed separately, then combines the results

This two-dimensional optimization significantly reduces output-layer memory peaks while maintaining numerical consistency.

4. Training Validation and Efficiency Gains

Using the above approach, end-to-end validation can be performed on a supernode environment. On a 128-accelerator supernode with FSDP2 and expert parallelism set to 64, LoRA fine-tuning of the ultra-large-scale model was successfully validated. The system demonstrated the following capabilities:

• Support for base training and fine-tuning of multi-trillion-parameter sparse MoE models
• Support for integrating and fusing operators in hybrid attention architectures
• Support for efficient All-to-All communication under high expert parallelism
• Support for stable training with hierarchical memory management and large vocabularies

From an engineering efficiency perspective, an automated driver-development workflow can compress the cycle for new model integration and iteration to a matter of days. For teams that need to quickly evaluate newly released models, this capability is highly valuable.

5. Reusable Lessons and General Practices

Although the practices above are based on a model of a specific scale, many of the insights apply broadly to similar ultra-large-scale sparse models:

1. Define the parallel topology before optimizing individual kernels. The parallelism strategy determines overall communication boundaries and memory distribution. Decide on the combination of data parallelism, expert parallelism, and sequence parallelism early to avoid optimizing on the wrong topology.

2. Identify fragmented stages first. Use profiling to find stages with dense kernel launches and poor memory access efficiency, then prioritize them for fusion. Attention, normalization, and loss computation are usually high-yield starting points.

3. Communication-computation overlap matters more than optimizing communication alone. In sparse MoE models, communication and computation naturally alternate. Use pipelining, asynchronous execution, and prefetching to increase overlap.

4. Memory optimization requires a layered design. Activation recomputation or parameter sharding alone often solves only part of the problem. Combining hierarchical storage, chunked loss, and selective caching achieves a better balance between training speed and memory usage.

5. Automation determines iteration speed. New model architectures change frequently. Manual integration is costly and error-prone. Building a configuration-driven or agent-driven automated development pipeline can significantly shorten the cycle from model release to training validation.

6. Conclusion

Ultra-large-scale sparse MoE models bring substantial capability gains but also impose higher demands on training infrastructure. This article has examined the core challenges of hybrid attention, sparse routing, communication bottlenecks, and extreme memory pressure, and outlined an end-to-end approach covering parallelism strategy, kernel fusion, communication optimization, and memory management. This approach has been validated for LoRA fine-tuning of a multi-trillion-parameter model on a several-hundred-accelerator supernode, with new-model integration cycles compressed to days.

For teams looking to build or improve ultra-large-scale model training capability, the key is to establish a closed-loop process: understanding model architecture, locating bottlenecks, designing parallelism, optimizing kernels, and validating the system. Model architectures will continue to evolve, but bottleneck-driven system optimization remains the essential method for adapting to rapid change.