Setup Guide

Skill: .agents/skills/cosmos3-setup/SKILL.md

System Requirements

• NVIDIA GPUs with Ampere architecture (RTX 30 Series, A100) or newer — Hopper (H100) or Blackwell (B200) recommended for full training throughput
• NVIDIA driver compatible with CUDA version
• NVIDIA CUDA >=12.8
• Linux x86-64/aarch64
• glibc >=2.35 (e.g. Ubuntu >=22.04)
• Python >=3.10
• Multi-node training additionally requires a working NCCL setup (IB/RoCE recommended) and a shared filesystem visible to all ranks for checkpoint I/O
• Free disk: ~150 GiB recommended for a first-run inference or training workflow (Hugging Face cache ~90 GiB, uv cache ~20 GiB, run outputs ~30 GiB). See FAQ → Expected disk footprint for the breakdown and how to relocate caches.

Recommended Base Image

Recommended Base Image

For CUDA 13 builds, the NVIDIA NGC PyTorch container is the recommended starting point — it bundles PyTorch + CUDA 13 + cuDNN + NCCL tuned for NVIDIA hardware, plus Apex, TransformerEngine, and Megatron utilities that training infra users commonly need.

FROM nvcr.io/nvidia/pytorch:26.06-py3

For CUDA 12.8 builds, pin to an earlier NGC tag (e.g. nvcr.io/nvidia/pytorch:25.06-py3) that still ships CUDA 12.

Installation

If you encounter issues, see Troubleshooting.

Clone the repository:

git clone git@github.com:NVIDIA/cosmos-framework.git
cd cosmos-framework

The two supported install paths are the recommended base image and the Docker container. For other paths (standalone venv, custom torch/cuda) see Advanced.

Quickstart: From the Recommended Base Image

Quickstart: From the Recommended Base Image

If you started from the recommended base image (nvcr.io/nvidia/pytorch:26.06-py3), the following commands set up the full environment in one go. Run them from the root of this repository (i.e. inside the cosmos-framework/ directory you just cloned):

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

curl -LsSf https://astral.sh/uv/0.11.29/install.sh | sh
source $HOME/.local/bin/env

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

Check the GPU and media runtime before downloading weights. Do not clear LD_LIBRARY_PATH unconditionally: an older host driver may depend on the container's CUDA compatibility libraries.

Docker Container

Docker Container

Please make sure you have access to Docker on your machine and the NVIDIA Container Toolkit is installed.

Build the container:

image_tag=$(docker build -q .)

Run the container:

docker run -it --runtime=nvidia --ipc=host --rm \
  -v .:/workspace -v /workspace/.venv \
  -v /root/.cache:/root/.cache \
  -e HF_TOKEN="$HF_TOKEN" \
  $image_tag

For multi-node training, also bind-mount your shared dataset and checkpoint directories so all ranks see the same filesystem.

Optional arguments:

• --ipc=host: Use host system's shared memory, since parallel torchrun consumes a large amount of shared memory. If not allowed by security policy, increase --shm-size (documentation).
• -v /root/.cache:/root/.cache: Mount host cache to avoid re-downloading cache entries.

If you get docker: Error response from daemon: unknown or invalid runtime name: nvidia, you need to configure docker:

sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

See docker/README.md for additional images and build options.

Advanced

Advanced

Use these paths only when the recommended base image or Docker container are not viable for your environment.

Virtual Environment

Virtual Environment

Install system dependencies:

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

Install uv:

curl -LsSf https://astral.sh/uv/0.11.29/install.sh | sh
source $HOME/.local/bin/env

Install the package using one of the following methods:

UV Sync: fully reproducible environment

Choose the dependency group that matches your CUDA toolkit (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

UV Pip: virtual environment

# Create virtual environment (skip if using an existing environment)
uv venv --clear && source .venv/bin/activate

uv pip install -r pyproject.toml --all-extras --group=cu130-train
uv pip install -e .

UV Pip: system environment

uv pip install --system --break-system-packages -r pyproject.toml --all-extras --group=cu130-train

Custom torch/cuda versions

cuda_name=cu130
torch_name=torch210

# 1. Create and activate the virtual environment
uv venv --clear && source .venv/bin/activate

# 2. Install the desired torch/cuda versions
uv pip install "torch==2.10.0" "torchvision" --torch-backend=$cuda_name

# 3. Install the package with desired extras
uv pip install -r pyproject.toml --all-extras --group=cu130-train

# 4. Install one of the following attention backends:
# * Blackwell
uv pip install "natten==0.21.6.dev6+$cuda_name.$torch_name" -f https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0/natten
# * Hopper
uv pip install "flash-attn-3-nv==1.0.3+$cuda_name.$torch_name" -f https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0/flash-attn-3-nv
# * Ada/Ampere
uv pip install "flash-attn==2.7.4.post1+$cuda_name.$torch_name" -f https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0/flash-attn

If there is no attention backend wheel for your torch/cuda versions, you can build one using cosmos-dependencies.

Optional package extras:

• train: Training infrastructure (FSDP, parallelism, checkpointing, datasets)

CUDA Variants

This repository is training-focused, so the *-train dependency groups are the supported install path. Inference-only groups exist for evaluating trained checkpoints in-tree but are not required for training.

CUDA Version	Training (recommended)	Notes
CUDA 13.0 (recommended)	--group=cu130-train	NVIDIA Driver
CUDA 12.8	--group=cu128-train	NVIDIA Driver

CUDA 13 normally requires driver 580 or newer. With an older driver, such as 570, use the CUDA 12.8 group unless your platform supplies a supported CUDA forward-compatibility setup. If that setup requires a cuda-compat library directory on LD_LIBRARY_PATH, preserve it. Installing a new CUDA toolkit alone does not upgrade the host driver. The CUDA version displayed by nvidia-smi is the driver's supported version, not the version installed in the Python environment; exact equality with torch.version.cuda is not required.

Environment Variables

Export the following before downloading checkpoints or launching training. See environment_variables.md for the full reference.

Variable	Purpose
HF_TOKEN	Hugging Face access token for gated model/dataset downloads. Alternative to uvx hf auth login.
HF_HOME	Cache directory for Hugging Face models and datasets. Recommend ≥ 1 TB free.
IMAGINAIRE_OUTPUT_ROOT	Output root for training DCP checkpoints and logs. Recommend ≥ 1 TB free.
UV_CACHE_DIR	Cache directory for uv-managed dependencies.
LD_LIBRARY_PATH	Preserve required driver/compatibility libraries; remove conflicting library directories only when diagnosing an import failure.

Verify the GPU and media runtime

Run these checks inside the same container and allocation that will run the model:

uv --version
nvidia-smi
ffmpeg -version
python - <<'PY'
import cv2
import torch
from cosmos_framework.utils.generator.torchcodec_dependencies import load_torchcodec

print(f"torch={torch.__version__}, CUDA build={torch.version.cuda}, OpenCV={cv2.__version__}")
torch.cuda.init()
torch.ones(1, device="cuda").add_(1)
torch.cuda.synchronize()
print(f"GPU: {torch.cuda.get_device_name()}")
print(f"TorchCodec: {load_torchcodec().__version__}")
PY

The selected uv executable must match tool.uv.required-version in pyproject.toml (currently 0.11.29). If another version is already installed, run uv self update 0.11.29 or reinstall that exact version and check command -v uv. An import-only check does not establish that CUDA initializes or video decoding works. After downloading an example video, verify decoding too:

export CHECK_VIDEO=/path/to/one/example.mp4
python - <<'PY'
import os
from cosmos_framework.utils.generator.torchcodec_dependencies import load_torchcodec

codec = load_torchcodec()
decoder = codec.decoders.VideoDecoder(os.environ["CHECK_VIDEO"], device="cpu")
print(decoder[0].shape)
PY

The apt packages above supply FFmpeg shared libraries and OpenCV's OpenGL/GLib runtime dependencies. A Python wheel installation does not replace these system libraries. See CUDA troubleshooting for TorchCodec's NPP loading.

Downloading Base Checkpoints

Training in this repo typically starts from a pretrained base checkpoint that you fine-tune or post-train. The recommended source is the Hugging Face Hub.

1. Get a Hugging Face Access Token with Read permission.
2. Authenticate using either mechanism (they are equivalent — pick one, do not set both with different tokens):
3. HF_TOKEN environment variable — preferred for Docker and non-interactive shells. Export it once and any huggingface_hub call (CLI or library) picks it up.
4. uvx hf auth login — preferred for local interactive use. Writes the token to ~/.cache/huggingface/token, persisted across sessions (and across Docker runs if you bind-mount /root/.cache).
5. Accept the license for any gated model you intend to use (e.g. the NVIDIA Open Model License Agreement where applicable).
6. Test access:

uvx hf@latest download --repo-type model nvidia/Cosmos-Guardrail1 \
  --revision d6d4bfa899a71454a700907664f3e88f503950cf --include "README.md"

If you encounter issues:

1. Check which account is active with hf auth whoami. If HF_TOKEN is set, it takes precedence over the saved login; keep these credentials consistent. Keep token values out of logs and support requests.
2. Check that your token has sufficient permissions.

Checkpoints are downloaded on demand during training and evaluation. To change the cache location, set HF_HOME. See training.md for DCP conversion and Hugging Face safetensors export.

Troubleshooting

PyTorch Import Issue

Errors:

• ImportError: cannot import name '_functionalization' from 'torch._C'

An inherited NGC library directory can conflict with the environment's PyTorch. Inspect LD_LIBRARY_PATH and remove the conflicting entries. If the host driver already supports the selected CUDA build and no compatibility directory is needed, test a single process with an empty path:

LD_LIBRARY_PATH= python -c 'import torch; torch.cuda.init(); print(torch.cuda.get_device_name())'

Otherwise retain the platform's cuda-compat directory while removing conflicting cuDNN/NCCL/PyTorch directories. Do not add an unconditional path reset to your shell startup or Dockerfile. Re-run the GPU check after any change.

If this doesn't fix the issue, try reinstalling venv.

Dependency Issue

Errors:

• ModuleNotFoundError: No module named <module_name>

Reinstall venv:

uv sync --all-extras --group=cu130-train --reinstall
source .venv/bin/activate

If this doesn't fix the issue, try reinstalling uv.

Python Issue

Errors:

• fatal error: Python.h: No such file or directory

Reinstall uv and venv:

curl -LsSf https://astral.sh/uv/0.11.29/install.sh | sh
uv python install --reinstall
rm -rf .venv
uv sync --all-extras --group=cu130-train --reinstall
source .venv/bin/activate

CUDA Issue

For TorchCodec's CUDA builds, NPP is required even for CPU video decoding. The cu128 and cu130 groups install the matching NPP wheel. Framework's common video reader and inference conditioning decoder load that wheel before importing TorchCodec, so an empty LD_LIBRARY_PATH does not require a system NPP install. Standalone TorchCodec callers can use the same preparation:

from cosmos_framework.utils.generator.torchcodec_dependencies import load_torchcodec

torchcodec = load_torchcodec()

Install FFmpeg shared libraries that TorchCodec can load, as well as the ffmpeg executable.

• OSError: <lib_name>: cannot open shared object file: No such file or directory

Check which library is missing before installing a toolkit. Missing FFmpeg, OpenGL or GLib libraries require the system packages above. CUDA initialization errors require checking the host driver and CUDA variant, including any compatibility-library path. If a system CUDA toolkit library is actually missing, install the toolkit matching the environment's CUDA build:

sudo apt-get install -y --no-install-recommends "cuda-toolkit-${CUDA_TOOLKIT_VERSION:?set a version such as 13-0}"

Alternatively, use the Docker container.

TorchCodec Cannot Find CUDA 12 NPP

With cu128 or cu128-train, import torchcodec can fail with libnppicc.so.12: cannot open shared object file when the host only exposes another CUDA major version. The environment already installs nvidia-npp-cu12, but its wheel directory is not necessarily on the system linker's search path. Do not substitute a CUDA 13 library for a CUDA 12 library.

For a session-local workaround, use the matching libraries from the active environment. This also works when a launch script clears LD_LIBRARY_PATH:

source .venv/bin/activate
COSMOS_NPP_LIBDIR="$(python -c 'from importlib.metadata import distribution; print(distribution("nvidia-npp-cu12").locate_file("nvidia/npp/lib"))')"
export LD_PRELOAD="$COSMOS_NPP_LIBDIR/libnppc.so.12:$COSMOS_NPP_LIBDIR/libnppicc.so.12"
LD_LIBRARY_PATH= python -c 'import torch; import torchcodec; print(torchcodec.__version__)'

Apply this only to the CUDA 12 environment that needs it. After the run, unset LD_PRELOAD and COSMOS_NPP_LIBDIR (or restore their previous values if set). This does not install or replace system CUDA libraries. A matching CUDA 12 container or toolkit, as described above, is the alternative.