Apprentissage par renforcement avec NVIDIA NeMo-RL : Megatron-Core pour un débit d’entraînement optimisé

TL;DR

• NVIDIA NeMo-RL ajoute le backend Megatron-Core dans la version v0.3 pour optimiser le throughput post-formation des modèles très grands.
• Megatron-Core exploite des kernels GPU optimisés, un parallélisme 6D et des techniques comme le packing de séquences et l’échantillonnage d’importance pour augmenter la vitesse de formation sans compromettre la convergence.
• Le contexte long est pris en charge, avec des exemples montrant Llama 3.3-70B sur 16k de longueur de séquence via le backend Megatron.
• La configuration est simplifiée via policy.megatron_cfg; activer Megatron nécessite enabled=True et le lancement reste identique à DTensor.
• Les modèles densés et MoE (Mixture of Experts) sont pris en charge, élargissant l’applicabilité.

Contexte et historique

La première version de NeMo-RL incluait le support de l’entraînement via PyTorch DTensor (aussi appelé FSDP2). Cet backend permet une intégration native avec l’écosystème HuggingFace, des expérimentations rapides et une scalabilité avec les parallélismes natifs de PyTorch (FSDP2, parallélisme tensoriel, parallélisme de séquences et parallélisme de contexte). Toutefois, lorsque les modèles atteignent des centaines de milliards de paramètres, le chemin DTensor devient insuffisant. La mémoire d’activation des grands modèles provoque une surcharge de recomputation et des temps d’étape trop longs. Le chemin DTensor manque également de kernels CUDA optimisés et d’autres améliorations de performance nécessaires pour un débit optimal. Ces défis démontrent le besoin d’une solution plus efficiente, précisément ce que la bibliothèque NVIDIA Megatron-Core est conçue pour fournir. Consultez la dernière version de NeMo-RL (v0.3) qui propose une documentation détaillée, des scripts d’exemple et des fichiers de configuration pour post-entraîner des modèles volumineux avec le backend Megatron-Core. Construit avec des techniques optimisées pour GPU et des améliorations de débit, Megatron-Core permet un entraînement fluide des modèles de langage massifs. La bibliothèque repose sur une stratégie de parallélisme en 6D qui optimise les schémas de communication et les motifs de calcul et prend en charge une large gamme d’architectures de modèles. NeMo-RL a ajouté le support de Megatron-Core, permettant aux développeurs d’utiliser ces optimisations lors du post-entrainement. Bien que Megatron-Core offre de nombreuses configurations de bas niveau, les régler peut être intimidant pour les débutants. NeMo-RL simplifie ce processus en gérant automatiquement une grande partie de l’optimisation complexe et en présentant une série d’options de configuration plus simples et intuitives. Pour activer l’entraînement basé Megatron, ajoutez la section policy.megatron_cfg dans votre YAML : voyez un exemple complet en fonctionnement. Toutes les options du config seront transmises à Megatron pendant l’entraînement. Après avoir ajouté la section Megatron à votre config et défini enabled=True, vous êtes prêt à entraîner un modèle. Le démarrage de l’entraînement se fait de la même manière que pour DTensor, comme décrit dans le README ou notre guide pour reproduire DeepScaleR. L’entraînement Megatron porte aussi bien sur des modèles dense que MoE. La publication décrit une analyse temps par étape pour GRPO (Group Relative Policy Optimization) sur plusieurs modèles couramment utilisés. Le temps mesuré est une moyenne sur les pas 22-29 de chaque entraînement. En tirant parti des optimisations de Megatron-Core, nous obtenons des performances d’entraînement supérieures à DTensor tout en conservant les mêmes propriétés de convergence. Les commandes utilisées pour générer ces courbes de récompense sont fournies. Ces sessions utilisent des optimisations de performance et de convergence afin d’atteindre à la fois le débit optimal et une convergence fiable. Nous pouvons également utiliser le parallélisme de contexte avec Megatron-Core et DTensor pour l’entraînement à contexte long. Par exemple, les résultats actuels pour Llama 3.3-70B sur une longueur de 16K utilisent le backend Megatron. Des longueurs de contexte encore plus longues sont prises en charge, et les optimisations pour le contexte long continuent.

Pourquoi cela compte pour le RL à grande échelle

Le Tableau 2 (ci-dessous) met en évidence comment les backends Megatron-Core peuvent améliorer significativement le débit d’entraînement en RL par rapport à PyTorch DTensor, en particulier pour les modèles extrêmement grands. Ces améliorations conservent les propriétés de convergence et prennent en charge des architectures dense et MoE.

Points de performance avec Megatron-Core

• Fenêtres de contexte long de 16K démontrées sur des modèles volumineux comme Llama 3.3-70B.
• Kernels optimisés pour GPU, parallélisme 4D, packing de séquences et sampling d’importance qui contribuent à un débit supérieur et à une convergence robuste.
• Le contexte long reste une priorité avec des optimisations en cours pour des longueurs encore plus grandes. Pour plus de détails et des étapes pratiques, reportez-vous à la documentation NVIDIA NeMo RL, aux configs d’exemple et aux scripts pour démarrer le post-formation de vos grands modèles avec les optimisations Megatron-Core.

Ce qui est nouveau

La sortie NeMo-RL v0.3 introduit le support du backend Megatron-Core, offrant des flux de travail post-formation optimisés pour les grands modèles de langage. Éléments clés :

• Backend Megatron-Core intégré à NeMo-RL v0.3 pour des flux post-formation avec ajustement automatique et simplification de la configuration.
• Une section policy.megatron_cfg dans YAML pour configurer les options Megatron ; activation via enabled=True. Tous les arguments du config seront transmis à Megatron pendant l’entraînement.
• Compatibilité avec les modèles dense et MoE, élargissant le spectre des architectures pris en charge.
• Support du contexte long, avec démonstrations à 16K de longueur pour des modèles lourds (ex. Llama 3.3-70B).
• Parallélisme 6D et kernels GPU optimisés, avec packing de séquences et sampling d’importance, pour augmenter le débit. NeMo-RL permet aussi le parallélisme de contexte avec Megatron-Core et DTensor pour des formations à contexte long, offrant des configurations flexibles. La version v0.3 met l’accent sur des flux de travail pratiques de post-formation, accompagnés de documentation détaillée, scripts d’exemple et fichiers de configuration pour aider les équipes à post-entrainer leurs grands modèles avec les optimisations Megatron-Core.

Aperçu rapide du flux de configuration

• Ajoutez policy.megatron_cfg à votre YAML.
• Définissez enabled=True.
• Lancez l’entraînement comme pour DTensor.
• Dirigez les arguments du config vers Megatron pour activer les optimisations Megatron-Core.

Modèles et résultats de référence

La documentation relate des améliorations du temps par étape pour GRPO sur plusieurs modèles en utilisant Megatron-Core par rapport à DTensor, avec des propriétés de convergence inchangées. Un exemple cité est Llama 3.3-70B avec 16K de contexte. Ces résultats illustrent une augmentation du débit.

Pourquoi cela compte pour les développeurs/entreprises

Pour les développeurs et les entreprises qui visent l’échelle, Megatron-Core dans NeMo-RL offre une voie pratique pour post-entraînement de modèles de langage massifs avec une efficacité accrue. Grâce à des kernels GPU optimisés, un parallélisme 6D et des techniques avancées comme le packing de séquences et l’échantillonnage d’importance, les équipes peuvent accélérer les cycles d’itération, valider et déployer plus rapidement les grands modèles. Cette approche renforce aussi l’objectif de NeMo-RL: démocratiser le post-formation efficace pour une large palette de modèles, y compris les architectures dense et MoE, et montre l’engagement envers des outils qui évoluent avec la taille des modèles.

Détails techniques ou Implémentation

• Backend: Megatron-Core intégré dans NeMo-RL v0.3 pour les flux post-formation.
• Parallélisme: stratégie 6D pour optimiser communication et calcul.
• Contexte: entraînement de contexte long supporté, démontré jusqu’à 16K pour des modèles lourds.
• Types de modèles: prend en charge les modèles dense et MoE.
• Ajustement: NeMo-RL automatise une grande partie de l’optimisation complexe, tout en exposant une surface de configuration simple.
• Configuration: pour activer, ajoutez policy.megatron_cfg au YAML et définissez enabled=True. Tous les arguments Megatron dans le config seront transposés au Megatron pendant l’entraînement. Le démarrage suit le flux DTensor.
• Caractéristiques de performance: kernels optimisés pour GPU, parallélisme 4D, packing de séquences et sampling d’importance accélèrent le throughput et assurent une convergence stable.
• Parallélisme de contexte: disponible pour les entraînements à contexte long lorsque utilisé avec Megatron-Core et DTensor.

Tableau

| Backend | Avantages clés | Remarques |---|---|---| | DTensor (PyTorch DTensor) | Intégration native avec HuggingFace; supporte FSDP2, parallélisme tensoriel, séquence et contexte | Mémoire d’activation et recomputation augmentent à l’échelle; manque d’optimisations Megatron |Megatron-Core | Kernels GPU optimisés; parallélisme 6D; packing de séquences; sampling d’importance; supporte dense et MoE; contexte long | Débit plus élevé avec convergence préservée; nécessite policy.megatron_cfg pour activer |

Points-clés

• NeMo-RL v0.3 ajoute Megatron-Core pour améliorer le throughput de entraînement de modèles gigantesques.
• Le chemin Megatron-Core offre un débit supérieur à DTensor tout en conservant les propriétés de convergence.
• Un contexte long est pris en charge, avec démonstrations à 16K de longueur sur des modèles lourds comme Llama 3.3-70B.
• Activation Megatron via policy.megatron_cfg et enabled=True, avec un flux de lancement identique à DTensor.
• Supporte à la fois les modèles dense et MoE, avec parallélisme 6D et kernels GPU optimisés.

Références

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