Skip to content

About

Our inference and training framework to run on the Cosmos Models

Resources

Contributing

Security policy

Stars

577 stars

Watchers

10 watching

Forks

Latest commit

Β 

History

181 Commits

Folders and files

NameName
Last commit message
Last commit date
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

NVIDIA Cosmos

NVIDIA Cosmos | πŸ€— Cosmos 3

Part of the NVIDIA Cosmos project family β€” the training and serving framework repository.

Cosmos-Framework

Cosmos-Framework is an end-to-end framework for training and serving world models, including the Cosmos3 model family. Everything lives in a single top-level cosmos_framework/ Python package:

  • Training β€” distributed FSDP / TP / CP / PP trainer, native DCP checkpoints with HuggingFace safetensors import/export, JSONL / WebDataset / LeRobot dataset adapters. Entry point: cosmos_framework.scripts.train. See docs/training.md.
  • Inference β€” Diffusers / Transformers / vLLM backends with offline batch generation and online serving (Ray + Gradio). Entry point: cosmos_framework.scripts.inference. Ecosystem-facing shim libraries (lightweight standalone wrappers for downstream projects) live under packages/.
  • Policy serving β€” a server-client pipeline for deploying Cosmos3 policies. See docs/action_policy_droid_server.md.
  • Post-training and RL β€” supervised, preference-based, reinforcement-learning, distillation, diffusion, and embodied-policy optimization under cosmos_framework.posttrain. See docs/posttrain/.

Cosmos 3

Cosmos 3 is our newest model family [Report] [Website]. It is a suite of omnimodal world models designed to jointly process and generate language, images, video, audio, and action sequences within a unified Mixture-of-Transformers architecture. By supporting highly flexible input-output configurations, it seamlessly unifies critical modalities for Physical AI β€” effectively subsuming vision-language models, video generators, world simulators, and world-action models into a single framework. For a guided experience to test out Cosmos3, please visit [Cosmos].

Framework Documentation

Setup

For more details and alternative installation methods, see Setup. Before installing, make sure your machine meets the System Requirements. If you want a curated PyTorch + CUDA environment, start from the recommended NVIDIA NGC base image.

Install system dependencies:

sudo apt-get install -y --no-install-recommends curl ffmpeg git-lfs libgl1 libglib2.0-0 libx11-dev tree wget

Install the package with the required uv version, 0.11.29 (pick a dependency group supported by your driver β€” see CUDA Variants):

# CUDA 13.0 (recommended)
uv sync --all-extras --group=cu130-train
# Or, for CUDA 12.8:
# uv sync --all-extras --group=cu128-train
source .venv/bin/activate

Run the GPU and media checks before downloading weights. Preserve any CUDA compatibility libraries required by your driver; do not clear LD_LIBRARY_PATH unconditionally.

If you are starting from the recommended NGC image (nvcr.io/nvidia/pytorch:26.06-py3), see the one-shot quickstart.

Training

For the full guide (data preparation, base-checkpoint conversion, parallelism strategies, mixed precision, resuming), see Training. The number of GPUs required depends on the recipe; the shipped recipes under examples/ are 8-GPU configurations (tested on 8Γ— H100 80 GB) launched via their paired launch shells, e.g.:

bash examples/launch_sft_vision_nano.sh

Users may adjust the GPU count to match their model and underlying hardware architecture β€” tune NPROC_PER_NODE and the parallelism degrees (DP/CP/FSDP shard) in the recipe accordingly.

Transfer-Auto and Sim-Auto

  • Cosmos3-Nano-Transfer-Auto: Start from Cosmos3-Nano to reproduce the training curriculum, or directly fine-tune the released Transfer-Auto checkpoint on your own data. See the Transfer-Auto training skill for both paths.
  • Cosmos3-Nano-Sim-Auto: For domain or sensor adaptation, fine-tune Transfer-Auto first, then distill that adapted checkpoint through TF β†’ TF-dCM β†’ SF-DMD. See the Sim-Auto training skill and training guide.

We recommend completing adaptation before distillation, rather than applying ordinary supervised fine-tuning (SFT) to the released distilled Sim-Auto checkpoint, to preserve its few-step generation behavior. The released Sim-Auto checkpoint can be used directly for inference.

Inference

See Inference for the full guide β€” launch commands, supported modes, parallelism presets, and troubleshooting.

Quick single-GPU launch:

python -m cosmos_framework.scripts.inference \
    --parallelism-preset=latency \
    -i "inputs/omni/t2v.json" \
    -o outputs/omni_nano \
    --checkpoint-path Cosmos3-Nano \
    --seed=0

Policy Server

See Policy Server for the full guide.

Post-training and RL

cosmos_framework.posttrain provides a common runtime for offline and online post-training, from single-host recipes to distributed Ray execution. Start with the Posttrain documentation or follow the step-by-step container run guide.

The repository has two SFT paths. Use the main Training path for Cosmos3 generator, reasoner, and policy recipes built on the framework's native configuration, checkpoint, and distributed-training stack. Use Posttrain SFT for Hugging Face-based models outside the Cosmos3 recipes. It currently supports Qwen text models, dense Qwen3.5/3.8 multimodal models, and Qwen3-VL/Qwen3-VL-MoE, with a common Hugging Face foundation for extending support to additional compatible models.

Algorithm Description Guide
DPO Offline preference optimization DPO
GRPO Group-relative online policy optimization GRPO
PPO Actor-critic policy optimization, including embodied-agent training PPO Β· Embodied PPO examples
OPD On-policy distillation from a teacher model OPD
DiffusionNFT Online reinforcement learning for diffusion generators DiffusionNFT
SFT Supervised fine-tuning β€”

Model integrations include Cosmos3, GR00T-N1.7, Ο€0.5, Alpamayo through AlpaGym, Qwen and Qwen-VL families, and general Hugging Face text and vision-language backends. Algorithm support is defined by runtime and data contracts rather than a fixed model allowlist; shipped recipes demonstrate qualified combinations, while additional compatible models can use the documented model, DataPacker, reward, and simulator extension interfaces.

Agent Skills

Coding agents (Claude Code, Codex CLI, Cursor, and other AGENTS.md-aware tools) can load task-specific instructions for this repo from the bundled AgentSkills. Each skill is a self-contained SKILL.md that the agent invokes automatically when the user's request matches its description.

Skills live in .agents/skills/ (canonical) and are mirrored under .claude/skills/ for Claude Code:

Skill When it activates
cosmos3-setup Installation, environment setup, checkpoint downloads, Docker.
cosmos3-codebase-nav "Where is X" / "where do I change parameter Y" questions across cosmos_framework/.
cosmos3-inference Running offline or online inference, parallelism, sampling parameters.
cosmos3-post-training SFT post-training end-to-end: data prep, DCP conversion, launch, export.
cosmos-automl Reasoner hyperparameter search and GEPA prompt optimization; setup and release status.
cosmos3-env-troubleshoot Diagnosing install/runtime errors (ImportError, CUDA, Docker, checkpoint failures).

See AGENTS.md for the canonical repo map that agents load first; the skills above are referenced from there.

Reference

Topic What it covers
Setup Hardware/software prerequisites, uv install paths, CUDA variants, Docker base image, and base-checkpoint downloading.
Code Structure Repository layout and a per-subpackage tour of cosmos_framework/ β€” where each concern lives and where to add new code.
Training Launching multi-GPU and multi-node runs; parallelism strategies; mixed precision; resuming.
Inference (from a trained checkpoint) Loading a trained checkpoint into one of the inference backends.
Policy Server Running the server-client pipeline for Cosmos3-Policy-DROID.
Post-training and RL SFT, preference optimization, online RL, distillation, diffusion RL, embodied policies, and extension guides.
FAQ Troubleshooting (OOM, NCCL hangs, slow training), environment variables, and common pitfalls.
AGENTS.md + Agent Skills Repo map and task-specific SKILL.md files loaded automatically by AGENTS.md-aware coding agents (Claude Code, Codex CLI, Cursor, etc.).

About

Our inference and training framework to run on the Cosmos Models

Resources

Contributing

Security policy

Stars

577 stars

Watchers

10 watching

Forks

Releases

Packages

Contributors

Languages