NeMo-RL v0.3 dévoile un gain de performance drastique avec Megatron-Core pour l’entraînement accéléré de modèles LLM à grande échelle

La version 0.3 de NVIDIA NeMo-RL introduit un support majeur pour Megatron-Core, une bibliothèque optimisée pour les GPU, permettant un entraînement accéléré de modèles linguistiques massifs. Initialement basé sur PyTorch DTensor (FSDP2), le framework offrait une intégration fluide avec Hugging Face et une scalabilité via les parallélismes natifs de PyTorch. Toutefois, à mesure que les modèles atteignent des centaines de milliards de paramètres, le chemin DTensor se révèle insuffisant : la mémoire d’activation croissante entraîne un surcoût de recalcul important, ralentissant considérablement les étapes d’entraînement, et manque de kernels CUDA optimisés pour une performance maximale. Megatron-Core, conçu pour surmonter ces limites, exploite une stratégie de parallélisme 6D (data, tensor, pipeline, sequence, expert, contexte) qui optimise la communication et le calcul, tout en supportant à la fois les modèles denses et les Mixture of Experts (MoE). NeMo-RL l’intègre désormais, permettant aux développeurs d’exploiter ces optimisations lors de l’entraînement postérieur. Pour simplifier la configuration, NeMo-RL gère automatiquement les paramètres complexes de Megatron-Core, offrant une interface intuitive via un fichier YAML. Il suffit d’ajouter une section policy.megatron_cfg avec enabled: True pour activer le backend. Les résultats montrent une amélioration significative de la throughput. Sur un modèle Llama 3.1-8B, le temps par étape passe de 285 s (DTensor) à 187 s (Megatron), avec une réduction de 34 %. Pour Llama 3.1-70B, la baisse est encore plus marquée : de 2309 s à 1484 s, soit une économie de 35 %, sans perte de convergence. Ces gains sont confirmés par les courbes de récompense, qui montrent une progression similaire entre les deux backends, mais avec une stabilité et une rapidité supérieures sous Megatron-Core. Deux techniques clés renforcent ces performances : la sequence packing, qui regroupe plusieurs séquences dans une longueur maximale (4096), réduisant les tokens de padding et permettant une réduction d’environ 1x du temps d’étape pour Llama 70B ; et l’importance sampling, qui corrige les écarts entre les probabilités d’inférence et d’entraînement, améliorant la stabilité et la convergence. Le support du contexte long est également amélioré : pour Llama 3.3-70B à 16 384 tokens, le temps par étape est de 445 s, démontrant la capacité de Megatron-Core à gérer des fenêtres contextuelles étendues. Des optimisations continues sont en cours pour aller encore plus loin. En outre, NeMo-RL v0.3 inclut des fonctionnalités comme la gestion des récompenses personnalisées, des outils d’évaluation avancés et des scripts de reproduction simplifiés. De nouveaux éléments sont annoncés prochainement : prise en charge des modèles MoE avec gestion dynamique des experts, améliorations de l’efficacité mémoire, et intégration avec des systèmes de déploiement distribués. En résumé, NeMo-RL v0.3, avec son backend Megatron-Core, représente une avancée majeure pour l’entraînement par renforcement à grande échelle. Il allie performance GPU-optimisée, scalabilité, et facilité d’utilisation, rendant l’entraînement de modèles comme Llama 70B ou Qwen3 32B plus accessible et efficace. Les utilisateurs sont invités à explorer la documentation, les exemples et les scripts disponibles pour commencer leurs propres expérimentations.