A empresa Moore Threads tornou o projeto Megatron-MUSA-Titan público, apresentando uma infraestrutura para treinamento de grandes modelos que ela desenvolveu para uso em suas próprias unidades de processamento gráfico. Este projeto é destinado ao trabalho com modelos do tipo 'mistura de especialistas' (mixture-of-experts) com um trilhão de parâmetros em clusters compostos por mais de 10.000 GPUs da Moore Threads, abrangendo tanto o pré-treinamento quanto o treinamento por reforço.
O código do projeto está disponível no GitHub. A empresa afirma que este sistema foi usado continuamente para pré-treinar MoE com um trilhão de parâmetros em clusters de 10.000 GPUs, onde sessões de treinamento prolongadas garantiram alta utilização de FLOPs do modelo e estabilidade do processo.
O repositório apresenta um exemplo de configuração para um modelo MoE no estilo Kimi-K2 com 1 trilhão de parâmetros. O princípio central do design foi manter a proximidade com soluções existentes. O Megatron-MUSA-Titan é construído sobre o framework Megatron-LM da NVIDIA. Graças à integração do PyTorch via torch_musa e à compatibilidade MUSA C com CUDA C, os desenvolvedores que possuem o ambiente MUSA podem clonar os repositórios oficiais Megatron-LM, slime e miles e iniciar o treinamento nas GPUs Moore Threads sem a necessidade de modificar o código-fonte.
O projeto indica que algoritmos gerais e correções de erros devem ser direcionados aos projetos superiores, enquanto a Moore Threads se concentra em dois tipos de alterações: otimizações em larga escala adaptadas ao hardware e clusters da Moore Threads, e soluções alternativas para problemas específicos do MUSA.
Os recursos de pré-treinamento incluem treinamento em bloco no formato FP8, que depende do MT-TransformerEngine da Moore Threads. A roteamento e agregação de tokens MoE são acelerados com o DeepEP, que opera sobre o Async Copy Engine da empresa, permitindo também a comunicação dos especialistas em paralelo com os cálculos. Outras opções incluem o descarregamento de ativações usando paralelismo de pipeline virtual e o agendamento de uma passagem direta e uma passagem reversa, além da comunicação paralela de pipeline combinada com cálculos, graças ao Mooncake RDMA.
A ferramenta de análise de rastreamento distribuído lê os rastreamentos PyTorch de cada GPU para detectar gargalos, como bolhas de pipeline e nós de baixo desempenho. O projeto afirma ser capaz de analisar uma execução de MoE com 1 trilhão de parâmetros em 10.000 GPUs em apenas algumas horas.
Para o treinamento por reforço, o stack suporta os frameworks slime e miles; o Megatron cuida do treinamento, e a compilação MUSA-SGLang gera as passagens, permitindo que um único sistema base atenda tanto ao pré-treinamento quanto ao RL. A otimização para inferência autônoma não faz parte do escopo deste projeto.
Para desenvolvedores que consideram aceleradores domésticos, a questão prática é qual parte do código de treinamento existente permanecerá funcional. A resposta da Moore Threads é que o código do ecossistema Megatron não requer modificações.
