ci(gpu): run the whole GPU test workflow at 4 GPUs

.github/workflows/gpu-tests.yml

@@ -1,22 +1,29 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
-# Cosmos3-Nano GPU test suite on a self-hosted 8×H200 runner.
+# Cosmos3-Nano GPU test suite on a self-hosted H200 runner.
+#
+# Every job is 4-GPU: each sets ``TEST_MAX_GPUS: "4"``, which is what
+# ``cosmos_framework/inference/fixtures/args.py`` reads to build
+# ``ALL_NUM_GPUS = (0, 1, MAX_GPUS)`` -- the set of widths ``--num-gpus`` and the
+# ``gpus(N)`` markers may name. The env var and the ``--num-gpus`` flag must
+# agree, or collection either errors on an out-of-range choice or deselects
+# everything.
 #
 # A single ``pre-commit`` (lint) job runs first; the six GPU jobs all
 # ``needs:`` it, so they wait on ONE pre-commit run and are skipped if lint
 # fails — the single self-hosted runner is never spent on a lint-failing commit.
 # The six GPU jobs then run (one at a time on the single runner):
 #   * training-smoke                — Nano SFT pipeline (convert -> train 5 -> export -> t2i)
-#   * generator-training-regression — vision_sft_nano loss vs goldens (4-GPU subset)
+#   * generator-training-regression — vision_sft_nano loss vs goldens
 #   * generator-inference-smoke     — Nano + Edge multi-modality inference (t2vs/t2v + policy + forward_dynamics) + Nano ModelOpt FP8
-#   * distilled-inference-smoke     — Super 4-step T2I + I2V inference (4-GPU subset)
-#   * reasoner-inference-smoke      — Nano reasoner golden + Edge reasoner smoke + Qwen conversion coverage (4-GPU)
-#   * reasoner-training-regression  — llava_ov loss vs goldens (4-GPU subset)
+#   * distilled-inference-smoke     — Super 4-step T2I + I2V inference
+#   * reasoner-inference-smoke      — Nano reasoner golden + Edge reasoner smoke + Qwen conversion coverage
+#   * reasoner-training-regression  — llava_ov loss vs goldens
 #
 # Requires:
-#   * a self-hosted runner labelled [self-hosted, gpu, h200] with 8 GPUs,
-#     NVIDIA drivers, and `uv` on PATH;
+#   * a self-hosted runner labelled [self-hosted, gpu, h200] with at least 4
+#     GPUs, NVIDIA drivers, and `uv` on PATH;
 #   * an `HF_TOKEN` repository secret (gated dataset/model downloads, incl. the
 #     streamed LLaVA-OneVision-Data dataset). The Nano ModelOpt FP8 cases also
 #     need it to read nvidia/Cosmos3-Experimental, which is not public; they skip
@@ -49,6 +56,11 @@ jobs:
     env:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DISABLE_XET: "1"
+      # 4-GPU SFT run. vision_sft_nano_5iter.toml sets
+      # data_parallel_shard_degree = -1 (auto from WORLD_SIZE), so this is a
+      # valid FSDP layout; each rank holds twice the shard it did at 8, which is
+      # the memory headroom to watch if this job starts OOM-ing.
+      TEST_MAX_GPUS: "4"
     steps:
       - uses: actions/checkout@v6
 
@@ -59,12 +71,12 @@ jobs:
 
       # Full SFT pipeline: download + convert Nano->DCP, train 5 steps (loss
       # trend), export to HF safetensors, then a t2i generation from the export.
-      # MAX_GPUS defaults to 8. -s streams the live process log.
-      - name: Nano SFT pipeline smoke (convert -> train 5 -> export -> t2i, 8 GPU)
+      # -s streams the live process log.
+      - name: Nano SFT pipeline smoke (convert -> train 5 -> export -> t2i, 4 GPU)
         run: |
           export LD_LIBRARY_PATH=
           uv run --all-extras --group=cu128-train python -m pytest -v -s \
-            tests/nano_training_smoke_test.py --num-gpus=8 --levels=2 -o addopts=
+            tests/nano_training_smoke_test.py --num-gpus=4 --levels=2 -o addopts=
 
       # Clear the heavy artifacts (even on failure): examples/checkpoints (the
       # Cosmos3-Nano DCP + Wan VAE, ~30 GB) and the pytest tmp dirs (the SFT
@@ -83,7 +95,7 @@ jobs:
     env:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DISABLE_XET: "1"
-      # Select the 4-GPU regression test variant (uses 4 of the 8 GPUs).
+      # Select the 4-GPU regression test variant.
       TEST_MAX_GPUS: "4"
     steps:
       - uses: actions/checkout@v6
@@ -96,7 +108,7 @@ jobs:
       # Generator (vision_sft_nano) loss vs the h100 goldens, plus a Nano DCP ->
       # Diffusers conversion that validates the exported component files and indexes.
       # Both tests share the module-scoped staged Nano DCP. -s streams the live log.
-      - name: Generator regression and Nano Diffusers export (4-GPU subset)
+      - name: Generator regression and Nano Diffusers export (4 GPU)
         run: |
           export LD_LIBRARY_PATH=
           uv run --all-extras --group=cu128-train python -m pytest -v -s \
@@ -114,11 +126,14 @@ jobs:
     needs: pre-commit
     runs-on: [self-hosted, gpu, h200]
     # 90 (not 60) since the Nano step gained the two ModelOpt FP8 cases: a 20 GB
-    # checkpoint on a cold cache plus two more 8-GPU runs.
+    # checkpoint on a cold cache plus two more full-width runs.
     timeout-minutes: 90
     env:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DISABLE_XET: "1"
+      # Both steps use the `throughput` preset, whose layout is derived from
+      # world_size (cfgp=1, cp=1, dp_shard=world_size), so the width is free.
+      TEST_MAX_GPUS: "4"
     steps:
       - uses: actions/checkout@v6
 
@@ -133,26 +148,26 @@ jobs:
       # access-controlled nvidia/Cosmos3-Experimental repo, so those two cases SKIP
       # unless HF_TOKEN can read it — check the log for "no access to
       # nvidia/Cosmos3-Experimental" before trusting a green run to have covered FP8.
-      # MAX_GPUS defaults to 8. -s streams the live process log.
+      # -s streams the live process log.
       # Reuse the same input-asset cache dir as the unittest job.
-      - name: Nano inference smoke (t2vs + action policy + forward_dynamics + FP8, 8 GPU)
+      - name: Nano inference smoke (t2vs + action policy + forward_dynamics + FP8, 4 GPU)
         run: |
           export LD_LIBRARY_PATH=
           export COSMOS_DOWNLOAD_CACHE_DIR="$RUNNER_WORKSPACE/cosmos_input_cache"
           uv run --all-extras --group=cu128-train python -m pytest -v -s \
-            tests/nano_inference_smoke_test.py --num-gpus=8 --levels=2 -o addopts=
+            tests/nano_inference_smoke_test.py --num-gpus=4 --levels=2 -o addopts=
 
       # Same three-modality smoke on the smaller Cosmos3-Edge checkpoint: a plain
       # t2v (Edge has no sound, so t2vs is swapped for t2v and its 480/121/24
       # generation defaults are asserted) plus action policy + forward_dynamics.
       # Reuses the same input-asset cache dir; the Edge checkpoint downloads once
       # and is reused afterward.
-      - name: Edge inference smoke (t2v + action policy + forward_dynamics, 8 GPU)
+      - name: Edge inference smoke (t2v + action policy + forward_dynamics, 4 GPU)
         run: |
           export LD_LIBRARY_PATH=
           export COSMOS_DOWNLOAD_CACHE_DIR="$RUNNER_WORKSPACE/cosmos_input_cache"
           uv run --all-extras --group=cu128-train python -m pytest -v -s \
-            tests/edge_inference_smoke_test.py --num-gpus=8 --levels=2 -o addopts=
+            tests/edge_inference_smoke_test.py --num-gpus=4 --levels=2 -o addopts=
 
       # Inference writes only the pytest tmp dir (the t2vs video + logs); the
       # checkpoint download stays in the HF cache (kept). No examples/ artifacts.
@@ -202,7 +217,7 @@ jobs:
     env:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DISABLE_XET: "1"
-      # 4-GPU reasoner inference test (uses 4 of the 8 GPUs); the gpus(4) marker
+      # 4-GPU reasoner inference test; the gpus(4) marker
       # requires MAX_GPUS == 4.
       TEST_MAX_GPUS: "4"
     steps:
@@ -265,7 +280,7 @@ jobs:
     env:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DISABLE_XET: "1"
-      # Select the 4-GPU regression test variant (uses 4 of the 8 GPUs).
+      # Select the 4-GPU regression test variant.
       TEST_MAX_GPUS: "4"
     steps:
       - uses: actions/checkout@v6
@@ -276,7 +291,7 @@ jobs:
         run: uv sync --all-extras --group=cu128-train
 
       # Reasoner (llava_ov) loss vs the h100 goldens. -s streams the live log.
-      - name: Reasoner regression (llava_ov, 4-GPU subset)
+      - name: Reasoner regression (llava_ov, 4 GPU)
         run: |
           export LD_LIBRARY_PATH=
           uv run --all-extras --group=cu128-train python -m pytest -v -s \
@@ -304,6 +319,11 @@ jobs:
     env:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DISABLE_XET: "1"
+      # This job passes no --num-gpus, so it only collects the gpus(0)/gpus(1)
+      # cases either way. Pinning the width anyway keeps the whole workflow on
+      # one ALL_NUM_GPUS, so a gpus(8) marker added later fails collection here
+      # rather than passing lint and dying on the runner.
+      TEST_MAX_GPUS: "4"
     steps:
       - uses: actions/checkout@v6
 

tests/edge_inference_smoke_test.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
-"""8-GPU multi-modality inference smoke test for Cosmos3-Edge.
+"""Multi-modality inference smoke test for Cosmos3-Edge (``MAX_GPUS`` ranks).
 
 Mirrors ``nano_inference_smoke_test.py`` but targets the smaller Cosmos3-Edge
 checkpoint, adapted to its capabilities: Edge has ``sound_gen=false`` (so the
@@ -34,10 +34,15 @@
 goldens). The checkpoint + its tokenizers download from the HF Hub on first run
 and are reused afterward.
 
-Invocation (inside the inference container, from the repo root, on an 8-GPU
-node)::
+The run is ``MAX_GPUS`` ranks wide (``TEST_MAX_GPUS``, default 8). The
+``throughput`` preset derives its layout from ``world_size`` (cfgp=1, cp=1,
+dp_shard=world_size), so 4 and 8 are both valid; what changes is the FSDP shard
+width, not what the test covers.
+
+Invocation (inside the inference container, from the repo root)::
 
     pytest -s tests/edge_inference_smoke_test.py --num-gpus=8 --levels=2 -o addopts=
+    TEST_MAX_GPUS=4 pytest -s tests/edge_inference_smoke_test.py --num-gpus=4 --levels=2 -o addopts=
 
 Without ``--num-gpus``/``--levels`` (e.g. the no-GPU pre-commit CI) the test is
 not collected.
@@ -156,30 +161,31 @@ def _assert_valid_action(content: dict, where: str) -> None:
 
 
 @pytest.fixture(scope="module", autouse=True)
-def _require_8_gpus() -> None:
-    """Skip the module unless we can launch an 8-GPU run here."""
+def _require_gpus() -> None:
+    """Skip the module unless we can launch a ``MAX_GPUS``-wide run here."""
     if shutil.which("torchrun") is None:
         pytest.skip("torchrun not on PATH -- must run inside the inference container")
     try:
         import torch
     except Exception as exc:  # pragma: no cover -- surfaces during dev only
         pytest.skip(f"torch unavailable ({exc!r})")
-    if not torch.cuda.is_available() or torch.cuda.device_count() < 8:
-        pytest.skip(f"requires 8 visible CUDA devices, found {torch.cuda.device_count()}")
+    if not torch.cuda.is_available() or torch.cuda.device_count() < MAX_GPUS:
+        pytest.skip(f"requires {MAX_GPUS} visible CUDA devices, found {torch.cuda.device_count()}")
 
 
-# Defined only when the active MAX_GPUS is 8 -- the conftest rejects ``gpus(N)``
-# markers outside ``ALL_NUM_GPUS = (0, 1, MAX_GPUS)``.
-if MAX_GPUS == 8:
+# Markers use MAX_GPUS because the conftest rejects ``gpus(N)`` outside
+# ``ALL_NUM_GPUS = (0, 1, MAX_GPUS)``. Both supported widths are listed so the
+# module is collected under either; nothing here pins a parallelism degree.
+if MAX_GPUS in (4, 8):
 
     @pytest.mark.level(2)
-    @pytest.mark.gpus(8)
+    @pytest.mark.gpus(MAX_GPUS)
     def test_edge_inference_omni(tmp_path: Path) -> None:
         """Throughput run over t2v + policy + forward_dynamics on Cosmos3-Edge."""
         out_dir = tmp_path / "out"
         cmd = [
             "torchrun",
-            "--nproc_per_node=8",
+            f"--nproc_per_node={MAX_GPUS}",
             f"--master_port={_free_port()}",
             "-m",
             "cosmos_framework.scripts.inference",

tests/nano_inference_smoke_test.py

@@ -194,10 +194,15 @@
 # hooks, which is the path that does not exist upstream, and it is the only
 # layout a Super-class FP8 model fits in. ``replicated`` guards the
 # single-device-weights path that worked before those hooks were added.
+# Both layouts follow MAX_GPUS so the product of the parallel degrees equals
+# WORLD_SIZE, which is what ParallelDims validates. `sharded` stays pure FSDP and
+# `replicated` pure context-parallel -- the distinction each case exists to
+# cover -- at whichever width the run has. cp_size is bounded by MAX_CP_SIZE=32,
+# so 4 and 8 are both in range.
 _FP8_LAYOUTS = {
     "sharded": (
         "--parallelism-preset=throughput",
-        "--dp-shard-size=8",
+        f"--dp-shard-size={MAX_GPUS}",
         "--dp-replicate-size=1",
         "--cp-size=1",
         "--cfgp-size=1",
@@ -206,7 +211,7 @@
         "--parallelism-preset=latency",
         "--dp-shard-size=1",
         "--dp-replicate-size=1",
-        "--cp-size=8",
+        f"--cp-size={MAX_GPUS}",
         "--cfgp-size=1",
     ),
 }
@@ -389,10 +394,8 @@ def _require_gpus() -> None:
 
 
 # Markers use MAX_GPUS because the conftest rejects ``gpus(N)`` outside
-# ``ALL_NUM_GPUS = (0, 1, MAX_GPUS)``. Cases that genuinely need 8 ranks carry
-# their own skipif below: on a smaller host only the 4-rank case can run, and
-# forcing the others onto fewer ranks would change what they cover (ParallelDims
-# fails when the parallelism product does not equal world_size).
+# ``ALL_NUM_GPUS = (0, 1, MAX_GPUS)``. Every case here derives its parallelism
+# from the active width, so both supported widths are listed.
 if MAX_GPUS in (4, 8):
 
     @pytest.mark.level(2)
@@ -572,11 +575,6 @@ def test_nano_inference_multi_control_transfer(tmp_path: Path) -> None:
 
     @pytest.mark.level(2)
     @pytest.mark.gpus(MAX_GPUS)
-    # Unlike the cases above, this one cannot simply follow MAX_GPUS: the FP8
-    # setup pins cp_size=8 / dp_shard_size=8, which ParallelDims rejects against
-    # a smaller WORLD_SIZE. Running it below 8 needs a different parallelism
-    # layout, i.e. a different thing under test.
-    @pytest.mark.skipif(MAX_GPUS < 8, reason="FP8 setup pins 8-way cp/dp_shard")
     @pytest.mark.parametrize("layout", sorted(_FP8_LAYOUTS))
     def test_nano_fp8_inference(tmp_path: Path, layout: str) -> None:
         """text2video from the ModelOpt static-FP8 Nano checkpoint, once per layout.

tests/nano_training_smoke_test.py

@@ -1,10 +1,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
-"""8-GPU Cosmos3-Nano SFT pipeline smoke test (train -> export -> infer).
+"""Cosmos3-Nano SFT pipeline smoke test (train -> export -> infer), ``MAX_GPUS`` ranks.
 
 Runs the documented Vision SFT (Cosmos3-Nano) lifecycle from ``docs/training.md``
-end to end on 8 GPUs and validates each artifact:
+end to end on ``MAX_GPUS`` GPUs and validates each artifact:
 
   1. Step 1 -- download the bridge-v2 subset dataset + the Wan2.2 VAE.
   2. Step 2 -- ``convert_model_to_dcp`` Cosmos3-Nano -> DCP; check DCP completeness.
@@ -23,10 +23,17 @@
 ``examples/checkpoints/``, cached across runs); run output goes under the pytest
 tmp dir. Steps 1-2 are skipped when their artifacts already exist.
 
-Invocation (inside the training container, from the repo root, on an 8-GPU
-node)::
+The run is ``MAX_GPUS`` ranks wide (``TEST_MAX_GPUS``, default 8).
+``vision_sft_nano_5iter.toml`` sets ``data_parallel_shard_degree = -1`` (auto
+from ``WORLD_SIZE``), so a narrower run is a valid FSDP layout -- each rank just
+holds a proportionally larger shard, which is the thing to watch for OOM when
+lowering it. The loss check is a trend, not a golden, so the batch-size change
+that comes with it does not invalidate the assertion.
+
+Invocation (inside the training container, from the repo root)::
 
     pytest -s tests/nano_training_smoke_test.py --num-gpus=8 --levels=2 -o addopts=
+    TEST_MAX_GPUS=4 pytest -s tests/nano_training_smoke_test.py --num-gpus=4 --levels=2 -o addopts=
 
 Without ``--num-gpus``/``--levels`` (e.g. the no-GPU pre-commit CI) the test is
 not collected.
@@ -391,8 +398,8 @@ def _assert_valid_image(path: Path) -> None:
 
 
 @pytest.fixture(scope="module", autouse=True)
-def _require_8_gpus() -> None:
-    """Skip the module unless we can launch an 8-GPU training run here."""
+def _require_gpus() -> None:
+    """Skip the module unless we can launch a ``MAX_GPUS``-wide training run here."""
     if shutil.which("torchrun") is None:
         pytest.skip("torchrun not on PATH -- must run inside the training container")
     if shutil.which("uvx") is None:
@@ -401,14 +408,17 @@ def _require_8_gpus() -> None:
         import torch
     except Exception as exc:  # pragma: no cover
         pytest.skip(f"torch unavailable ({exc!r})")
-    if not torch.cuda.is_available() or torch.cuda.device_count() < 8:
-        pytest.skip(f"requires 8 visible CUDA devices, found {torch.cuda.device_count()}")
+    if not torch.cuda.is_available() or torch.cuda.device_count() < MAX_GPUS:
+        pytest.skip(f"requires {MAX_GPUS} visible CUDA devices, found {torch.cuda.device_count()}")
 
 
-if MAX_GPUS == 8:
+# Markers use MAX_GPUS because the conftest rejects ``gpus(N)`` outside
+# ``ALL_NUM_GPUS = (0, 1, MAX_GPUS)``. Both supported widths are listed so the
+# module is collected under either.
+if MAX_GPUS in (4, 8):
 
     @pytest.mark.level(2)
-    @pytest.mark.gpus(8)
+    @pytest.mark.gpus(MAX_GPUS)
     def test_nano_sft_train_export_infer(tmp_path: Path) -> None:
         """Full Cosmos3-Nano SFT pipeline: convert -> train 5 -> export -> t2i infer."""
         # 1-2. Inputs + HF->DCP convert, then DCP completeness.
@@ -424,7 +434,7 @@ def test_nano_sft_train_export_infer(tmp_path: Path) -> None:
             extra_env={
                 "MASTER_PORT": str(_free_port()),
                 "OUTPUT_ROOT": str(tmp_path / "launcher_out"),
-                "NPROC_PER_NODE": "8",
+                "NPROC_PER_NODE": str(MAX_GPUS),
             },
         )
         assert rc == 0, f"SFT launch failed (exit {rc}):\nLog tail:\n{out[-4000:]}"
@@ -490,7 +500,7 @@ def test_nano_sft_train_export_infer(tmp_path: Path) -> None:
         infer_out = tmp_path / "exported_out"
         rc, out = _run(
             [
-                "torchrun", "--nproc_per_node=8", f"--master_port={_free_port()}",
+                "torchrun", f"--nproc_per_node={MAX_GPUS}", f"--master_port={_free_port()}",
                 "-m", "cosmos_framework.scripts.inference",
                 "--parallelism-preset=throughput",
                 "-i", "inputs/omni/t2i.json",