Nvidia acelera treino de modelos MoE em 10,4 vezes com JAX e Transformer Engine

A arquitetura MoE divide a rede neural em especialistas menores ativados sob demanda por um roteador que escolhe os melhores caminhos para cada conjunto de tokens. Modelos como DeepSeek, Qwen e Mixtral adotam essa divisão para igualar ou superar redes densas gastando uma fração do poder computacional. O problema surge quando o roteador aprende a concentrar trabalho em determinados especialistas. Essa preferência desequilibra o volume de dados entre as GPUs. O tráfego gera tensores irregulares que travam bibliotecas comuns, preparadas apenas para matrizes retangulares e uniformes.

Comunicação consumia 84% do tempo dos chips

No teste inicial da Nvidia com o DeepSeek-V3 no hardware GB300, a linha de base sem otimização alcançou apenas 103 TFLOPS por GPU. A comunicação entre os chips consumiu 84% de todo o tempo acumulado de processamento dos kernels. As placas ficavam paradas esperando a troca de mensagens na rede. Para eliminar esse gargalo, a equipe aplicou a abordagem sem descarte chamada dropless MoE. Esse método processa cada token no especialista escolhido sem descartar nada nem preencher matrizes com zeros para forçar formatos retangulares. A qualidade do modelo fica intacta.

A base teórica veio do projeto MegaBlocks, que reformulou o cálculo dos especialistas como multiplicação de matrizes com esparsidade em blocos. A biblioteca Transformer Engine colocou essa lógica em prática com kernels agrupados de GEMM que resolvem contagens variáveis de tokens em uma única chamada. Na camada de rede, a extensão NCCL EP acelerou o paralelismo de especialistas. O mecanismo funde as etapas de envio e recombinação de tensores e elimina tokens duplicados para aliviar o tráfego de dados. A rede trabalha menos.

Escala sustenta 97% de eficiência em 1.024 placas

O sistema combinou técnicas no nível do framework JAX para evitar que a memória esgotasse durante as execuções mais pesadas. A equipe utilizou host offloading no JAX para transferir dados da memória de vídeo e ativou fluxos múltiplos de comunicação coletiva via XLA. O compilador sobrepôs a transmissão física pelas conexões de rede NVLink e InfiniBand ao cálculo matemático das placas. Nos números divulgados pela Nvidia, o conjunto sustentou 97% de eficiência de escalabilidade em 1.024 GPUs no hardware GB300 NVL72 ao treinar a versão de 671 bilhões de parâmetros do DeepSeek-V3 671B. O resultado impressiona no papel.

Os números precisam de leitura cautelosa porque saíram dos próprios laboratórios da fabricante, com hardware de ponta e infraestrutura controlada. A Nvidia não detalhou na publicação o comportamento do sistema fora da topologia NVL72 ou em placas de gerações anteriores. Quem quiser reproduzir os testes precisa acessar o contêiner MaxText com a biblioteca Transformer Engine habilitada no catálogo NGC da empresa. A Nvidia também disponibilizou um guia de configuração do MaxText MoE e a documentação oficial da Transformer Engine para instruir a montagem do ambiente.