Training and Optimization Practices for Adapting Trillion-Parameter Sparse MoE Models on Ascend Compute Clusters

1. Background: Training Challenges Posed by Trillion-Parameter Sparse Models

In recent years, large language models and multimodal models have continued to scale toward higher parameter counts. Particularly under sparse Mixture-of-Experts (MoE) architectures, total model parameters can reach trillions, while the number of parameters actually activated per token accounts for only a small fraction. This design increases model capacity while significantly raising the complexity of the training system.

Consider a typical sparse expert model with approximately 2.8 trillion total parameters. It employs a stable latent expert routing mechanism with hundreds of routed experts, and each token activates only a small subset of them. The model also combines linear attention with sparse attention: roughly three-quarters of the attention layers use linear attention, while the remaining layers incorporate Top-k sparse attention, supporting million-token context and native image understanding. Such models demonstrate strong capabilities in long-horizon autonomous programming, complex agent tasks, 3D content generation, financial research, and deep knowledge work, but they place higher demands on the underlying training infrastructure.

From a training perspective, the challenges come primarily from three directions. First, the model weights are enormous—at BF16 precision alone, the weights require approximately 5.6 TB of device memory, which a single card or small-scale node cannot accommodate. Second, sparse expert routing introduces imbalance in both communication and computation; token distribution across experts fluctuates considerably, easily forming small per-expert batches that reduce computational efficiency. Third, the hybrid attention structure diversifies operator types—linear attention’s state recurrence and sparse attention’s non-contiguous memory access are interleaved, increasing the difficulty of operator fusion and execution scheduling.

2. Overall Design Approach for the Training Platform

To address these challenges, the training platform needs to provide higher inter-node bandwidth, lower communication latency, and unified memory addressing at the architectural level. Ideally, the cluster should have high-throughput intra-node and inter-node interconnects, allowing parameters to be distributed across multiple accelerator nodes while maintaining high data exchange efficiency.

At the software level, the training system must build full-stack capabilities covering model adaptation, hybrid parallelism, operator acceleration, communication optimization, and memory management. The model adaptation stage requires rapidly parsing model structure, parameter naming, and tensor shapes, and mapping them into representations recognizable by the local parallel framework. The hybrid parallelism stage needs to determine the combination of data parallelism, expert parallelism, and tensor parallelism based on model scale, number of experts, and cluster topology. The operator acceleration stage must complete fusion and high-performance implementations for key modules such as attention, expert computation, and normalization. The communication optimization stage needs to reduce all-to-all exchange overhead during expert routing and overlap communication with computation as much as possible. The memory management stage must lower peak memory usage through recomputation, parameter offloading, and tiered caching.

In practical engineering, an automated development workflow can also be introduced—using an agent-driven approach to complete new model onboarding, structure parsing, operator mapping, and benchmarking—thereby shortening the adaptation cycle for new models to days.

3. Parallelism Strategy: FSDP and Expert Parallelism Working Together

For trillion-parameter sparse models, a single parallelism strategy cannot simultaneously satisfy memory and computational efficiency requirements. In practice, a combination of FSDP and expert parallelism is commonly adopted. FSDP shards the dense (non-expert) parameters across different accelerator cards, completing updates through parameter gathering in the forward pass and gradient synchronization in the backward pass. Expert parallelism distributes different experts across different devices, with each device storing only a subset of expert parameters; after token routing, data is sent only to the device hosting the target expert.

Taking a 128-card cluster as an example, a configuration of FSDP with an expert parallelism degree of 64 can be used to complete LoRA fine-tuning validation. The LoRA approach only requires training low-rank adapter matrices, substantially reducing optimizer state and gradient storage requirements, making it more suitable for rapidly validating ultra-large models on a limited number of nodes.

In expert parallelism, a single token activating only a few experts means that the number of tokens each expert receives is stochastic. To improve computational granularity, tokens can be reordered and grouped after routing, aggregating tokens sent to the same expert into contiguous batches for subsequent batched matrix multiplication. For extremely uneven load scenarios, expert capacity limits and auxiliary balancing strategies can also be introduced to prevent individual experts from becoming training bottlenecks.

4. Attention Operator Fusion and Implementation

The hybrid attention architecture causes both forward and backward computation to involve multiple operator types, such as normalization, matrix multiplication, Softmax, state recurrence, and sparse indexing. If small operators are executed segment by segment, large amounts of intermediate results are repeatedly written to and read from high-bandwidth memory, reducing effective computational density.

One effective approach is to split the attention forward computation into a large fused forward operator, integrating multiple computation stages into a single kernel to reduce the number of memory accesses for intermediate tensors. The same applies to backpropagation—using a large fused backward operator that merges multiple steps in gradient computation. For the sparse attention portion, a dedicated fused operator can handle index construction, attention score computation, and weighted aggregation. Compared to the original small-operator chain, such fused implementations can achieve more than double the performance.

For relatively novel attention residual structures, compilation frameworks and domain-specific compilers can be leveraged to rapidly develop forward and backward operators. The compilation framework maps high-level tensor expressions to low-level parallel computation graphs, and the compiler then performs loop unrolling, instruction scheduling, and memory optimization based on target hardware characteristics. This approach can significantly shorten the onboarding cycle for new structures while maintaining good portability.

5. Sparse Expert Computation and Communication Optimization

The forward path of a sparse expert model typically includes three stages: token dispatch, expert computation, and result aggregation. Token dispatch sends tokens to the corresponding experts based on routing results; expert computation performs the linear transformation or gating computation for each expert; result aggregation recombines expert outputs in the original token order.

For this pipeline, three categories of operators need focused optimization. The first is Permute and Unpermute, used for token reordering before and after routing. The second is grouped matrix multiplication, used to improve small-batch computation efficiency across different experts. The third is all-to-all communication, used to transfer tokens and expert results across nodes. Optimization strategies include packaging Permute with subsequent transfers, merging small-batch matrix multiplications from multiple experts into grouped matrix multiplication, and overlapping communication with forward computation.

To further alleviate load imbalance, the number of tokens received by each expert can be tracked, and balancing strategies dynamically adjusted during training. When a single expert’s computation batch becomes too small, grouped matrix multiplication efficiency drops noticeably; therefore, a minimum batch threshold can be set, and experts below the threshold can use batch merging or deferred processing.

6. Memory Management: Tiered Caching and Computational Optimization

In ultra-large model training, memory pressure comes not only from the weights themselves but also from activations, optimizer states, and temporary buffers. To reduce memory usage, a Swap Core operator-level selective caching mechanism can be employed, applying secondary caching to computation results of specified matrix multiplications or grouped matrix multiplications, forming tiered storage between HBM and DDR. By temporarily storing some intermediate results in slower but larger-capacity memory, high-bandwidth memory can be freed at the cost of moderately increased transfer overhead. Under the tested configuration, this mechanism delivers an average training performance improvement of approximately 15%.

For loss computation, large vocabularies introduce significant memory usage. The traditional approach requires computing full log probabilities over the large vocabulary matrix, consuming substantial temporary memory. Two optimization methods can be adopted: sequence-chunked loss and vocabulary-chunked loss. Sequence-chunked loss splits the input along the sequence dimension into multiple chunks, computing cross-entropy chunk by chunk and accumulating gradients; vocabulary-chunked loss splits along the vocabulary dimension, loading only a portion of vocabulary vectors at a time to complete computation. The two methods can be combined to significantly reduce peak memory usage without compromising numerical stability.

7. Practical Deployment Recommendations

For teams looking to train or fine-tune sparse expert models on large-scale clusters, the following steps can serve as a reference:

1. Structure Parsing: First complete automatic parsing of model weights, tensor shapes, and network structure. Confirm which layers belong to the expert portion and which belong to the dense portion, and which attention layers use linear attention versus sparse attention.
2. Parallelism Selection: Based on model parameter count, available cards, and network topology, initially determine the FSDP sharding granularity and expert parallelism degree. It is recommended to first run single-step forward and backward validation on a small-scale node setup to confirm memory feasibility before scaling to full training.
3. Operator Adaptation: Prioritize implementing key operators such as attention, grouped matrix multiplication, and token dispatch/aggregation. Functional correctness verification can be completed first, then gradually replaced with high-performance fused versions.
4. Communication Overlap: Add communication-computation overlap logic to routing and data synchronization paths, hiding all-to-all operations within expert computation or parameter updates as much as possible.
5. Memory and Loss Optimization: Combine tiered caching, sequence chunking, and vocabulary-chunked loss to control peak memory. Reserve memory headroom in training configurations to handle temporary peaks from long texts and complex inputs.
6. Monitoring and Tuning: Continuously observe expert load distribution, communication time proportion, operator throughput, and memory levels during training. Target bottleneck stages for optimization rather than blindly stacking parallelism degrees.

Through these methods, large-scale sparse expert models can be adapted and LoRA fine-tuned under existing cluster conditions, providing an engineering foundation for subsequent full-parameter training and more complex tasks. In the future, as more efficient communication topologies, finer-grained expert scheduling, and more intelligent compilation optimization mature, training efficiency for trillion-parameter models will continue to improve.