NVIDIA NeMo AutoModel acelera el ajuste fino de MoE 3.7x con paralelismo de expertos y DeepEP

NVIDIA ha lanzado NeMo AutoModel, una biblioteca de código abierto dentro del marco NeMo que aporta importantes ganancias de rendimiento al ajuste fino de modelos de mezcla de expertos (MoE). Construida directamente sobre Hugging Face Transformers v5, la biblioteca ofrece un rendimiento de entrenamiento de 3.4 a 3.7 veces mayor y reduce el consumo de memoria GPU entre un 29 y un 32 por ciento en comparación con la mejor configuración disponible de Transformers v5, todo mientras mantiene la misma API from_pretrained(). Los usuarios solo tienen que cambiar una línea de importación para obtener estos beneficios.

Antecedentes: El desafío de eficiencia de MoE

El auge de los modelos MoE ha introducido nuevos obstáculos para el entrenamiento eficiente. Enrutar tokens a través de cientos de expertos, fusionar multiplicaciones de matrices de expertos en kernels individuales, distribuir pesos entre GPUs y superponer comunicación con cómputo requiere una infraestructura que las bibliotecas de uso general normalmente no proporcionan de serie.

Transformers v5 agregó soporte de primera clase para MoE, incluyendo backends de expertos, carga dinámica de pesos y planes de paralelismo tensorial para ejecución distribuida. Pero piezas clave de rendimiento, como la superposición eficiente de comunicación, quedaron sin resolver. NeMo AutoModel llena ese vacío.

Cómo funciona NeMo AutoModel

NeMo AutoModel subclasifica AutoModelForCausalLM e incorpora tres optimizaciones importantes: Paralelismo de Expertos (EP), despacho fusionado all-to-all de DeepEP y kernels de TransformerEngine. DeepEP superpone la comunicación con el cómputo de expertos, una capacidad que Transformers v5 no posee. Debido a que NeMo AutoModel aprovecha la conversión reversible de pesos de v5 para cargar modelos, concentra el esfuerzo de ingeniería en operaciones centrales reutilizables en lugar de en la plomería de checkpoints por modelo.

La biblioteca es completamente compatible con la API de Hugging Face Transformers. Cargar un modelo solo requiere un cambio en la declaración de importación, como muestran los ejemplos de código de NVIDIA.

Evaluaciones comparativas de rendimiento

Multi-nodo: Nemotron 3 Ultra 550B A55B

Para el ajuste fino completo del modelo Nemotron 3 Ultra A55B de 550B parámetros en 16 nodos H100 (128 GPUs), NeMo AutoModel con Paralelismo de Expertos configurado en 64 alcanzó 815 tokens por segundo por GPU, 293 TFLOPs por segundo por GPU y un pico de memoria de 58.2 GiB. Es importante destacar que Transformers v5 se queda sin memoria a esta escala, lo que hace imposible una comparación directa sin EP.

Nodo único: Qwen3-30B-A3B

En un solo nodo con 8 GPUs H100 de 80 GB, NeMo AutoModel (EP=8) registró 11,340 TPS/GPU, un salto de 3.69x sobre los 3,075 TPS/GPU de Transformers v5. El pico de memoria cayó de 68.2 GiB a 48.1 GiB, una reducción del 29 por ciento. Notablemente, Transformers v4 se bloqueó por completo en este modelo debido al envoltorio incompatible de FSDP de módulos expertos individuales.

Nodo único: Nemotron 3 Nano 30B A3B

Para el modelo Nemotron Nano de 30B, NeMo AutoModel alcanzó 15,421 TPS/GPU, una mejora de 3.36x sobre los 4,583 TPS/GPU de v5. El consumo máximo de memoria disminuyó de 62.1 GiB a 42.5 GiB, una reducción del 32 por ciento.

Fuentes de aceleración

Las ganancias de rendimiento provienen de tres optimizaciones coordinadas:

• Paralelismo de Expertos fragmenta los pesos de los expertos entre las GPUs. Con EP=8 en 8 GPUs, cada GPU posee solo un octavo de los parámetros de los expertos. Para un modelo con aproximadamente 55 GiB de pesos de expertos, EP reduce la huella por GPU a aproximadamente 6.8 GiB, lo que permite el entrenamiento donde los enfoques solo con FSDP se quedan sin memoria.
• DeepEP fusiona el despacho de tokens y la combinación en kernels de GPU optimizados, superponiendo la comunicación con el cómputo de expertos en lugar de usar colectivos separados AllGather y ReduceScatter.
• Kernels de TransformerEngine aceleran operaciones centrales como atención fusionada, capas lineales y RMSNorm, ofreciendo aceleraciones consistentes en todos los tipos de capas.

Aprovechando las características de Transformers v5

NeMo AutoModel se basa en varias características clave de v5, incluyendo el sistema de backends de expertos. v5 introdujo el parámetro experts_implementation con tres opciones: eager (depuración con bucle for), batched_mm (GEMM por lotes único) y grouped_mm (ejecución de tokens ordenados). NeMo AutoModel extiende eso usando el despacho fusionado all-to-all de DeepEP combinado con kernels GEMM agrupados y capas lineales de TransformerEngine.

La carga dinámica de pesos a través de WeightConverter y WeightRenaming de v5 permite que los checkpoints de MoE se almacenen en tensores 3D fusionados. Más de 20 tipos de modelos usan este mecanismo, incluyendo Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3 y OLMoE. Las conversiones son completamente reversibles: save_pretrained() produce safetensors en formato estándar de Hugging Face que herramientas posteriores como vLLM y SGLang pueden cargar.

Conclusión y disponibilidad

NVIDIA NeMo AutoModel ofrece una ruta de actualización sin fricciones para los usuarios de Hugging Face que escalan el entrenamiento de modelos. Al cambiar una sola línea de importación, los usuarios obtienen una instancia de modelo que es más de tres veces más rápida con una huella de memoria mucho menor. El código, las configuraciones y los scripts de evaluación comparativa están disponibles en el repositorio de NeMo AutoModel. Los contribuyentes principales del trabajo incluyen a Adil Asif, Hemil Desai, Alexandros Koumparoulis y Huiying Li.