Aprendizaje por Refuerzo con NVIDIA NeMo-RL: Megatron-Core para un rendimiento de entrenamiento optimizado

TL;DR

• NeMo-RL ahora soporta Megatron-Core como backend para aumentar el rendimiento de entrenamiento en modelos muy grandes.
• Megatron-Core ofrece paralelismo en 6D, kernels optimizados para GPU y funciones como empaquetado de secuencias y muestreo de importancia.
• El entrenamiento sigue siendo accesible a través del flujo DTensor, con una configuración más simple para el entrenamiento basado en Megatron.
• Existe soporte para entrenamiento con contexto largo (hasta 16K longitudes de secuencia), con resultados demostrados en modelos como Llama 3.3 70B.
• Se soportan modelos densos y MoE (Mezcla de Expertos); habilitar el entrenamiento Megatron es tan sencillo como actualizar la configuración YAML. publicación de NVIDIA.

Contexto y antecedentes

La versión inicial de NVIDIA NeMo-RL incluyó soporte de entrenamiento a través de PyTorch DTensor (conocido también como FSDP2). Este backend facilita la integración nativa con el ecosistema HuggingFace, la experimentación rápida y el escalado con paralelismos nativos de PyTorch (FSDP2, paralelismo de tensores, paralelismo de secuencia y de contexto). Sin embargo, cuando los modelos alcanzan cientos de miles de millones de parámetros, el camino DTensor resulta insuficiente: la memoria de activación y el overhead de recomputación hacen que los tiempos por paso sean impracticablemente lentos. Además, el camino DTensor carece de kernels NVIDIA CUDA optimizados y de otras mejoras de rendimiento necesarias para un rendimiento máximo. Estas limitaciones impulsaron la necesidad de una solución más eficiente, precisamente la que ofrece la biblioteca Megatron-Core. Para quienes quieran explorar las últimas novedades, NeMo-RL v0.3 incorpora soporte de Megatron-Core con documentación detallada, scripts de ejemplo y archivos de configuración para entrenar modelos grandes con Megatron-Core. publicación de NVIDIA.

Novedades

NeMo-RL v0.3 introduce soporte para Megatron-Core como backend, construido con técnicas optimizadas para GPU y mejoras de rendimiento para permitir el entrenamiento fluido de modelos lingüísticos masivos. Megatron-Core implementa una estrategia de paralelismo en 6D que optimiza los patrones de comunicación y cómputo, soportando una amplia gama de arquitecturas de modelo. Esto facilita un post-entrenamiento eficiente para modelos extremadamente grandes. NeMo-RL facilita la experiencia de configuración para que los desarrolladores aprovechen estas optimizaciones de Megatron sin necesidad de ajustes de bajo nivel.

• Megatron-Core ofrece múltiples configuraciones de bajo nivel; NeMo-RL oculta gran parte de esos ajustes y propone un flujo de configuración más intuitivo.
• Habilitar el entrenamiento basado en Megatron es sencillo: añade la sección policy.megatron_cfg a tu YAML y establece enabled=True. Todos los argumentos del config se transferirán a Megatron durante el entrenamiento.
• Los backends Megatron-Core funcionan junto a DTensor; el lanzamiento del entrenamiento se realiza de la misma forma que con DTensor, solo con configuraciones extendidas para Megatron.
• Megatron-Core soporta tanto modelos densos como MoE.
• Esta versión incluye datos de rendimiento y ejemplos para ayudar a los desarrolladores a reproducir y optimizar sus ejecuciones.
• El soporte para contexto largo está confirmado, con resultados para Llama 3.3 70B a 16K de longitud; contextos más largos son compatibles y se siguen optimizando.
• Además del backend Megatron, NeMo-RL v0.3 introduce características para democratizar el post-entrenamiento eficiente a través de modelos de distintas escalas.

Por qué importa (impacto para desarrolladores/empresas)

Para equipos que trabajan con grandes modelos de lenguaje, el Megatron-Core integrado en NeMo-RL ofrece un camino hacia un mayor rendimiento de entrenamiento en comparación con DTensor, especialmente para modelos extremadamente grandes. La combinación de kernels optimizados de GPU, estrategias de paralelismo y características como empaquetamiento de secuencias y muestreo de importancia se traduce en ciclos de iteración más rápidos y propiedades de convergencia más manejables a escala. Las empresas que buscan post-entrenar y ajustar modelos masivos pueden aprovechar Megatron-Core para lograr entrenamiento eficiente con ventanas de contexto más amplias, hasta 16K secuencias, sin sacrificar la convergencia. Este esfuerzo facilita que equipos obtengan resultados reproducibles en entrenamiento a gran escala. El trabajo también demuestra que la integración NeMo-RL + Megatron-Core ofrece mejoras de rendimiento en comparación con DTensor, manteniendo el mismo comportamiento de convergencia. El contexto largo permanece como foco, con resultados confirmados y optimizaciones en curso.

Detalles técnicos o Implementación

• Megatron-Core provee una estrategia de paralelismo en 6D para optimizar la comunicación y el cómputo en el entrenamiento a gran escala, soportando una amplia gama de arquitecturas, incluyendo modelos densos y MoE.
• La integración NeMo-RL equipa Megatron-Core con un flujo de configuración simplificado: los usuarios añaden policy.megatron_cfg al YAML y, al activar enabled=True, el entrenamiento continúa con Megatron-Core sin cambios significativos.
• Los backends Megatron-Core trabajan junto a DTensor; el lanzamiento del entrenamiento es similar al de DTensor, con configuraciones extendidas para Megatron.
• Megatron-Core incluye kernels optimizados por GPU y optimizaciones de rendimiento como paralelismo 4D, empaquetado de secuencias y muestreo de importancia para aumentar throughput y convergencia.
• El soporte para contexto largo está confirmado, con resultados para Llama 3.3 70B a 16K; contextos más largos son compatibles y las optimizaciones para contexto largo continúan.
• El diseño soporta tanto modelos densos como MoE, ampliando el conjunto de arquitecturas que pueden ser post-entrenadas con eficiencia.

Conclusiones clave

• El backend Megatron-Core mejora significativamente el throughput de entrenamiento para modelos muy grandes frente al camino DTensor.
• Permite paralelismo en 6D para optimizar la comunicación y el cómputo en gran escala.
• La integración se centra en la usabilidad, con configuración simplificada a través de policy.megatron_cfg y enabled=True.
• El entrenamiento con contexto largo es soportado, con demostraciones de 16K de longitud de secuencia.
• Soporta modelos densos y MoE, ampliando las arquitecturas aptas para post-entrenamiento eficiente.
• La versión incluye datos de rendimiento y ejemplos para ayudar a las comunidades a reproducir y optimizar ejecuciones.

Referencias

• https://developer.nvidia.com/blog/reinforcement-learning-with-nvidia-nemo-rl-megatron-core-support-for-optimized-training-throughput