Tempo de leitura: 6 minutos
Olá, é um prazer inenarrável estar mais uma vez escrevendo este post. Treinar modelos de inteligência artificial com arquitetura Mixture of Experts (MoE) pode reduzir o custo computacional, mas também cria um problema difícil: cada especialista recebe uma quantidade diferente de tokens. Em vez de uma operação perfeitamente organizada, o hardware precisa lidar com dados irregulares, comunicação entre GPUs e cargas que mudam a cada etapa.
Em uma publicação técnica datada de 14 de setembro de 2026, a NVIDIA apresentou otimizações para treinamento MoE sem descarte em JAX usando o Transformer Engine. Segundo a empresa, a combinação elevou o desempenho de um treinamento do DeepSeek-V3 de 103 para 1.068 TFLOPS por GPU em uma configuração baseada em NVIDIA GB200, um ganho de 10,4 vezes. O resultado depende de uma cadeia de otimizações, e não de um único recurso milagroso.
Por que o treinamento MoE em JAX é difícil
Modelos MoE utilizam vários especialistas menores e um roteador que decide quais deles devem processar cada token. Essa abordagem evita executar uma grande rede densa para todos os tokens, mas introduz uma distribuição desigual de trabalho. Um especialista pode receber muitos tokens, enquanto outro recebe poucos.
Essa irregularidade produz os chamados tensores ragged, estruturas cujos grupos têm tamanhos diferentes. Bibliotecas de alto desempenho normalmente trabalham melhor com matrizes retangulares e previsíveis. Quando o formato dos dados muda constantemente, o sistema pode perder tempo organizando os tokens, transferindo informações e esperando pela comunicação entre GPUs.
O problema fica ainda mais evidente no expert parallelism. Os tokens precisam ser enviados para as GPUs que hospedam seus especialistas, processados e depois devolvidos à posição original. Se esse caminho não for otimizado, as GPUs passam parte do tempo paradas, aguardando dados em vez de executando cálculos.
O que muda no MoE sem descarte
No MoE baseado em capacidade, cada especialista recebe um limite fixo de tokens. Quando esse limite é ultrapassado, alguns tokens podem ser descartados ou o sistema pode adicionar preenchimento. O primeiro caminho economiza recursos, mas elimina parte dos dados do treinamento. O segundo preserva os tokens, porém desperdiça memória e processamento.
O treinamento MoE sem descarte segue outra lógica: todos os tokens selecionados pelo roteador são processados, mesmo quando a distribuição entre especialistas é desigual. Isso tende a preservar melhor o fluxo de dados do modelo, mas exige kernels capazes de trabalhar com tamanhos variáveis.
É como organizar uma fila em que cada guichê recebe um número diferente de pessoas. Em vez de colocar cadeiras vazias para fingir que todas as filas têm o mesmo tamanho, o sistema precisa transportar e processar apenas quem realmente está esperando. A eficiência depende de fazer isso sem criar uma nova fila de gargalos.
Como o Transformer Engine otimiza o treinamento MoE sem descarte
A NVIDIA descreve três blocos principais para tornar esse caminho viável em JAX: quantização MXFP8 orientada a grupos, operações de grouped GEMM para as matrizes dos especialistas e operações otimizadas de expert parallelism para dispatch e combine.
Grouped GEMM para lidar com tamanhos variáveis
Em uma rede feed-forward densa, os tokens passam pela mesma matriz de pesos. No MoE, cada especialista recebe uma quantidade diferente de tokens, o que quebra o formato regular esperado por muitas implementações tradicionais.
Uma alternativa simples seria executar um GEMM separado para cada especialista. Porém, esse método pode exigir cópias de informações do dispositivo para a CPU e impedir o uso eficiente de CUDA Graphs. Outra possibilidade seria preencher todos os grupos até o maior tamanho, desperdiçando processamento.
O grouped GEMM reúne as operações em uma única chamada de kernel e calcula somente as regiões que contêm tokens válidos. De acordo com a NVIDIA, o Transformer Engine usa implementações relacionadas ao cuBLAS e ao cuBLASLt para aproveitar os Tensor Cores mesmo quando as dimensões dos especialistas são irregulares.
Dispatch e combine com expert parallelism
O dispatch envia os tokens para os especialistas corretos, possivelmente localizados em GPUs diferentes. Depois do processamento, o combine reúne os resultados e restaura a ordem necessária para a continuação do modelo.
Em uma implementação sem fusão, essas etapas podem formar uma sequência de operações separadas, com leituras e gravações repetidas na memória. A solução apresentada utiliza uma implementação de expert parallelism integrada ao NCCL EP, backend desenvolvido para padrões de comunicação irregulares associados ao roteamento de tokens.
O sistema também pode deduplicar tokens quando o mesmo conteúdo precisa atravessar uma determinada região de rede mais de uma vez. Assim, o token é enviado uma vez e replicado no destino quando necessário, reduzindo o tráfego. Na prática, o grouped GEMM cuida do trabalho dentro dos especialistas, enquanto o expert parallelism organiza o que acontece ao redor deles.
Outras otimizações usadas pela NVIDIA
O ganho não vem apenas dos kernels de multiplicação de matrizes e da comunicação entre especialistas. A configuração apresentada também combina recursos do JAX e do XLA para reduzir a pressão sobre memória e sobrepor operações que poderiam ficar serializadas.
Offloading de ativações no JAX
Durante o treinamento, algumas ativações intermediárias precisam ser preservadas para a etapa de retropropagação. Em vez de manter tudo na memória da GPU durante todo o percurso, o JAX permite enviar determinados dados para a memória do host e rematerializá-los quando necessário.
Segundo a publicação, resultados de projeções de consulta e valor foram escolhidos para offloading no treinamento do DeepSeek-V3. Essa estratégia pode aliviar gargalos de memória de alta largura de banda, embora também exija atenção ao custo das transferências.
Coletivas em múltiplos fluxos
Operações coletivas de comunicação, como all-reduce e all-gather, podem ser executadas em um único fluxo por padrão. Isso faz com que operações independentes esperem umas pelas outras.
Com o multistreaming do XLA, coletivas independentes podem usar fluxos CUDA diferentes. Dessa forma, transferências entre nós pela InfiniBand podem ser sobrepostas à comunicação dentro do nó por NVLink. O escalonador de ocultação de latência analisa quais operações podem ser combinadas sem criar riscos de deadlock.
Os resultados apresentados para o DeepSeek-V3
A NVIDIA afirma ter observado um ganho de aproximadamente dez vezes no throughput final do treinamento do DeepSeek-V3 com 671 bilhões de parâmetros. O caminho de otimização incluiu grouped GEMM baseado em cuBLAS, coletivas em múltiplos fluxos no XLA, quantização MXFP8, offloading de ativações e uma implementação otimizada de expert parallelism.
Em outro resultado divulgado no artigo, a pilha teria mantido 97% de eficiência de escalonamento em 1.024 GPUs NVIDIA GB300 NVL72. A publicação também apresenta medições em uma faixa de 128 a 2.048 GPUs, mas esses números devem ser interpretados como resultados da configuração e do ambiente testados pela NVIDIA, não como uma garantia para qualquer cluster.
Essa ressalva é importante. O desempenho de modelos MoE depende de fatores como tamanho do modelo, topologia da rede, quantidade de especialistas, estratégia de paralelismo, precisão numérica e versão dos componentes de software. Um resultado obtido em hardware especializado não se transfere automaticamente para uma máquina menor.
Como experimentar o caminho otimizado
As otimizações descritas estão disponíveis, segundo a NVIDIA, em uma imagem do NVIDIA NGC MaxText com o Transformer Engine integrado. A publicação recomenda começar com uma configuração de referência, validar a correção em um modelo MoE pequeno e só depois aumentar a escala.
Entre as métricas sugeridas estão tempo por etapa, TFLOPS por GPU, utilização efetiva do hardware, latência do grouped GEMM e tempo gasto nas fases de dispatch e combine. A documentação de configuração do MoE no MaxText reúne os parâmetros necessários para habilitar o caminho correspondente.
Também é possível consultar a documentação oficial do Transformer Engine antes de reproduzir os testes. Como as configurações de reprodução do DeepSeek-V3 são específicas para aquele modelo e ambiente, copiar os parâmetros sem medir o próprio sistema pode produzir resultados diferentes.
O que esse avanço mostra sobre o futuro do treinamento de modelos
O ponto mais interessante não é apenas o número de 10,4 vezes. A melhoria mostra que o desempenho de modelos MoE depende de várias camadas trabalhando juntas. Um kernel mais rápido não resolve uma rede congestionada; uma comunicação eficiente não compensa operações que desperdiçam memória; e uma boa arquitetura de modelo pode continuar limitada por detalhes da implementação.
O treinamento MoE sem descarte preserva a ideia de processar todos os tokens encaminhados pelos especialistas, mas cobra um preço em complexidade. Para tornar essa abordagem prática, o sistema precisa compreender formatos irregulares, mover dados com menos etapas intermediárias e utilizar melhor computação e comunicação ao mesmo tempo.
A resposta para o problema apresentado no início, portanto, não é simplesmente escolher mais GPUs. É coordenar roteamento, kernels, quantização, memória e interconexões como partes de um mesmo sistema. Quando essa coordenação funciona, arquiteturas MoE podem aproveitar sua computação condicional sem transformar a comunicação entre especialistas no principal gargalo do treinamento.
