Skip to content

Commit 2f603cb

Browse files
pengcuoclaudelfengad
authored
Add cu130-torch213 dependency group (torch 2.13 on CUDA 13.0) (#134)
## What Adds `cu130-torch213` and `cu130-torch213-train` dependency groups pinning **torch 2.13.0+cu130** (with matching torchvision 0.28.0, torchcodec 0.14.0, torchao 0.17.0, triton 3.7.1, `natten 0.21.6+cu130.torch213` on aarch64, and the cu13 NVIDIA runtime libs), plus a `conflicts` entry making it mutually exclusive with the other backend groups. `uv.lock` is regenerated for the new closure. ## Why the code change torchvision 0.28 removed `torchvision.io.read_video`, so `cosmos_framework/inference/vision.py` switches to a torchcodec-based THWC uint8 decoder. It's forced onto CPU so an active default-CUDA context during generation doesn't route torchcodec's internal frame-index tensor to CUDA (which would raise `NotImplementedError`). ## Docs Bumps the recommended NGC base image to `nvcr.io/nvidia/pytorch:26.06-py3` in `README.md` and `docs/setup.md` to match the torch 2.13 line. ## Verification On 4× GB200 (aarch64): - `uv sync --all-extras --group=cu130-torch213-train` resolves clean; `torch.__version__ == 2.13.0+cu130`, `torch.version.cuda == 13.0`. - Cosmos3-Nano `throughput` inference over **all 13** `inputs/omni/*.json` samples (t2i, t2v, i2v, v2v, t2vs, i2vs, and the 6 action modes) produced valid outputs: videos decode with real pixel variation (std 44–76, well above the degenerate floor), audio modes carry an audio track, and all action arrays are non-empty and finite. ### Known limitations (in code comments) - x86_64 has no torch2.13 wheels yet for flash-attn / TE / natten — those are commented out / left on the torch2.10 build; this line is validated on **aarch64**. - torch 2.13.0+cu130 hard-pins cuDNN to 9.20.0.48; the cuDNN attention backend wants ≥9.22 (else it falls back to NATTEN). Override at runtime if needed. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: lfengad <liangf@nvidia.com>
1 parent 9853713 commit 2f603cb

6 files changed

Lines changed: 3348 additions & 1929 deletions

File tree

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ uv sync --all-extras --group=cu130-train
5757
source .venv/bin/activate && export LD_LIBRARY_PATH=
5858
```
5959

60-
If you are starting from the recommended NGC image (`nvcr.io/nvidia/pytorch:25.09-py3`), see the [one-shot quickstart](./docs/setup.md#quickstart-from-the-recommended-base-image).
60+
If you are starting from the recommended NGC image (`nvcr.io/nvidia/pytorch:26.06-py3`), see the [one-shot quickstart](./docs/setup.md#quickstart-from-the-recommended-base-image).
6161

6262
## Training
6363

cosmos_framework/inference/inference.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
import cattrs.preconf.json
1515
import safetensors.torch
1616
import torch
17-
import torchvision.io
1817
from PIL import Image
1918
from qwen_vl_utils.vision_process import smart_nframes
2019
from torch.utils._pytree import tree_map_only
@@ -46,6 +45,7 @@
4645
from cosmos_framework.inference.vision import (
4746
build_conditioned_video_batch,
4847
build_image_edit_batch,
48+
decode_video_thwc_uint8,
4949
load_conditioning_image,
5050
load_conditioning_video,
5151
load_prompt_upsampling_image,
@@ -574,17 +574,16 @@ def _get_prompt_sample_data(sample_args: OmniSampleArgs, model: OmniMoTModel, *,
574574
def _decode_reasoner_video(vision_path: str, video_fps: float | None) -> dict[str, Any]:
575575
"""Decode a local video file into the frame-list payload the Qwen3-VL processor expects.
576576
577-
Returns ``{"frames": [PIL.Image, ...], "fps": float}``. Uses the same
578-
``torchvision.io.read_video`` decode the rest of the inference path relies on
579-
(no ``decord`` dependency), then uniformly samples frames toward ``video_fps``
580-
(default 2.0) via Qwen's ``smart_nframes``. The repo ``Qwen3VLProcessor`` runs
581-
with ``do_sample_frames=False``, so it consumes this pre-sampled frame list
582-
as-is and handles its own per-frame resize."""
583-
frames, _, info = torchvision.io.read_video(str(vision_path), pts_unit="sec") # [T,H,W,C] uint8
577+
Returns ``{"frames": [PIL.Image, ...], "fps": float}``. Uses the same TorchCodec
578+
decode the rest of the inference path relies on (no ``decord`` dependency), then
579+
uniformly samples frames toward ``video_fps`` (default 2.0) via Qwen's
580+
``smart_nframes``. The repo ``Qwen3VLProcessor`` runs with ``do_sample_frames=False``,
581+
so it consumes this pre-sampled frame list as-is and handles its own per-frame resize."""
582+
frames, src_fps = decode_video_thwc_uint8(Path(vision_path)) # [T,H,W,C] uint8
584583
total_frames = int(frames.shape[0])
585584
if total_frames == 0:
586585
raise ValueError(f"Decoded zero frames from reasoner video: {vision_path}")
587-
src_fps = float(info.get("video_fps") or 0.0) or 1.0
586+
src_fps = src_fps or 1.0
588587
target_fps = video_fps if video_fps is not None else 2.0
589588
nframes = smart_nframes({"fps": target_fps}, total_frames=total_frames, video_fps=src_fps)
590589
idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()

cosmos_framework/inference/vision.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
import numpy as np
99
import torch
10-
import torchvision.io
1110
import torchvision.transforms.functional as TF
1211
from PIL import Image
1312

@@ -76,6 +75,26 @@ def load_conditioning_image(image_path: Path, target_h: int, target_w: int) -> t
7675
return img_tensor.unsqueeze(1) # [3,1,target_h,target_w]
7776

7877

78+
def decode_video_thwc_uint8(path: Path) -> tuple[torch.Tensor, float]:
79+
"""Read all frames of a video as a uint8 [T, H, W, C] tensor plus fps.
80+
81+
TorchCodec replacement for ``torchvision.io.read_video``; mirrors its frame layout so
82+
callers are unchanged. As of torchvision 0.28 the public wheel no longer binds that
83+
symbol — ``torchvision/io/__init__.py`` imports it from the Meta-internal
84+
``pytorch.vision.fb.io.video`` under ``except ImportError: pass``, so ``import
85+
torchvision.io`` still succeeds but calling ``read_video`` raises ``AttributeError``.
86+
"""
87+
from cosmos_framework.utils.generator.torchcodec_video import decode_frames_nhwc_uint8, probe_video
88+
89+
# torchcodec's core ops are CPU-only. Force CPU tensor creation so an active default-CUDA
90+
# device context (torch.set_default_device during generation) doesn't route torchcodec's
91+
# internal frame-index tensor to CUDA and raise NotImplementedError.
92+
with torch.device("cpu"):
93+
num_frames = probe_video(path).num_frames
94+
frames_nhwc, meta = decode_frames_nhwc_uint8(path, list(range(num_frames)))
95+
return torch.from_numpy(frames_nhwc), float(meta.average_fps)
96+
97+
7998
def load_conditioning_video(
8099
video_path: Path,
81100
target_h: int,
@@ -88,7 +107,7 @@ def load_conditioning_video(
88107
89108
``keep`` selects which ``max_frames`` to take when the input is longer.
90109
"""
91-
frames, _, _ = torchvision.io.read_video(str(video_path), pts_unit="sec")
110+
frames, _ = decode_video_thwc_uint8(video_path)
92111
frames = frames[-max_frames:] if keep == "last" else frames[:max_frames] # [T,H,W,3]
93112
frames_tchw = frames.permute(0, 3, 1, 2).float() # [T,3,H,W]
94113
frames_resized = _resize_and_center_crop(frames_tchw, target_h, target_w) # [T,3,target_h,target_w]
@@ -185,9 +204,8 @@ def read_media_frames(path: Path, max_frames: int) -> tuple[torch.Tensor, float]
185204
return frames, 1.0
186205
if ext not in _VIDEO_EXTENSIONS:
187206
raise ValueError(f"Unsupported media extension: {ext}")
188-
frames, _, info = torchvision.io.read_video(str(path), pts_unit="sec")
207+
frames, fps = decode_video_thwc_uint8(path)
189208
frames = frames[:max_frames].permute(0, 3, 1, 2).permute(1, 0, 2, 3)
190-
fps = float(info.get("video_fps", 24.0))
191209
return frames, fps
192210

193211

docs/setup.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ ______________________________________________________________________
4646
For CUDA 13 builds, the [NVIDIA NGC PyTorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) 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.
4747

4848
```dockerfile
49-
FROM nvcr.io/nvidia/pytorch:25.09-py3
49+
FROM nvcr.io/nvidia/pytorch:26.06-py3
5050
```
5151

5252
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.
@@ -70,7 +70,7 @@ The two supported install paths are the recommended base image and the Docker co
7070

7171
### Quickstart: From the Recommended Base Image
7272

73-
If you started from the [recommended base image](#recommended-base-image) (`nvcr.io/nvidia/pytorch:25.09-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/` directory you just cloned):
73+
If you started from the [recommended base image](#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/` directory you just cloned):
7474

7575
```shell
7676
apt-get update

pyproject.toml

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,46 @@ cu130-train = [
233233
"transformer-engine==2.12.0+cu130.torch210",
234234
"triton==3.6.0",
235235
]
236+
# torch 2.13 on CUDA 13.0. Same layout as cu130 but pinned to torch 2.13.
237+
cu130-torch213 = [
238+
#"flash-attn-3-nv==1.0.3+cu130.torch210; platform_machine == 'x86_64'", # TODO: torch213 build (cu130.torch210 wheels are cp313-only + torch210 ABI)
239+
#"flash-attn==2.7.4.post1+cu130.torch210; platform_machine == 'x86_64'", # TODO: torch213 build (cu130.torch210 wheels are cp313-only + torch210 ABI)
240+
# torch2.13 build (natten 0.21.6, sm 80/86/89/90/100/103/110/120/121) — matches this group's torch/ABI.
241+
# Both arches now ship a torch213 wheel, so no platform split is needed.
242+
"natten==0.21.6+cu130.torch213",
243+
"torch==2.13.0+cu130",
244+
"torchcodec==0.14.0+cu130", # https://github.com/meta-pytorch/torchcodec/releases
245+
"torchvision==0.28.0+cu130", # TODO: release matching torch 2.13 — https://github.com/pytorch/vision/releases
246+
# Dependencies determined from 'uv pip compile --group cu130-torch213' — these are the exact
247+
# ==-pins that torch==2.13.0+cu130 (download.pytorch.org) declares. NOTE: this is the CUDA 13.0/13.1
248+
# line the public wheel ships, NOT the CUDA 13.3 libs in nvcr.io/nvidia/pytorch:26.06.
249+
# Issue: https://github.com/astral-sh/uv/issues/14237
250+
"nvidia-cublas==13.1.1.3",
251+
"nvidia-cuda-cupti==13.0.85",
252+
"nvidia-cuda-nvrtc==13.0.88",
253+
"nvidia-cuda-runtime==13.0.96",
254+
# NOTE: torch 2.13.0+cu130 hard-pins this to ==9.20.0.48. The framework's cuDNN attention backend
255+
# needs >= 9.22 (else it falls back to NATTEN, which is fine). To use the cuDNN backend, override at
256+
# runtime: `uv pip install nvidia-cudnn-cu13==9.25.0.15` (can't pin here without a global override).
257+
"nvidia-cudnn-cu13==9.20.0.48",
258+
"nvidia-cufft==12.0.0.61",
259+
"nvidia-cufile==1.15.1.6",
260+
"nvidia-curand==10.4.0.35",
261+
"nvidia-cusolver==12.0.4.66",
262+
"nvidia-cusparse==12.6.3.3",
263+
"nvidia-cusparselt-cu13==0.8.1",
264+
"nvidia-nccl-cu13==2.29.7",
265+
"nvidia-npp==13.0.0.50",
266+
"nvidia-nvjitlink==13.3.33",
267+
"nvidia-nvshmem-cu13==3.4.5",
268+
"nvidia-nvtx==13.0.85",
269+
]
270+
cu130-torch213-train = [
271+
{include-group = "cu130-torch213"},
272+
"torchao==0.17.0+cu130; platform_machine == 'x86_64'", # 0.18.0+cu130 not published; no aarch64 build; https://github.com/pytorch/ao/issues/2919
273+
"transformer-engine==2.12.0+cu130.torch210", # TODO: torch213 build
274+
"triton==3.7.1", # torch==2.13.0+cu130 pins triton==3.7.1
275+
]
236276
# LIBERO simulator dependencies for the closed-loop eval client.
237277
# Mirrors packages/cosmos-policy
238278
# `libero` group; `libero` from PyPI declares `robosuite` transitively so we
@@ -274,10 +314,13 @@ repository = "https://github.com/NVIDIA/cosmos-framework"
274314
required-version = ">=0.11.3"
275315
conflicts = [
276316
[
317+
{group = "vllm"},
277318
{group = "cu128"},
278319
{group = "cu128-train"},
279320
{group = "cu130"},
280321
{group = "cu130-train"},
322+
{group = "cu130-torch213"},
323+
{group = "cu130-torch213-train"},
281324
],
282325
]
283326
override-dependencies = [

0 commit comments

Comments
 (0)