Skip to content

Cosmos-RL is a flexible and scalable Reinforcement Learning framework specialized for Physical AI applications.

License

Notifications You must be signed in to change notification settings

nvidia-cosmos/cosmos-rl

Repository files navigation

NVIDIA Cosmos Header

Getting Started

Cosmos-RL is a flexible and scalable Reinforcement Learning framework specialized for Physical AI applications.

Documentation.

System Architecture

Cosmos-RL provides toolchain to enable large scale RL training workload with following features:

  1. Parallelism
    • Tensor Parallelism
    • Sequence Parallelism
    • Context Parallelism
    • FSDP Parallelism
    • Pipeline Parallelism
  2. Fully asynchronous (replicas specialization)
    • Policy (Consumer): Replicas of training instances
    • Rollout (Producer): Replicas of generation engines
    • Low-precision training (FP8) and rollout (FP8 & FP4) support
  3. Single-Controller Architecture
    • Efficient messaging system (e.g., weight-sync, rollout, evaluate) to coordinate policy and rollout replicas
    • Dynamic NCCL Process Groups for on-the-fly GPU [un]registration to enable fault-tolerant and elastic large-scale RL training

Policy-Rollout-Controller Decoupled Architecture

License and Contact

This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use.

NVIDIA Cosmos source code is released under the Apache 2 License.

NVIDIA Cosmos models are released under the NVIDIA Open Model License. For a custom license, please contact cosmos-license@nvidia.com.

About

Cosmos-RL is a flexible and scalable Reinforcement Learning framework specialized for Physical AI applications.

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Contributors 22

Languages