NVIDIA NeMo-RL v0.3、Megatron-Core対応で大規模モデルの強化学習訓練を高速化
NVIDIAが提供する強化学習フレームワーク「NeMo-RL」の最新バージョンv0.3では、大規模言語モデルの効率的後学習を実現するため、GPU最適化ライブラリ「Megatron-Core」のサポートが追加された。当初のNeMo-RLはPyTorchのDTensor(FSDP2)をバックエンドとして採用しており、Hugging Faceとの連携やスケーラビリティに優れていたが、モデルサイズが数百億パラメータを超える段階では、活性化メモリの再計算オーバーヘッドが顕著になり、処理速度が著しく低下するという課題があった。また、NVIDIA CUDAカーネルの最適化や高性能演算のサポートが不足していた。 これに対し、Megatron-CoreはGPUに特化した6D並列戦略(データ、Tensor、パイプライン、コンテキスト、Sequence、Expert並列)を採用し、通信と計算の最適化を実現。特に大規模モデルにおけるトレーニングステップ時間の短縮が顕著で、Llama 3.1-8Bや70Bモデルでの実測では、DTensorに比べて最大で約20%の高速化を達成。70Bモデルでは、4ウェイのTensor並列と4ウェイのパイプライン並列を活用し、効率的なスケーリングを実現。Qwen3 32Bや30B-A3Bモデルでも同様の性能向上が確認された。 NeMo-RLは、Megatron-Coreの複雑な設定を抽象化し、ユーザーが直感的なYAML設定で利用できるように設計。policy.megatron_cfgセクションを追加するだけで、Megatronベースのトレーニングが可能。さらに、シーケンスパッキングや重要度サンプリングといった機能により、パディングの削減や学習・推論の確率差の緩和が実現。これにより、収束性の安定と性能の向上が両立。 また、16,384トークンという長文シーケンスに対応するコンテキスト並列もサポート。Llama 3.3-70B Instructで16K長文処理を実現し、性能の継続的な最適化が進行中。 このように、NeMo-RL v0.3は、大規模モデルの強化学習トレーニングにおいて、性能と使いやすさの両立を実現。開発者は、Megatron-Coreの高効率性を活用しながら、シンプルな設定で高スループットな学習を実現できる。公式ドキュメントやサンプルスクリプトから、すぐに導入が可能。