Spaces:
Running on L40S
Running on L40S
File size: 16,158 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 | # Inference
> **Skill:** `.agents/skills/cosmos3-inference/SKILL.md`
<!--TOC-->
______________________________________________________________________
**Table of Contents**
- [Quick Start](#quick-start)
- [Single-GPU](#single-gpu)
- [Multi-GPU](#multi-gpu)
- [Cosmos3-Nano](#cosmos3-nano)
- [Cosmos3-Super](#cosmos3-super)
- [Models](#models)
- [Modes](#modes)
- [Parallelism Arguments](#parallelism-arguments)
- [Sample Arguments](#sample-arguments)
- [Text](#text)
- [Vision (Image/Video)](#vision-imagevideo)
- [Action](#action)
- [Custom Defaults](#custom-defaults)
- [Guardrails](#guardrails)
- [Troubleshooting](#troubleshooting)
- [Checkpoint Issue](#checkpoint-issue)
- [Torch CUDA Out of Memory Error](#torch-cuda-out-of-memory-error)
- [NCCL Issue](#nccl-issue)
- [NCCL Plugin Issue](#nccl-plugin-issue)
______________________________________________________________________
<!--TOC-->
Prerequisites:
- [Setup](../README.md#setup)
- [Environment Variables](./environment_variables.md)
- [FAQ](./faq.md) — troubleshooting (OOM, NCCL hangs), defaults, common pitfalls.
Arguments:
- `-i`, `--input-files`: Path to the sample argument file(s) (JSON, JSONL, YAML). Accepts quoted glob patterns (e.g. `"inputs/*.json"`).
- `-o`, `--output-dir`: Output directory.
Outputs:
- `<sample_name>/`
- `sample_args.json`: Sample arguments.
- `sample_outputs.json`: Generation status, action (if enabled).
- `vision.jpg`, `vision.mp4`: Vision output (if enabled).
To see all available arguments:
```shell
python -m cosmos_framework.scripts.inference --help
```
## Quick Start
### Single-GPU
Use `python -m` directly. Suitable for `--parallelism-preset=latency` on a single GPU, or for quick experimentation:
```shell
python -m cosmos_framework.scripts.inference \
--parallelism-preset=latency \
-i "inputs/omni/t2v.json" \
-o outputs/omni_nano \
--checkpoint-path Cosmos3-Nano \
--seed=0
```
**Note:** Cosmos3-Super (32B) does not fit on a single 80 GB H100 — see [Cosmos3-Super](#cosmos3-super) for the multi-GPU recipes.
### Multi-GPU
Use `torchrun --nproc-per-node=N` when launching across multiple GPUs (N > 1). By default the model weights are sharded (FSDP) across all N GPUs, so any model fits. The `throughput` preset runs that single sharded replica over a batch; the `latency` preset additionally needs `--dp-shard-size=1` on multiple GPUs so the ranks are free for context parallelism (see [Parallelism Arguments](#parallelism-arguments)).
#### Cosmos3-Nano
```shell
torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \
--parallelism-preset=throughput \
-i "inputs/omni/*.json" \
-o outputs/omni_nano \
--checkpoint-path Cosmos3-Nano \
--seed=0
```
**Note:** The progress bar only prints on rank 0.
**Note:** With the default full-GPU sharding, this same command also runs Cosmos3-Super (32B) on 8×80 GB H100 — the weights are sharded (FSDP) across all 8 GPUs. See [Cosmos3-Super](#cosmos3-super) for the explicit-axis variants and the 4-GPU recipe.
#### Cosmos3-Super
Cosmos3-Super (32B) must be sharded across multiple GPUs to fit in 80 GB H100 memory. The default already shards the model across every visible GPU (FSDP), so the `throughput` preset fits it directly; the commands below pin the axes explicitly (pure FSDP, no context- or CFG-parallelism overlay) and add the 4-GPU recipe.
**4 GPUs:**
```shell
torchrun --nproc-per-node=4 -m cosmos_framework.scripts.inference \
--parallelism-preset=throughput \
--dp-shard-size=4 --dp-replicate-size=1 \
--cp-size=1 --cfgp-size=1 \
-i "inputs/omni/*.json" \
-o outputs/omni_super \
--checkpoint-path Cosmos3-Super \
--seed=0
```
**8 GPUs:**
```shell
torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \
--parallelism-preset=throughput \
--dp-shard-size=8 --dp-replicate-size=1 \
--cp-size=1 --cfgp-size=1 \
-i "inputs/omni/*.json" \
-o outputs/omni_super \
--checkpoint-path Cosmos3-Super \
--seed=0
```
The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `--parallelism-preset`. Super supports `text2image`, `text2video`, and `image2video` (see [Models](#models)).
## Models
| Model | Arguments | Modes |
| ------------- | --------------------------------- | ---------------------------------------------- |
| Cosmos3-Nano | `--checkpoint-path=Cosmos3-Nano` | All |
| Cosmos3-Super | `--checkpoint-path=Cosmos3-Super` | `text2image`, `text2video`, `image2video` |
## Modes
`model_mode` selects the generation modality. The table below lists every supported mode with its required sample fields and a paired example file.
| `model_mode` | Inputs | Outputs | Required sample fields | Example |
| ------------------ | ------------------------------------------ | ------------------------------------------------------------------------------------------ | ------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
| `text2image` | text prompt | `vision.jpg` | `prompt` | [`inputs/omni/t2i.json`](../inputs/omni/t2i.json) |
| `text2video` | text prompt | `vision.mp4` | `prompt` | [`inputs/omni/t2v.json`](../inputs/omni/t2v.json) |
| `image2video` | text prompt + image | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/i2v.json`](../inputs/omni/i2v.json) |
| `video2video` | text prompt + video | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/v2v.json`](../inputs/omni/v2v.json) |
| `forward_dynamics` | observation image/video + prompt + actions | future visual rollout in `vision.mp4` | `domain_name`, `vision_path`, `action_path` | [`inputs/omni/action_forward_dynamics_av.json`](../inputs/omni/action_forward_dynamics_av.json), [`inputs/omni/action_forward_dynamics_camera.json`](../inputs/omni/action_forward_dynamics_camera.json), [`inputs/omni/action_forward_dynamics_robot.json`](../inputs/omni/action_forward_dynamics_robot.json), [`inputs/omni/action_forward_dynamics_batch.jsonl`](../inputs/omni/action_forward_dynamics_batch.jsonl) |
| `inverse_dynamics` | observation video + prompt | predicted action sequence in `sample_outputs.json` | `domain_name`, `vision_path` | [`inputs/omni/action_inverse_dynamics_av.json`](../inputs/omni/action_inverse_dynamics_av.json), [`inputs/omni/action_inverse_dynamics_robot.json`](../inputs/omni/action_inverse_dynamics_robot.json), [`inputs/omni/action_inverse_dynamics_batch.jsonl`](../inputs/omni/action_inverse_dynamics_batch.jsonl) |
| `policy` | observation image/video + prompt | predicted action sequence in `sample_outputs.json` + future visual rollout in `vision.mp4` | `domain_name`, `vision_path` | [`inputs/omni/action_policy_av.json`](../inputs/omni/action_policy_av.json), [`inputs/omni/action_policy_robot.json`](../inputs/omni/action_policy_robot.json), [`inputs/omni/action_policy_batch.jsonl`](../inputs/omni/action_policy_batch.jsonl) |
Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To run every example in one batch, use `-i "inputs/omni/*.json"`.
## Parallelism Arguments
By default the model weights are sharded (FSDP) across **all** visible GPUs (`dp_shard_size = WORLD_SIZE`, `dp_replicate_size = 1`), so any model fits regardless of size. Override any axis with the `--dp-shard-size` / `--dp-replicate-size` / `--cp-size` / `--cfgp-size` flags.
- `--parallelism-preset`
- `latency`: Minimize wall-clock per sample by splitting each sample across GPUs with **context parallelism**. On multiple GPUs, also pass `--dp-shard-size=1` so the ranks are used for context/CFG parallelism instead of weight sharding. Used for real-time jobs.
- `throughput`: No context parallelism (`cp=cfgp=1`); the model is sharded across all GPUs and a single replica processes the batch. Used for batch jobs.
- `--dp-shard-size`: Number of ranks the model is sharded over (FSDP). Defaults to all ranks (`WORLD_SIZE`).
- `--max-num-seqs`: Maximum number of samples batched together per replica.
## Sample Arguments
Sample arguments are read from multiple sources (in priority order):
- CLI overrides (e.g. `--model-mode=text2video`): Overrides for all samples.
- Input files (e.g. `--input-files "inputs/omni/*t2i*.json"`): Single sample per input.
- Defaults: `cosmos_framework/inference/defaults/<model_mode>`: Defaults for all samples.
For debugging, the full set of sample arguments is saved to `<output_dir>/<sample_name>/sample_args.json`.
Common arguments:
- `model_mode`: Generation modality. See [Modes](#modes) above for all options.
- `seed`: Random seed for reproducibility.
**Note:** Condition file paths are relative to the input file.
### Text
- `prompt`: Inline text prompt.
### Vision (Image/Video)
Common arguments:
- `fps`: Condition and output frames per second.
- `resolution` (`"256"`, `"480"`, `"720"`): Condition and output resolution (height in pixels).
- `aspect_ratio` (`1,1`, `4,3`, `"3,4`, `16,9`, `9,16`): Condition and output aspect ratio. Defaults to `16,9`.
Condition arguments:
- `vision_path`: Path to an image or video file (local path or URL).
Generation arguments:
- `num_frames`: Number of output frames. `1` = image; `≥24` = video. Default 189; resolution-dependent max — see [FAQ § How many frames can I generate?](./faq.md#q-how-many-frames-can-i-generate).
Outputs `vision.jpg` or `vision.mp4` depending on `num_frames`.
### Action
Common arguments:
- `action_chunk_size`: Number of action steps in the chunk. The action media loader reads at most `action_chunk_size + 1` observation frames.
- `domain_name`: Domain name passed to the action domain registry, such as `bridge_orig_lerobot`, `camera_pose`, or `av`.
- `view_point`: Viewpoint description injected into the action prompt, such as `ego_view`.
Condition arguments:
- `action_path`: JSON action sequence. Required for `forward_dynamics`; each row is one action step and each column is one raw action dimension.
- `image_size`: Action input resize bucket. The value is passed as the action media resolution bucket; examples use `256` for LIBERO and `480` for AV.
The action output is written to `sample_outputs.json`.
See the [Modes](#modes) table above for the action mode inputs/outputs and example files.
### Custom Defaults
To use your own default values instead of the built-in presets, pass a JSON file via the `defaults_file` field in your sample arguments:
```json
{
"defaults_file": "my_defaults.json",
"prompt": "..."
}
```
The custom defaults file has the same format as the built-in presets. Fields you set explicitly in the sample argument file still take precedence over the custom defaults file.
## Guardrails
Inference ships with guardrails enabled by default, sourced from [nvidia/Cosmos-Guardrail1](https://huggingface.co/nvidia/Cosmos-Guardrail1). Active filters: text blocklist (better-profanity + fuzzy match), text safety classifier ([Qwen/Qwen3Guard-Gen-0.6B](https://huggingface.co/Qwen/Qwen3Guard-Gen-0.6B)), video content-safety classifier, and RetinaFace face-blur post-processor. Pass `--no-guardrails` to disable, or `--offload-guardrail-models` to keep them on CPU between calls (saves GPU memory, adds latency).
## Troubleshooting
### Checkpoint Issue
If you encounter failures downloading checkpoints, refer to [Downloading Base Checkpoints](./setup.md#downloading-base-checkpoints).
Checkpoint download commands are printed to the console. You can run them manually to debug issues.
### Torch CUDA Out of Memory Error
Error: `torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate X MiB`
[Optimize memory allocation](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf):
```shell
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
```
If that's not enough, see [FAQ § OOM during inference](./faq.md#q-i-get-torchcudaoutofmemoryerror-during-inference) for the full ladder (`--dp-shard-size`, `--device-memory-utilization`, `--offload-guardrail-models`).
### NCCL Issue
Error:
```shell
[rank0]:[W415 18:57:09.249883195 ProcessGroupNCCL.cpp:5138] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
Fatal Python error: Segmentation fault
```
Re-run with debugging enabled:
```shell
export NCCL_DEBUG=INFO
export TORCH_DISTRIBUTED_DEBUG=DETAIL
export CUDA_LAUNCH_BLOCKING=1
```
#### NCCL Plugin Issue
Error:
```shell
NCCL INFO Failed to initialize NET plugin Libfabric
Fatal Python error: Segmentation fault
```
Fix:
```shell
export NCCL_NET_PLUGIN=none
```
|