Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUs

TL;DR

This work addresses the challenge of training extremely large sparse MoE language models on Ascend NPUs by combining a simulation-driven architecture search with specialized system optimizations. It introduces Pangu Ultra MoE, a near-trillion-parameter MoE model achieving a practical Model Flops Utilization (MFU) of $30.0\%$ on a $6{,}000$-node Ascend cluster, with throughput rivaling state-of-the-art systems. Key contributions include a two-level MoE design (architecture-search-informed) and a suite of system optimizations—hierarchical EP communication, adaptive pipe overlap, fine-grained recomputation, tensor swapping, and dynamic load balancing—that collectively reduce communication and memory bottlenecks. The results demonstrate that Ascend NPUs can harness full training stages for large-scale sparse LLMs, with extensive analyses offering guidelines for future MoE training and deployment.

Abstract

Sparse large language models (LLMs) with Mixture of Experts (MoE) and close to a trillion parameters are dominating the realm of most capable language models. However, the massive model scale poses significant challenges for the underlying software and hardware systems. In this paper, we aim to uncover a recipe to harness such scale on Ascend NPUs. The key goals are better usage of the computing resources under the dynamic sparse model structures and materializing the expected performance gain on the actual hardware. To select model configurations suitable for Ascend NPUs without repeatedly running the expensive experiments, we leverage simulation to compare the trade-off of various model hyperparameters. This study led to Pangu Ultra MoE, a sparse LLM with 718 billion parameters, and we conducted experiments on the model to verify the simulation results. On the system side, we dig into Expert Parallelism to optimize the communication between NPU devices to reduce the synchronization overhead. We also optimize the memory efficiency within the devices to further reduce the parameter and activation management overhead. In the end, we achieve an MFU of 30.0% when training Pangu Ultra MoE, with performance comparable to that of DeepSeek R1, on 6K Ascend NPUs, and demonstrate that the Ascend system is capable of harnessing all the training stages of the state-of-the-art language models. Extensive experiments indicate that our recipe can lead to efficient training of large-scale sparse language models with MoE. We also study the behaviors of such models for future reference.

Figures (14)

• Figure 1: Experts number ablation. Since increasing experts from 256 to 512 brings limited gains, we select 256 experts for Pangu Ultra MoE to balance performance and practical efficiency.
• Figure 2: Shared expert ablation. We compare the 20B baseline MoE (topk=8) and its no shared expert variant (topk=9). The former leads to slightly lower training loss.
• Figure 3: Simulation workflow for LLM performance and optimal search.
• Figure 4: Simulated model normalized throughput of training and inference.
• Figure 5: Comparison of the load balancing auxiliary loss and cross-entropy loss during training across different auxiliary loss strategies and potentially varying regularization strengths, i.e., $\alpha \in \{\text{1e-2}, \text{1e-4}\}$. Clearly, setting $\alpha$ to 1e-2 more effectively controls the expert load, with the auxiliary loss remaining close to its minimum value (around 1.0). However, the right plot shows that a stronger load constraint leads to a higher pretraining loss. On the other hand, when changing the scope of tokens used to calculate the expert selection frequency, it is noticeable that different levels of auxiliary loss exhibit varying optimization difficulties. Specifically, the Sequence level is the hardest to optimize, while the DP-Group level is the easiest.