NVIDIA Transformer Engine acelera treinamento dropless de modelos MoE em JAX
A arquitetura conhecida como Mixture of Experts, abreviada como MoE, tornou-se uma das principais tendências no treinamento de modelos de inteligência artificial em larga escala. Modelos como DeepSeek, Qwen e Mixtral são exemplos dessa abordagem e conseguem igualar ou superar o desempenho de modelos densos tradicionais consumindo apenas uma fração do poder computacional necessário para o treinamento. Essa eficiência é alcançada por meio de um conceito chamado computação condicional, em que diferentes partes da rede neural, chamadas de especialistas, são acionadas conforme a necessidade em vez de processar todas as informações de forma uniforme.
Em vez de utilizar uma única rede feed-forward densa compartilhada por todo o modelo, a abordagem MoE distribui o trabalho entre múltiplos especialistas especializados. Um mecanismo de roteamento, geralmente chamado de gate, decide quais especialistas devem processar cada trecho de informação, chamado de token. Isso permite que o modelo tenha um grande número total de parâmetros, mas apenas uma parte deles seja efetivamente utilizada a cada passo de inferência ou treinamento. O resultado é uma redução significativa do custo computacional sem perda de capacidade representativa.
Apesar das vantagens, treinar modelos MoE em larga escala apresenta desafios técnicos importantes. Um dos principais está relacionado à forma como os tokens são distribuídos entre os especialistas. A abordagem mais comum em frameworks como JAX e MaxText mantém os tensores de cada especialista com um formato fixo e descarta os tokens que excedem a capacidade de processamento de cada especialista. Esse método, conhecido como dropless, busca preservar a qualidade do modelo, mas esbarra em um limite de memória que torna seu uso inviável em implementações puras baseadas em JAX.
Para superar essa limitação, a NVIDIA publicou em seu blog técnico uma abordagem que utiliza o Transformer Engine, biblioteca voltada para a aceleração de modelos Transformer em GPUs da empresa. A solução combina kernels otimizados para a arquitetura MoE com suporte nativo ao framework JAX, permitindo o treinamento dropless de modelos MoE sem a necessidade de descartar tokens. O Transformer Engine também oferece suporte a formatos de precisão reduzida, como o FP8, que diminuem o uso de memória e aumentam o throughput de processamento em arquiteturas como Hopper, Ada e Blackwell.
A integração do Transformer Engine com o JAX permite ainda que desenvolvedores aproveitem blocos de construção otimizados para a arquitetura Transformer em conjunto com implementações especializadas para camadas MoE. Entre os blocos oferecidos estão mecanismos de atenção fundidos, camadas lineares otimizadas e operações de normalização, todos projetados para reduzir a sobrecarga de inicialização de kernels e o tráfego de memória entre a GPU e o sistema principal. Isso resulta em tempos de treinamento menores e maior eficiência no uso dos recursos computacionais disponíveis.
De acordo com a abordagem descrita pela NVIDIA, o uso combinado do Transformer Engine com kernels especializados para MoE torna viável o treinamento dropless em JAX, algo que anteriormente exigia soluções alternativas com maior consumo de memória. A expectativa é que essa integração facilite a adoção de modelos MoE em projetos de pesquisa e aplicações industriais, ampliando o acesso a arquiteturas de maior capacidade com menor custo computacional. A tendência reforça o papel do MoE como peça central na evolução dos grandes modelos de linguagem e no desenvolvimento de sistemas de inteligência artificial mais eficientes.
