NVIDIA acelera treino de MoE em JAX: 10x mais throughput no DeepSeek-V3

Mixture of experts (MoE) se tornou uma das tendências arquiteturais mais importantes no treinamento de modelos de IA em larga escala. DeepSeek, Qwen e Mixtral são exemplos de modelos MoE que igualam ou superam equivalentes densos usando uma fração do custo de treinamento — graças à computação condicional: em vez de uma única rede feed-forward compartilhada por todos os tokens, o MoE usa vários especialistas menores e um roteador aprendido que decide quais deles ativar.

O desafio é tornar esse treinamento eficiente em escala. No treinamento do DeepSeek-V3 em GPUs NVIDIA GB200, uma linha de base não otimizada atingia apenas 103 TFLOPS/GPU, com a comunicação entre GPUs consumindo 84% do tempo acumulado de kernel. Com a biblioteca Python JAX e as otimizações direcionadas do Transformer Engine, da NVIDIA, esse número subiu para 1.068 TFLOPS/GPU — uma melhoria de 10,4 vezes.

Por que o MoE é difícil de otimizar

O treinamento de MoE em produção introduz gargalos que não existem em modelos densos: roteamento de tokens, despacho e coleta de especialistas, comunicação all-to-all e GEMMs “ragged” (irregulares). O problema se agrava porque o roteador é aprendido: ao longo do treinamento, a distribuição pode ficar fortemente enviesada conforme o roteador desenvolve preferências. Cada especialista recebe um número diferente de tokens, então não existe um GEMM retangular limpo para agrupar e despachar — o resultado são tensores irregulares.

Há duas formas de lidar com o roteamento. No MoE baseado em capacidade, cada especialista recebe um orçamento fixo de tokens, e o excesso é cortado ou preenchido — computação regular e amigável ao hardware, mas com um trade-off direto entre qualidade e eficiência. No MoE “dropless”, todo token é processado pelo especialista escolhido, não importa quão desigual seja a carga. Isso é melhor para a qualidade do modelo, mas muito mais exigente para o sistema.

As otimizações do Transformer Engine

A NVIDIA detalhou três blocos de construção que tornam o MoE dropless prático no JAX. O primeiro é o grouped GEMM, que resolve todas as multiplicações de matriz dos especialistas em uma única chamada de kernel, cada uma com seu número real de tokens — computando apenas as regiões com tokens válidos. Apoiado por cuBLAS e cuBLASLt, ele oferece utilização total dos Tensor Cores mesmo com formatos irregulares, e nas GPUs Blackwell abre caminho para quantização MXFP8 nos matmuls dos especialistas.

O segundo é a integração de dispatch e combine no expert parallelism (EP), alimentada pelo backend de comunicação NCCL EP, ajustado para os padrões irregulares de tráfego do roteamento. Ele ainda emprega deduplicação de tokens: quando um token é despachado para vários especialistas no mesmo rank, ele atravessa a rede apenas uma vez e é replicado no destino, economizando largura de banda.

O terceiro bloco reúne otimizações adicionais — offloading de ativações para memória do host via JAX e coletivos multistreaming no XLA, que permitem sobrepor transferências InfiniBand entre nós com comunicação NVLink intra-nó, em vez de serializá-las.

Resultados e como reproduzir

O ganho de ponta a ponta é de cerca de 10x no throughput do DeepSeek-V3 671B, e o sistema sustenta 97% de eficiência de escalonamento em 1.024 GPUs — um feito significativo, já que a degradação de comunicação tende a crescer mais rápido que o cálculo em clusters grandes.

As otimizações estão disponíveis no container NVIDIA NGC MaxText, com o Transformer Engine embutido, permitindo reproduzir os resultados diretamente. A configuração básica exige habilitar o bloco TE MoEBlock no MaxText com flags como te_moe_block: true e te_gmm_quantization: "te_mxfp8", além de ajustes específicos do DeepSeek-V3 documentados pela NVIDIA. Para o público técnico brasileiro, é uma referência concreta de como extrair o máximo de clusters de GPUs — um tema caro e central para quem treina modelos em larga escala por aqui.