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

[lfengad]

What

Drops the GPU workflow from a mixed 8-and-4 layout to 4 GPUs everywhere.

Four of the seven GPU jobs already pinned TEST_MAX_GPUS: "4". Only
training-smoke and generator-inference-smoke ran at the default 8. Because
ALL_NUM_GPUS = (0, 1, MAX_GPUS) admits exactly one multi-GPU width per
process, that split forced each multi-GPU test file to guard its own
definitions behind if MAX_GPUS == 4 / == 8, and made the 8-wide files
uncollectable under a 4-GPU invocation.

Why this is safe

Neither 8-wide job needs 8 ranks:

• edge_inference_smoke / nano_inference_smoke run with
  --parallelism-preset=throughput, and _build_context_parallelism derives
  that layout from world_size (cfgp=1, cp=1, dp_shard=world_size). The width
  is free.
• nano_training_smoke drives vision_sft_nano_5iter.toml, whose
  data_parallel_shard_degree = -1 means "auto from WORLD_SIZE", and asserts a
  loss trend rather than a golden — so the batch-size change a narrower run
  brings does not invalidate it.

The hardcoded 8s become MAX_GPUS; the collection guards accept either width,
matching what nano_inference_smoke_test.py already did. --num-gpus and
TEST_MAX_GPUS must agree or collection either errors on an out-of-range
choice or deselects everything, so both move together.

unittest passes no --num-gpus and carries only gpus(0)/gpus(1) markers,
so it is unaffected either way. It is pinned anyway so the workflow has a single
ALL_NUM_GPUS and a gpus(8) marker added later fails collection instead of
reaching a runner that cannot host it.

The FP8 commit, separately

_FP8_LAYOUTS pinned dp_shard_size=8 / cp_size=8 literally, which is why
the two ModelOpt FP8 cases carried skipif(MAX_GPUS < 8). Left alone they would
have silently skipped. The second commit makes both degrees follow MAX_GPUS,
keeping sharded a pure-FSDP layout and replicated a pure-CP one — the
distinction the two cases exist to draw. This genuinely changes the
configuration under test (a 4-way rather than 8-way all-gather through the
TorchAO static-FP8 FSDP2 shim), so it is split out and can be reverted alone.

Verification

Run on 4×GB200 (TEST_MAX_GPUS=4 --num-gpus=4 --levels=2), not the H200 CI arch:

• Collection is exactly the intended set — 15 tests across all seven GPU test
  files, including test_nano_fp8_inference[sharded] and [replicated], which
  previously skipped.
• TEST_MAX_GPUS=8 --num-gpus=8 still parses with no marker errors (all
  deselected on a 4-GPU host, as designed).
• ruff check clean on the three touched files. ruff format deviations in
  those files are pre-existing and not gated — pre-commit configures no ruff hook.
• Execution results are in a comment below.

What this gives up

Coverage of the 8-wide FSDP/CP path itself: a bug that only appears at a larger
world size is no longer caught here.

🤖 Generated with Claude Code

[lfengad]

Execution results

Filling in the placeholder from the PR body. Run on the real CI runner (H200), not just the GB200 box I used while developing.

generator-inference-smoke — the job this PR actually changes

The one result worth stating explicitly. The job's own comment warns that the ModelOpt FP8 cases skip rather than fail when HF_TOKEN cannot read nvidia/Cosmos3-Experimental, so a green run does not by itself prove FP8 was covered — and FP8 is the entire subject of the second commit. Checked the log:

• 4 passed in 560.42s for nano_inference_smoke_test.py — omni, multi_control_transfer, and both fp8[sharded] and fp8[replicated]. 0 skipped.
• 1 passed in 67.09s for edge_inference_smoke_test.py.
• grep -c 'no access to' → 0. The access guard never fired, so the FP8 cases genuinely ran.

The 4-wide layouts really did engage, rather than silently degrading to bf16:

Loaded ModelOpt FP8 weights into TorchAO (504 sharded / 0 replicated)   <- sharded layout
Loaded ModelOpt FP8 weights into TorchAO (0 sharded / 504 replicated)   <- replicated layout
Installed static-FP8 FSDP support (8 aten ops + all-gather hooks)

504 linears swapped in each layout, and the two layouts are mirror images as intended — sharded exercises the TorchAO static-FP8 FSDP2 shim at a 4-way all-gather, replicated the pure-CP path at cp_size=4. The resolved config shows data_parallel_shard_degree: 4.

training-smoke — the memory question

Passed in 11m31s. This was the one risk called out in the PR body: dropping dp_shard from 8 to 4 means each rank holds twice the shard plus optimizer state. On 141 GB H200s there is headroom. Worth noting this does not generalize downward — a smaller-memory 4-GPU box would need re-checking.

Everything else

Job	Result
pre-commit (x2)	pass
training-smoke	pass, 11m31s
generator-training-regression	pass, 11m3s
generator-inference-smoke	pass, 10m42s
distilled-inference-smoke	pass, 2m6s
reasoner-inference-smoke	pass, 3m13s
reasoner-training-regression	pass, 2m19s
unittest	still running at time of writing

Offline checks (4xGB200, aarch64 — different arch from CI)

• Collection under TEST_MAX_GPUS=4 --num-gpus=4 --levels=2 across all seven GPU test files: exactly 15 tests, including the two FP8 cases that previously skipped.
• TEST_MAX_GPUS=8 --num-gpus=8 still parses with no marker errors (all deselected on a 4-GPU host, as designed).
• The unittest job's exact invocation collects identically with TEST_MAX_GPUS=4 and unset: 1230/1233 collected, 3 deselected, same 3 pre-existing container-dependency collection errors either way. Confirms that pinning the width there is a no-op.
• ruff check clean on the three touched files. The ruff format deviations in those files are pre-existing and sit in code this PR does not touch; pre-commit configures no ruff hook, so they are not gated.

Follow-up, not for this PR

_detect_arch in launch_regression_test.py maps only GB200/H100/H200 and pytest.skips anything else. Moving the suite to a different GPU arch would make both regression jobs skip silently and still report green. Not a problem for a same-model-type capacity change, but worth turning into a hard failure separately.