Acelerando o Ajuste Fino de Transformers com NVIDIA NeMo AutoModel

NVIDIA NeMo AutoModel é uma biblioteca de código aberto que faz parte do framework NVIDIA NeMo, projetada para criar modelos generativos de IA personalizados em escala. Ela é construída sobre o Transformers v5, adicionando paralelismo de especialistas (Expert Parallelism), dispatch fused all-to-all do DeepEP e kernels do TransformerEngine, além de usar o carregamento dinâmico de pesos v5 para aplicar essas otimizações a uma ampla gama de famílias de modelos. O resultado é um aumento de 3,4 a 3,7 vezes na taxa de transferência de treinamento e uma redução de 29% a 32% no uso de memória GPU durante o ajuste fino de modelos MoE em comparação com o Transformers v5 nativo usando a mesma API from_pretrained() — basta alterar uma única linha de importação, sem modificar o restante do código. O NeMoAutoModelForCausalLM é uma subclasse do AutoModelForCausalLM, portanto, qualquer código que funcione com modelos Hugging Face também funciona com o AutoModel. Para arquiteturas MoE populares (Qwen3, NVIDIA Nemotron, GPT-OSS, DeepSeek V3), o NeMo AutoModel utiliza implementações especializadas com atenção do TransformerEngine, camadas lineares fundidas e kernels de especialistas personalizados. Para outros modelos, aplica-se um fallback otimizado com kernels Liger. Medições de desempenho mostraram: no ajuste fino completo do modelo Nemotron 3 Ultra 550B A55B em 16 nós (128 GPUs), o NeMo AutoModel atingiu 815 tokens/s por GPU e 293 TFLOP/s por GPU, com pico de consumo de memória de 58,2 GB, enquanto o Transformers v5 não conseguiu ser executado nessa escala devido à falta de memória. Em testes de nó único com 8 GPUs para o Qwen3-30B-A3B, o NeMo AutoModel obteve 11.340 tokens/s por GPU (contra 3.075 do v5, uma aceleração de 3,69x) e pico de memória de 48,1 GB (29% menos que os 68,2 GB do v5). Para o Nemotron 3 Nano 30B A3B, os resultados foram: 15.421 tokens/s (contra 4.583, aceleração de 3,36x) e 42,5 GB de memória (32% menos que os 62,1 GB do v5). A aceleração é alcançada por três fatores: o paralelismo de especialistas reduz a pressão na memória, o DeepEP combina comunicação com computação, e os kernels do TransformerEngine aceleram operações principais.