nvidia-cosmos/cosmos-rl

何を解決するか

Cosmos-RL は、身体的なAIアプリケーションにおける大規模強化学習（RL）トレーニングの複雑さに対応するように設計されており、エメーブドエージェントのトレーニングに伴う重い計算負荷を管理するためのスケーラブルで柔軟なフレームワークを提供します。

動作方法

このフレームワークは、トレーニングプロセスを専門的な役割に分離する、分離型の単一コントローラーアーキテクチャを使用しています：

• ポリシー（コンシューマー）： トレーニングインスタンス専用のレプリカ。
• ロールアウト（プロデューサー）： ジェネレーションエンジン専用のレプリカ。
• コントローラー： ポリシーとロールアウトレプリカ間での重みとロールアウトを調整する効率的なメッセージングシステム。

パフォーマンスを最大化するために、複数の並列化戦略（Tensor、Sequence、Context、FSDP、Pipeline）をサポートし、低精度トレーニングおよびロールアウト（FP8およびFP4）を活用してメモリと計算のオーバーヘッドを削減します。

対象ユーザー

ロボット工学、自動運転車両、スマートスペースなどの物理的AIに取り組む開発者や研究者で、大規模RLトレーニングワークロードに到達したい方を対象としています。

主な特徴

• 高度な並列化： GPU間でワークロードを分散するための5種類の並列化をサポート。
• 非同期アーキテクチャ： ポリシーとロールアウトレプリカを分離することで、より高い効率性を実現。
• エラスティックトレーニング： 動的NCCLプロセスグループを使用して、GPUの登録・非登録をリアルタイムで可能にし、障害耐性を確保。
• 低精度サポート： FP8およびFP8/FP4精度に最適化されており、トレーニングとロールアウトの高速化を実現。