NVIDIA NeMo-RL Megatron-Core: Rendimiento de entrenamiento optimizado

Visión general

NVIDIA NeMo-RL (NeMo Reinforcement Learning) comenzó con soporte para el entrenamiento a través de PyTorch DTensor (también conocido como FSDP2), lo que facilita la integración nativa con HuggingFace y la escalabilidad con paralelismos nativos de PyTorch, como DTensor/FSDP2, paralelismo de tensor, paralelismo de secuencia y paralelismo de contexto. A medida que los tamaños de los modelos se acercan a los cientos de miles de millones de parámetros, el camino DTensor resulta insuficiente debido al tamaño de la memoria de activación y al sobrecosto de recomputación, lo que conlleva tiempos de paso prohibitivamente lentos; además, carece de kernels CUDA optimizados y otras mejoras de rendimiento necesarias para un rendimiento óptimo. Esto destacó la necesidad de una solución más eficiente, que Megatron-Core está diseñado para proporcionar. En el lanzamiento de NeMo-RL v0.3, NVIDIA añade soporte para el backend Megatron-Core, con documentación detallada, scripts de ejemplo y archivos de configuración para entrenar modelos grandes de forma eficiente tras el preentrenamiento. Megatron-Core ofrece técnicas optimizadas para la GPU y mejoras de alto rendimiento, basadas en una estrategia de paralelismo 6D que optimiza los patrones de comunicación y cálculo y soporta una amplia gama de arquitecturas de modelos. NeMo-RL ahora permite a los desarrolladores aprovechar estas optimizaciones durante el post-entrenamiento. NeMo-RL simplifica el uso de Megatron-Core automatizando gran parte del ajuste complejo y exponiendo una configuración más simple e intuitiva. Activar el entrenamiento basado en Megatron es directo: añade la sección policy.megatron_cfg en tu configuración YAML y establece enabled: true. Un ejemplo completo se encuentra en la documentación, y todos los argumentos dentro de la configuración se enviarán a Megatron durante el entrenamiento. Después de añadir la sección Megatron y activar enabled, ya estás listo para entrenar un modelo. Lanzar el entrenamiento se realiza de la misma forma que con DTensor, como se describe en el README o en la guía para reproducir DeepScaleR. El entrenamiento basado en Megatron-Core admite modelos densos y Mixture-of-Experts (MoE). Los resultados muestran desgloses de tiempo por paso para GRPO (Group Relative Policy Optimization) en varios modelos, demostrando un rendimiento superior a DTensor manteniendo las mismas propiedades de convergencia. Se destacan kernels optimizados por GPU, paralelismo 4D, empaquetado de secuencias y muestreo de importancia como factores que mejoran el rendimiento, además del soporte para entrenamiento de contexto largo de hasta 16K longitudes de secuencia. El paralelismo de contexto con Megatron-Core y DTensor también es compatible para entrenamientos de contexto largo, como se observa en los resultados para Llama 3.3 70B. El mensaje es claro: entrenar con Megatron-Core mejora significativamente el rendimiento de RL para modelos muy grandes sin sacrificar la convergencia. Para investigadores y responsables técnicos que buscan post-entrenamiento de RL en modelos muy grandes, Megatron-Core ofrece un camino hacia un mayor rendimiento, mayor contexto y cobertura de modelos más amplia, apoyado por documentación, configuraciones de ejemplo y scripts para facilitar la adopción. La documentación de NeMo RL, las configuraciones de ejemplo y los scripts ayudan a iniciar el post-entrenamiento de sus grandes modelos con optimizaciones Megatron-Core.

Características clave

Soporte del backend Megatron-Core en NeMo-RL v0.3, habilitando optimizaciones de post-entrenamiento para modelos grandes.
Kernels optimizados para GPU y mejoras de alto rendimiento.
Estrategia de paralelismo de 6D para optimizar la comunicación y el cálculo; soporta una amplia variedad de arquitecturas de modelos.
Paralelismo 4D en conjunto con Megatron-Core, además de características como empaquetado de secuencias y muestreo de importancia.
Soporte para contexto largo, con pruebas hasta 16K longitudes de secuencia.
Paralelismo de contexto con Megatron-Core y DTensor para entrenamientos de contexto largo.
Soporte para modelos densos y MoE (Mixture-of-Experts).
Configuración simplificada mediante policy.megatron_cfg; ajuste automático de muchos parámetros.
Compatibilidad con flujos de trabajo DTensor existentes; el inicio del entrenamiento sigue el camino DTensor.
Documentación, ejemplos completos y scripts para ayudar a reproducir resultados como los mostrados para Llama 3.3 70B.

Casos de uso comunes

Post-entrenamiento de RL para modelos de lenguaje extremadamente grandes (tamaños de cientos de miles de millones de parámetros).
Escenarios que requieren entrenamiento de contexto largo (probado hasta 16K longitudes de secuencia) para mantener el contexto a lo largo de entradas largas.
Experimentos y flujos de trabajo de producción que involucran modelos densos y arquitecturas MoE.
Experimentación rápida e integración con el ecosistema HuggingFace, aprovechando las optimizaciones de Megatron-Core para el rendimiento.

Configuración e instalación

Para activar Megatron-Core en NeMo-RL, agregue la configuración Megatron a su YAML:

policy.megatron_cfg:
enabled: true

Esto sigue la directriz de que el entrenamiento basado en Megatron se activa al agregar la sección megatron al config y establecer enabled: true. Todos los argumentos del config se transmitirán a Megatron durante el entrenamiento. Después de añadir la sección Megatron y activarla, ya está listo para entrenar. Lanzar el entrenamiento se realiza de la misma forma que con DTensor, como se describe en el README o guía de reproducción DeepScaleR. El entrenamiento Megatron-Core también admite modelos densos y MoE. Para un inicio mínimo, asegúrese de que la sección Megatron esté presente y habilitada en el config y siga el flujo de entrenamiento DTensor documentado en el README del proyecto.

Quick start

Verifique que Megatron esté habilitado en su configuración de NeMo-RL:

policy.megatron_cfg:
enabled: true

Lance el entrenamiento con la misma ruta de entrada DTensor descrita en la guía DTensor. Consulte la guía DTensor y el README de NeMo-RL para los comandos exactos y scripts para reproducir DeepScaleR.
Para un ejemplo mínimo, use el snippet YAML anterior y, a partir de ahí, consulte la documentación para el ejemplo completo de configuración y los scripts proporcionados.

Pros y contras

Pros
Rendimiento de entrenamiento sustancialmente mayor para modelos grandes frente al camino PyTorch DTensor.
Kernels optimizados por GPU y paralelismo 4D/6D para reducir cuellos de botella de comunicación y mejorar la eficiencia computacional.
Soporte para contexto largo (probado hasta 16K) y para modelos de gran escala como Llama 3.3 70B.
Soporte para modelos densos y MoE.
Configuración simplificada a través de policy.megatron_cfg; ajuste automático facilita la adopción para nuevos usuarios.
Paralelismo de contexto con Megatron-Core y DTensor para escenarios de contexto largo.
Contras
Megatron-Core introduce múltiples opciones de configuración de bajo nivel; aunque NeMo-RL automatiza gran parte del ajuste, algunos usuarios avanzados pueden necesitar ajustar parámetros para cargas de trabajo específicas.
La configuración implica añadir y activar configuraciones Megatron específicas en YAML, lo que representa un paso adicional frente al flujo DTensor básico.

Alternativas (comparaciones rápidas)

| Aspecto | PyTorch DTensor (FSDP2) | Megatron-Core (NeMo-RL) |---|---|---| | Rendimiento en modelos muy grandes | Pasos más lentos debido a la memoria de activación y overhead de recomputación | Mayor rendimiento gracias a kernels optimizados para GPU y paralelismo 6D |Propiedades de convergencia | Mantiene las propiedades de convergencia estándar | Mantiene las propiedades de convergencia, optimizado para rendimiento |Complejidad / ajuste | Configuración de superficie más simple; puede requerir ajuste manual para escalar | Automatiza gran parte del ajuste; configuración más sencilla vía policy.megatron_cfg |Soporte de modelos | Amplio soporte a través de paralelismos PyTorch | Soporta modelos densos y MoE; contexto largo soportado |Longitud de contexto | Depende de las capacidades de DTensor | Demostrado hasta 16K longitudes de secuencia; mejoras en contexto largo en curso |