Scalable Training of Mixture-of-Experts Models with Megatron Core

Este artigo apresenta o Megatron Core, uma solução de código aberto que otimiza o treinamento escalável de modelos Mixture-of-Experts (MoE) através de técnicas integradas de memória, comunicação e computação, alcançando alto desempenho em clusters de milhares de GPUs para modelos de bilhões a trilhões de parâmetros.

O essencial

Imagine que você está tentando construir o maior e mais inteligente cérebro digital do mundo. Para isso, você precisa de uma arquitetura chamada Mixture of Experts (MoE), ou "Mistura de Especialistas".

Pense em uma grande empresa de consultoria. Em vez de ter um único funcionário super-humano que sabe fazer tudo (o que seria caro e lento), você contrata milhares de especialistas: um é ótimo em matemática, outro em poesia, outro em programação, outro em culinária. Quando chega um pedido (um "token" de texto), um Gerente de Roteamento (o Router) olha para o pedido e decide: "Ok, este problema de matemática precisa do Especialista em Álgebra e do Especialista em Estatística". Apenas esses dois trabalham; os outros 998 especialistas ficam descansando.

Isso é incrivelmente eficiente para o computador (ele gasta menos energia), mas cria um pesadelo logístico para os engenheiros da NVIDIA. O relatório que você leu, escrito pela equipe da NVIDIA, explica como eles resolveram esse pesadelo usando o Megatron Core.

Aqui está a explicação simplificada, usando analogias do dia a dia:

1. O Problema: As "Três Paredes"

Quando você tenta treinar esses modelos gigantes, você esbarra em três barreiras físicas, como se estivesse tentando encher um balde com água enquanto o balde tem buracos:

• A Parede da Memória (O Armazém): Mesmo que apenas 2 especialistas trabalhem por vez, você precisa ter todos os 256 especialistas (e seus cadernos de anotações) guardados na memória do computador ao mesmo tempo. É como ter que alugar um armazém gigante para guardar 10.000 ferramentas, mesmo que você só use 2 por dia. O balde transborda.
• A Parede da Comunicação (O Trânsito): Como os especialistas estão espalhados em diferentes computadores (GPUs), o Gerente precisa enviar o pedido para o especialista certo. Se o especialista está em outro prédio, o pedido tem que viajar por cabos. Se houver muitos pedidos indo para lugares diferentes ao mesmo tempo, o trânsito fica congestionado. O computador fica parado esperando a mensagem chegar.
• A Parede da Eficiência de Cálculo (O Motor): Os especialistas são muito pequenos e rápidos. O computador (GPU) é como um caminhão de carga gigante. Se você só colocar uma caixa pequena (um cálculo de um especialista) nele, o caminhão viaja quase vazio. É um desperdício de combustível. Além disso, o "motorista" (o processador principal) fica cansado de ligar e desligar o caminhão para cada pequena caixa.

2. A Solução: O "Dobramento de Paralelismo" (Parallel Folding)

Antes, os engenheiros eram forçados a usar a mesma estratégia para todas as partes do cérebro. Era como tentar usar o mesmo mapa de trânsito para uma cidade densa e para uma estrada deserta.

A NVIDIA criou uma técnica genial chamada Parallel Folding (Dobramento de Paralelismo).

• A Analogia: Imagine que a parte de "Atenção" (que lê o texto) gosta de estar em um grupo pequeno e unido (como uma equipe de rugby). Já a parte dos "Especialistas" (MoE) prefere estar espalhada em muitos grupos diferentes para não se chocarem.
• A Inovação: O Megatron Core permite que você use um mapa de trânsito para a equipe de rugby e um mapa totalmente diferente para os especialistas. Eles não precisam mais ficar presos às mesmas regras. Isso permite que o sistema cresça para milhares de GPUs sem quebrar.

3. Como Eles Quebraram as Paredes

Quebrando a Parede da Memória

• Recomputação Inteligente: Em vez de guardar todas as anotações do especialista na memória (o que ocupa espaço), o sistema decide: "Vou jogar essas anotações fora e, se precisar delas depois, vou refazer o cálculo rápido". É como um cozinheiro que joga fora os ingredientes cortados e, se precisar, corta de novo, em vez de ter uma geladeira cheia de ingredientes prontos.
• Precisão Reduzida (FP8/FP4): Em vez de usar números com 16 casas decimais (muito precisos, mas pesados), eles usam números com 8 ou até 4 casas. É como escrever um livro usando apenas letras maiúsculas e sem acentos: o texto fica menor e cabe mais páginas no mesmo livro, e o computador lê mais rápido, sem perder a "essência" da história.
• Offloading (Mudança de Casa): Quando a memória da GPU acaba, eles movem temporariamente os dados "dormindo" para a memória do processador (CPU), como se guardasse caixas no porão enquanto trabalha na sala.

Quebrando a Parede da Comunicação

• Entregadores Super-Rápidos (DeepEP e HybridEP): Eles criaram novos protocolos de envio de dados que são como drones de entrega que voam direto para a janela do especialista, ignorando o trânsito normal.
• Esconder o Tempo de Espera: Enquanto um especialista está trabalhando em um pedido, o sistema já está enviando o próximo pedido para outro especialista. É como um maestro de orquestra: enquanto os violinos tocam, os trompetes já estão se preparando. O tempo de espera some.

Quebrando a Parede da Eficiência

• Agrupamento (Grouped GEMM): Em vez de pedir para o caminhão ir buscar uma caixa de cada vez, eles agrupam todas as caixas dos especialistas em um único carregamento gigante. O caminhão viaja cheio e eficiente.
• CUDA Graphs (O Roteiro Fixo): O computador gasta muita energia decidindo "o que fazer a seguir" a cada milissegundo. O Megatron cria um "roteiro fixo" (um gráfico) de todas as tarefas. É como ter um roteiro de filme pronto: o ator (GPU) não precisa pensar no que fazer, ele apenas executa a cena. Isso elimina o tempo de decisão do "motorista".

4. O Resultado: Treinando o Futuro

Com todas essas otimizações, a NVIDIA conseguiu treinar modelos gigantes (como o DeepSeek-V3 com 685 bilhões de parâmetros) de forma extremamente rápida e eficiente.

• No hardware mais novo (GB300/GB200): Eles atingiram velocidades que seriam impensáveis antes, processando trilhões de operações por segundo.
• Para o Longo Prazo: Eles também resolveram como treinar modelos que leem livros inteiros de uma vez (contexto longo), onde a parte de "leitura" (atenção) domina o trabalho, exigindo estratégias diferentes das usadas para os especialistas.

Resumo Final

Este relatório é o "manual de instruções" de como transformar um caos logístico (milhares de especialistas trabalhando em diferentes computadores) em uma máquina de precisão suíça.

A NVIDIA não apenas construiu o motor (o hardware), mas redesenhou todo o sistema de trânsito, armazéns e gestão de tempo (o software) para garantir que, quando você pedir para o computador "pensar", ele faça isso com a máxima velocidade, sem gastar energia à toa e sem estourar o orçamento de memória. É a engenharia de sistemas no seu nível mais sofisticado, explicada para que qualquer pessoa possa entender a lógica por trás da magia da Inteligência Artificial moderna.

Resumo técnico

Resumo Técnico: Treinamento Escalável de Modelos Mixture-of-Experts (MoE) com Megatron Core

1. O Problema

O treinamento de modelos de linguagem grandes (LLMs) baseados em Mixture-of-Experts (MoE) apresenta desafios de sistemas fundamentais que não existem em modelos densos tradicionais. A natureza esparsa do MoE, onde apenas um subconjunto de "especialistas" (experts) é ativado por token, cria uma desconexão entre parâmetros e computação: o número total de parâmetros cresce muito mais rápido do que a computação por token.

Essa assimetria gera três barreiras críticas ("Three Walls") que limitam a eficiência do treinamento em escala:

1. Parede de Memória (Memory Wall): Todos os parâmetros dos especialistas (mesmo os não ativados), gradientes e estados do otimizador devem residir na memória da GPU. Isso gera pressão de memória muito superior à de modelos densos equivalentes em custo computacional.
2. Parede de Comunicação (Communication Wall): A paralelização de especialistas (Expert Parallelism - EP) exige comunicações all-to-all para rotear tokens entre GPUs. Em grandes escalas, esse volume de dados satura a largura de banda, especialmente quando atravessa limites de nós (inter-node), consumindo até 60% do tempo de treinamento.
3. Parede de Eficiência de Computação (Compute Efficiency Wall): A arquitetura MoE fina (muitos especialistas pequenos) resulta em operações GEMM (multiplicação de matrizes) pequenas que não utilizam totalmente os núcleos Tensor Core das GPUs. Além disso, o roteamento dinâmico e a necessidade de sincronização host-device introduzem sobrecarga de CPU, deixando a GPU ociosa entre o lançamento de kernels.

Além disso, existe um desencontro denso-esparso: as camadas de Atenção (densas) e as camadas MoE (esparças) exigem configurações de paralelismo conflitantes (ex: alta Tensor Parallelism beneficia a atenção, mas fragmenta os especialistas), tornando difícil otimizar ambas simultaneamente em frameworks tradicionais.

2. Metodologia e Arquitetura

O Megatron-Core MoE é uma pilha de treinamento integrada projetada para abordar simultaneamente essas três barreiras através de co-design de hardware e software.

2.1. Paralelismo Multidimensional e "Parallel Folding"

Para resolver o desencontro denso-esparso, o framework introduz o Parallel Folding.

• Desacoplamento: Permite que as camadas de Atenção e as camadas MoE usem configurações de paralelismo independentes.
• Mecanismo: Enquanto as camadas de atenção podem usar alta Tensor Parallelism (TP) e Context Parallelism (CP), as camadas MoE podem usar alta Expert Parallelism (EP) com TP=1 (para manter a eficiência dos GEMMs).
• Benefício: Remove a restrição tradicional onde EP deve ser menor ou igual a DP (Data Parallelism), permitindo mapeamentos flexíveis que otimizam o uso de topologias de hardware específicas (ex: manter comunicações EP dentro do domínio NVLink de alta velocidade).

2.2. Quebrando a Parede de Memória

O framework emprega uma combinação de técnicas para reduzir a pegada de memória:

• Permutação Eficiente de Memória: Reorganiza matematicamente o cálculo para eliminar tensores intermediários redundantes sem custo computacional.
• Recomputação Granular: Recalcula apenas as ativações mais custosas e baratas de computar (ex: funções de ativação, LayerNorm) durante o backward pass, em vez de armazená-las.
• Offloading de Ativações: Transfere seletivamente ativações para a memória da CPU (PCIe) quando a memória da GPU está cheia, sobrepondo a transferência com computação.
• Precisão Reduzida (FP8/FP4): Armazena ativações e estados do otimizador em precisão reduzida (BF16 moments, FP8/FP4), reduzindo drasticamente o consumo de memória.
• FSDP para MoE: Implementa Fully Sharded Data Parallelism adaptado para MoE, shardando parâmetros e estados do otimizador dentro de grupos de especialistas.

2.3. Quebrando a Parede de Comunicação

Para minimizar o tempo gasto em comunicações all-to-all:

• Dispatchers Otimizados (DeepEP e HybridEP): Kernels especializados que eliminam etapas de permutação redundantes e utilizam primitivas de hardware (como TMA e IBGDA) para maximizar a largura de banda, especialmente em topologias NVLink.
• Sobreposição de Comunicação e Computação: Implementa esquemas de 1F1B (Forward-Backward) e DualPipe para esconder a latência da comunicação all-to-all atrás da computação de especialistas de micro-batches adjacentes.
• Latência Oculta: Em topologias NVL72 (GB200), o HybridEP explora a largura de banda massiva para resolver a comunicação sem necessidade de sobreposição complexa.

2.4. Quebrando a Parede de Eficiência de Computação

Para maximizar a utilização da GPU e reduzir a sobrecarga de CPU:

• Grouped GEMM: Agrupa múltiplas computações de especialistas em um único kernel ou lança-os em streams paralelos para melhorar a ocupação dos núcleos.
• Fusão de Kernels: Funde operações de roteamento, permutação e cálculo de perda auxiliar em kernels únicos, reduzindo o número de lançamentos de kernel.
• CUDA Graphs: Captura a sequência de kernels para eliminar a sobrecarga de lançamento por iteração. Para MoE sem descarte de tokens (dropless), o framework introduz:
  • Kernels Iniciados pelo Dispositivo: Os kernels leem as formas dinâmicas diretamente da memória da GPU, eliminando a necessidade de sincronização com a CPU.
  • ECHO (Elastic Cloning): Clona especialistas populares para balancear a carga e reduzir a fragmentação de memória.
  • Paged Stashing: Gerencia memória de forma dinâmica dentro do CUDA Graph, alocando buffers apenas para o tamanho real dos dados, evitando alocação no pior caso.

2.5. Treinamento de Longo Contexto e RL

• Context Parallelism (CP) Dinâmico: Para sequências variáveis (comuns em RL), o sistema ajusta dinamicamente o tamanho do CP por micro-batch para equilibrar memória e computação, evitando desperdício de comunicação em sequências curtas.
• Suporte a RL: Integração com Megatron-Bridge para conversão de checkpoints (Hugging Face <-> Megatron), suporte a packed sequences e router replay para estabilidade no pós-treinamento por RL.

3. Principais Contribuições

1. Paradigma de Paralelismo Desacoplado: A introdução do Parallel Folding que permite otimizar independentemente camadas densas e esparsas, superando limitações de frameworks anteriores.
2. Stack de Otimização Integrada: Uma abordagem sistêmica que trata Memória, Comunicação e Computação como um sistema interconectado, onde otimizações em uma área habilitam melhorias em outras (ex: FP8 libera memória para permitir sobreposição de comunicação).
3. Soluções para MoE "Dropless": Técnicas inovadoras (kernels iniciados pelo dispositivo, ECHO, Paged Stashing) que permitem o uso de CUDA Graphs completos em MoE com roteamento dinâmico, eliminando gargalos de CPU.
4. Suporte a Precisão Reduzida Nativa: Implementação robusta de FP8 e FP4 (incluindo NVFP4 no Blackwell) com estratégias de precisão seletiva para manter a estabilidade do treinamento.
5. Recursos de Produção: Suporte a upcycling (conversão de modelos densos para MoE), distributed checkpointing agnóstico ao paralelismo e integração com RL.

4. Resultados de Desempenho

O framework foi validado em modelos de última geração (DeepSeek-V3 e Qwen3) em hardware NVIDIA (H100, GB200, GB300):

• DeepSeek-V3 (685B parâmetros, 256 especialistas):
  • GB300 (Blackwell): 1.233 TFLOPS/GPU.
  • GB200 (Blackwell): 1.048 TFLOPS/GPU.
  • H100 (Hopper): 368 TFLOPS/GPU (em 1024 GPUs).
• Qwen3-235B:
  • GB300: 974 TFLOPS/GPU.
  • GB200: 919 TFLOPS/GPU.
• Eficiência: O uso de FP8/FP4 e otimizações de comunicação reduziu o tempo de comunicação all-to-all de ~60% para menos de 5-10% do tempo total de iteração em configurações otimizadas.
• Escalabilidade: Sucesso no treinamento de modelos com trilhões de parâmetros em clusters de milhares de GPUs, mantendo alta eficiência de utilização (MFU).

5. Significado e Impacto

Este trabalho representa um marco para a comunidade de IA, fornecendo uma solução de código aberto, pronta para produção, capaz de treinar modelos MoE na escala de trilhões de parâmetros de forma eficiente.

• Viabilidade Econômica: Ao reduzir drasticamente o custo de treinamento e permitir o uso eficiente de hardware de última geração (Blackwell/GB200), torna viável o desenvolvimento de modelos maiores e mais capazes.
• Padrão da Indústria: Estabelece novas práticas para o treinamento de MoE, demonstrando que a combinação de paralelismo multidimensional, precisão reduzida e otimizações de baixo nível (kernels/graphs) é essencial para superar os gargalos de sistemas.
• Aceleração da Pesquisa: A disponibilidade do Megatron-Core MoE permite que pesquisadores e empresas foquem na arquitetura do modelo e nos dados, em vez de lutar contra limitações de infraestrutura de treinamento.

Em resumo, o Megatron-Core MoE transforma o treinamento de modelos esparsos de um desafio de engenharia de sistemas complexo em um processo escalável e eficiente, habilitando a próxima geração de modelos de inteligência artificial.