NVIDIA Opens Cosmos Framework for Training and Running World Models

When developers talk about multimodal networks, they usually mean a combination of text and images. NVIDIA went further and released the Cosmos 3 family — world models where text, video, audio, and sequences of physical actions are combined in a single Mixture-of-Transformers architecture. To work with this family, the company released a repository cosmos-framework.

What's Inside the Framework

The codebase is bundled into a single Python package cosmos_framework. The project is designed for engineers who need to fine-tune multimodal models for specific tasks or deploy them to production. If you're working on robotics or video generation with controllable behavior, this repository will save you a lot of time on writing boilerplate code.

All logic is divided into two major blocks.

The first block handles distributed training. The tool supports FSDP, tensor (TP), context (CP), and pipeline (PP) parallelism. Checkpoints are saved in DCP format, but the code includes conversion to the familiar HuggingFace safetensors. Data preparation is simplified with ready-made adapters for JSONL, WebDataset, HuggingFace, and LeRobot. The main launch script is located at cosmos_framework.scripts.train.

The second block handles inference. Under the hood, you can use Diffusers, Transformers, vLLM backends, or custom wrappers. For background generation, an offline mode is available, and for API creation, a Ray and Gradio stack is deployed. Inference is launched through the cosmos_framework.scripts.inference module.

For robotics tasks, the framework includes a Policy Server. It serves the Cosmos3-Policy-DROID model and passes generated commands directly to actuators.

Hardware Requirements and First Launch

The project's hardware appetite is substantial. Ready-made configs from the examples folder were tested on servers with 8 NVIDIA H100 80GB cards. You can adjust the degree of parallelism and sharding parameters through NPROC_PER_NODE and FSDP settings, but training a model on a single consumer GPU won't be possible.

The project uses the uv package manager for environment setup. The authors recommend starting with the official NVIDIA NGC Docker image (nvcr.io/nvidia/pytorch:25.09-py3).

System dependency installation looks like this:

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

Building a virtual environment for CUDA 13.0 is done with a single command:

uv sync --all-extras --group=cu130-train
source .venv/bin/activate && export LD_LIBRARY_PATH=

To test inference on a single GPU, the repository includes a ready-made preset:

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

Automatic Integration with Coding Agents

An interesting detail: the authors built AI assistant support into the repository structure. The .agents/skills/ directory contains ready-made instructions for Cursor, Claude Code, and Codex CLI.

Inside you'll find several files in SKILL.md format. Each instruction covers its own area:

• Environment setup, weight downloading, and working with Docker
• Navigating the cosmos_framework package files
• Inference and parallelism parameter configuration
• Full SFT fine-tuning cycle from data preparation to export
• Diagnosing CUDA errors, OOM, and container crashes

When you ask Claude Code or Cursor to fix an import error or modify a parameter in the launch script, the agent automatically loads the corresponding rules file from the .agents/skills/ directory.

Is It Worth Trying

The framework will be useful for robotics teams and generative video system developers. The repository is still young, but it already includes a complete set of tools for the model workflow: from dataset loading to physical hardware control.

If you have access to GPU clusters and are looking for a ready-made framework for experimenting with world models, check out cosmos-framework. Scripts from the examples folder will give you a good start.