Scaling MoE Model Adaptation on Domestic AI Compute: From 2.8T Parameters to 128-Card LoRA Fine-Tuning

Introduction: The Engineering Problem of Ultra-Large Open-Weight Models

A recent open-source release by an AI research team introduced a new flagship natively multimodal model with 2.8 trillion total parameters. The model employs a sparse Mixture-of-Experts (MoE) architecture, supports a context window of up to 1 million tokens, and includes native image understanding capabilities. Around the same time, a domestic AI compute platform completed adaptation of the model and validated LoRA fine-tuning on its supernode hardware.

The significance of this event goes beyond “another model running successfully.” It highlights several core tensions in large-scale model engineering: parameter size versus memory capacity, sparse routing versus communication efficiency, and heterogeneous attention mechanisms versus kernel fusion. This article analyzes the model architecture, breaks down these challenges, and describes the corresponding engineering solutions.

Model Architecture Analysis: Why Adaptation Is Harder Than with Conventional Dense Models

Sparse Expert Architecture Fundamentals

The model uses a Stable LatentMoE sparse expert architecture with 896 routed experts, activating 16 experts per token. This means:

• Total parameters: 2.8 trillion
• Activated parameters per token: approximately 1/56 of the total (16/896)
• Expert distribution: all 896 experts must be distributed across multiple accelerators or nodes

The design intent of sparse activation is to expand model capacity while controlling compute cost. However, this imposes a fundamental system requirement: expert parameters must be sharded across devices, and tokens must be routed to the devices hosting their selected experts. This introduces All-to-All communication—one of the most difficult communication patterns to optimize in distributed training.

Operator Complexity of Hybrid Attention

The model employs a hybrid of KDA (linear attention) and Gated DSA (sparse attention):

• Approximately 3/4 of attention blocks use KDA linear attention: the main advantage is reducing sequence-length complexity from quadratic to linear, but the state recurrence and chunked computation patterns are not straightforward to implement efficiently on GPU/NPU hardware.
• Top-2048 sparse attention mechanism: the DSA portion requires sparse indexing and non-contiguous memory access, which differs fundamentally from the access patterns of dense attention.
• 1 million token context support: KV cache management and memory footprint become serious concerns in long-sequence scenarios.

Engineering implication: The interleaving of KDA state recurrence and chunked computation with DSA sparse indexing and non-contiguous memory access makes kernel fusion and execution scheduling exponentially more complex. Implementing this with the traditional “one small operator per kernel” approach would cause kernel launch overhead and intermediate result read/write costs to dominate performance.

Additional Burden from Multimodal and Reinforcement Learning

The model has native image understanding and is post-trained with large-scale joint reinforcement learning. This means the training pipeline must handle image encoding and cross-modal alignment in addition to text tokens. Although this article focuses on the LoRA fine-tuning stage, the adaptation layer still needs to account for multimodal input preprocessing.

Three Core Challenges: Compute, Communication, and Memory

Compute: Kernel Fusion and Execution Scheduling

The KDA/DSA hybrid architecture imposes strict requirements on kernel fusion:

• KDA portion: state recurrence creates sequential dependencies where the current output depends on the previous hidden state. Chunked computation enables parallelism, but state must still be passed between chunks. Splitting this into many small operators forces each operator to read and write intermediate states, making memory bandwidth the bottleneck.
• DSA portion: Top-2048 sparse attention requires computing indices first, then gathering the corresponding KV pairs, and finally executing attention computation. This involves non-contiguous memory access, which is cache-unfriendly.

Solution direction: Design large forward and backward fused kernels that merge multiple computational stages into a single kernel. For example, fuse KDA chunked computation, state updates, and output projection into one kernel to reduce intermediate memory round-trips. For DSA, use fused operators such as LI/SFA/SFAG to combine sparse indexing, gather, and attention computation.

In actual validation, fused operators delivered over 2x performance improvement compared with the original small-operator implementation. This reflects a key reality: in ultra-large-scale model training, kernel launch overhead and memory bandwidth often become bottlenecks before raw FLOPs do.

Communication: Load Imbalance in Sparse Routing

Sparse routing across 896 experts typically produces two problems:

1. Uneven token distribution among experts: popular experts are selected by many tokens while less popular ones remain nearly idle. This overloads some devices and leaves others underutilized, reducing overall throughput.
2. Small per-expert batch size: the number of tokens routed to an individual expert may be very small, making the M dimension of matrix multiplications small and preventing full utilization of Tensor Core / matrix compute units.

Solution direction:

• Load-balancing auxiliary loss or capacity factor constraints can be introduced at the routing layer, but these are model training strategies. At the adaptation layer, the focus is on communication optimization.
• Deploy experts using Expert Parallelism (EP) combined with optimized All-to-All communication. The key is grouping tokens by destination expert and sending them in batches to reduce the number of communication rounds, while leveraging the large bandwidth and low latency of the supernode to reduce All-to-All overhead.
• For small per-expert batch sizes, token reordering can place tokens destined for the same expert contiguously, increasing the batch dimension of matrix multiplication.

Memory: Storing and Aggregating 2.8T Parameters

2.8 trillion parameters require approximately 5.6 TB in BF16 weights alone. This excludes optimizer states, gradients, and activations. Including full training state would multiply memory requirements severalfold.

Core conflict: single-accelerator memory is far from sufficient, so parameters must be sharded across multiple accelerators. However, sharding requires aggregating expert parameters during forward propagation and aggregating gradients during backpropagation, which again brings us back to communication.

Solution direction: use a hybrid of FSDP2 (Fully Sharded Data Parallel 2) and EP:

• FSDP2 shards non-expert parameters (attention layers, embedding layers) across all accelerators, performing All-Gather on demand during forward/backward passes.
• EP distributes the 896 experts across different accelerators, with each accelerator storing only a subset of expert parameters.
• Together, these enable distributed deployment and efficient aggregation of expert parameters.

Engineering Practice: 128-Card LoRA Fine-Tuning Validation

Parallel Strategy Configuration

In actual validation, a hybrid FSDP2 + EP64 parallel configuration was used to complete LoRA fine-tuning on a 128-accelerator supernode. EP64 means the experts are divided into 64 groups, each deployed on different devices.

Why LoRA fine-tuning instead of full fine-tuning?

• Full fine-tuning requires storing complete optimizer states, which demands extremely high memory.
• LoRA trains only low-rank adaptation matrices, significantly reducing the number of trainable parameters and easing memory pressure.
• For adaptation validation, LoRA fine-tuning is sufficient to verify correctness and performance of the forward/backward paths.

Full-Pipeline Training Capability

To address the training challenges of ultra-large MoE models, the domestic AI compute platform built a full-pipeline capability leveraging the supernode’s large bandwidth, low latency, and unified memory addressing:

Stage	Key Work
Model adaptation	Mapping the model structure to the framework, supporting KDA/DSA hybrid attention
Hybrid parallelism	FSDP2+EP combination, distributed deployment of expert parameters
Operator acceleration	Replacing small operators with fused kernels, achieving over 2x performance improvement
Communication optimization	Grouped batch sending for All-to-All, reducing communication overhead
Memory management	Parameter sharding, gradient sharding, activation recomputation

An agent-driven automated development pipeline enables day-level model onboarding and rapid iteration. This means the time window from model open-source release to completed adaptation validation is compressed to days.

Key Optimization Highlights

KDA linear attention optimization:

• Designed large forward/backward fused kernels
• Integrated chunked computation, state recurrence, and output projection into a single kernel
• Reduced intermediate result memory reads/writes

DSA sparse attention optimization:

• Used fused operators such as LI/SFA/SFAG
• Combined sparse indexing, gather, and attention computation
• Achieved over 2x performance improvement compared with small-operator implementations

Expert parallel communication optimization:

• EP64 grouping, 896 experts distributed across devices
• Tokens grouped by destination expert and sent in batches
• Leveraged supernode bandwidth to reduce All-to-All overhead

Lessons Learned and Reusable Engineering Principles

Several principles of general relevance to ultra-large-scale model training emerge from this adaptation practice:

Principle 1: Before the memory wall, sharding is the only way forward. Whether FSDP2 or EP, the essence is splitting parameters across multiple accelerators. The difference lies in the sharding dimension and aggregation method. FSDP2 shards along the parameter dimension; EP shards along the expert dimension. The two are complementary.

Principle 2: Communication optimization must be co-designed with compute optimization. All-to-All communication from sparse routing cannot be avoided, but communication latency can be hidden through token grouping, batch sending, and overlap with computation. The supernode’s high-bandwidth, low-latency architecture provides the hardware foundation for this.

Principle 3: Kernel fusion is a critical means of breaking through performance bottlenecks. In ultra-large-scale models, kernel launch overhead and memory bandwidth often become bottlenecks before FLOPs. Fused kernels reduce the number of kernels and intermediate result read/writes, delivering severalfold performance gains.

Principle 4: Automated adaptation pipelines determine iteration speed. Model architectures evolve rapidly, and manual adaptation costs are increasingly prohibitive. Agent-driven automated development pipelines can compress adaptation cycles from days to even shorter, which is essential for keeping pace with model open-source cadence.

Conclusion

Adapting and fine-tuning ultra-large-scale MoE models is a systematic engineering effort spanning compute, communication, and memory dimensions. This article analyzed the model architecture, broke down the specific challenges posed by KDA/DSA hybrid attention, 896-expert sparse routing, and the 2.8T parameter memory footprint, and described concrete engineering solutions including fused kernel design, FSDP2+EP hybrid parallelism, and All-to-All communication optimization.

For developers facing or about to face similar problems, the core approach can be summarized in one sentence: solve memory problems with sharding, solve compute problems with fusion, and solve communication problems with grouping and batching. Only when these three work together can usable performance be achieved in ultra-large-scale model training.