Moore Threads releases Megatron-MUSA-Titan for training trillion-parameter MoE models on proprietary GPUs
Read more
Pandaily
pandaily.com

Moore Threads releases Megatron-MUSA-Titan for training trillion-parameter MoE models on proprietary GPUs

Moore Threads has made the Megatron-MUSA-Titan project publicly available, presenting an infrastructure for training large models that it developed for use on its own graphics processors. This project is designed to work with trillion-parameter 'mixture-of-experts' (MoE) models on clusters consisting of over 10,000 Moore Threads GPUs, covering both pre-training and reinforcement learning.

The project code is available on GitHub. The company claims that this system was continuously used for pre-training a trillion-parameter MoE on 10,000-GPU clusters, where long training sessions ensured high FLOPs utilization and process stability.

The repository includes an example configuration for a Kimi-K2 style MoE model with 1 trillion parameters. A central design principle was maintaining closeness to existing solutions. Megatron-MUSA-Titan is built upon NVIDIA's Megatron-LM framework. Through integration of PyTorch via torch_musa and compatibility of MUSA C with CUDA C, developers with the MUSA environment can clone official Megatron-LM, slime, and miles repositories and run training on Moore Threads GPUs without modifying the source code.

The project indicates that general algorithms and bug fixes should be directed to upstream projects, while Moore Threads focuses on two types of changes: large-scale optimizations adapted for Moore Threads hardware and clusters, and workarounds for specific MUSA issues.

Pre-training features include block training in FP8 format, which relies on Moore Threads' MT-TransformerEngine. MoE token dispatching and merging are accelerated by DeepEP, running on the company's Async Copy Engine, which also allows expert communication parallel to computations. Other options include activation offloading using virtual pipeline parallelism and scheduling one forward and one backward pass, as well as pipeline parallel communication combined with computation, thanks to Mooncake RDMA.

A distributed tracing analysis tool reads PyTorch traces from each GPU to detect bottlenecks such as pipeline bubbles and low-performance nodes. The project states that it can analyze a 1-trillion-parameter MoE run on 10,000 GPUs in just a few hours.

For reinforcement learning, the stack supports the slime and miles frameworks; Megatron handles the training, and the MUSA build of SGLang generates runs, allowing one base system to serve both pre-training and RL. Optimization for autonomous inference is not within the scope of this project.

For developers considering domestic accelerators, the practical question is what part of the existing training code will remain functional. Moore Threads' answer is that the Megatron ecosystem code requires no changes.

Popular