test(inference): smoke-test FP8 mixed-precision diffusion steps (#226)

## What

Adds `test_nano_fp8_mixed_precision_inference` to
`tests/nano_inference_smoke_test.py`: the same sharded-layout text2video
run as `test_nano_fp8_inference`, plus the FP8 mixed-precision
diffusion-step flags from #217 (`--mixed-precision-first-steps=2
--mixed-precision-last-steps=2`, default `w8a16-cache=none` — the only
mode valid under FSDP sharding, which is exactly the layout used).

## Pass criterion

The schedule itself, not just a valid video: the test parses the
`MIXED_PRECISION_TRACE` log line and compares it exactly against the
expected `2x W8A16 / 6x W8A8 / 2x W8A16` sequence for the 10-step run
(`num_steps` is read back from the shared `_FP8_GENERATION_ARGS` so the
expectation cannot drift). A run where the flags silently never engaged
(all-W8A8 trace or no trace) fails even though its video would look
fine. `_assert_video_has_content` still catches the
numerically-broken-but-running case, and the existing FP8 swap-count
guard confirms the FP8 path engaged at all.

## CI wiring

No new job or trigger: the existing `generator-inference-smoke` job runs
the whole file, so the case is picked up automatically. Like the other
two FP8 cases it SKIPs (not fails) when `HF_TOKEN` cannot read
`nvidia/Cosmos3-Experimental`. Job timeout bumped 90 → 105 min for the
extra full-width run.

## Verification

`pytest --collect-only` inside the inference container collects 5 tests
including the new case; the GPU run itself is exercised by this PR's own
`generator-inference-smoke` job.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: pengcuo

.github/workflows/gpu-tests.yml

@@ -125,15 +125,22 @@ jobs:
   generator-inference-smoke:
     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 full-width runs.
-    timeout-minutes: 90
+    # 105 (not 60) since the Nano step gained the three ModelOpt FP8 cases (two
+    # parallelism layouts + mixed-precision diffusion steps): a 20 GB checkpoint
+    # on a cold cache plus three more full-width runs.
+    timeout-minutes: 105
     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"
+      # This is the job whose stated purpose includes FP8 coverage: promote the
+      # cannot-reach-nvidia/Cosmos3-Experimental skip to a hard failure, so a
+      # token rotation or a permissions change turns the job red instead of
+      # silently dropping all three FP8 cases. Forks and local runs (which do
+      # not set this) keep the skip behaviour.
+      REQUIRE_FP8: "1"
     steps:
       - uses: actions/checkout@v6
 
@@ -144,10 +151,12 @@ jobs:
 
       # One inference call over t2vs (+sound), action policy, and forward_dynamics; checks each output.
       # Plus the ModelOpt static-FP8 Nano checkpoint in both parallelism layouts
-      # (FSDP-sharded and replicated). The FP8 checkpoint lives in the
-      # 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.
+      # (FSDP-sharded and replicated) and once more with FP8 mixed-precision
+      # diffusion steps (W8A16 edge steps; exact schedule asserted from the
+      # MIXED_PRECISION_TRACE log). The FP8 checkpoint lives in the
+      # access-controlled nvidia/Cosmos3-Experimental repo; elsewhere those three
+      # cases SKIP unless HF_TOKEN can read it, but REQUIRE_FP8=1 above makes the
+      # skip a hard failure here, so a green run is guaranteed to have covered FP8.
       # -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, 4 GPU)

tests/nano_inference_smoke_test.py

@@ -37,6 +37,13 @@
    layout — the FP8 all-gather. Skipped when the checkpoint is not reachable
    (see ``_download_fp8_checkpoint``).
 
+4. One more ``text2video`` call against the same FP8 checkpoint with FP8
+   mixed-precision diffusion steps enabled (``--mixed-precision-first-steps`` /
+   ``--mixed-precision-last-steps``): the first/last N denoising steps run
+   W8A16 (dequantized weight + dense GEMM) while middle steps keep the TorchAO
+   W8A8 path. Asserts the exact per-step precision schedule from the
+   ``MIXED_PRECISION_TRACE`` log line plus a non-degenerate ``vision.mp4``.
+
 Smoke-level only (output validity, not numeric goldens). The checkpoint + its
 tokenizers download from the HF Hub on first run and are reused afterward.
 
@@ -216,6 +223,22 @@
     ),
 }
 
+# Mixed-precision diffusion steps (FP8 W8A16 edge steps) schedule for the
+# dedicated smoke case: with the 10-step ``_FP8_GENERATION_ARGS`` run this
+# selects 2x W8A16 / 6x W8A8 / 2x W8A16 (``use_w8a16_step``). Small enough to
+# stay cheap, large enough that first, middle, and last regions are all
+# non-empty.
+_MIXED_PRECISION_FIRST_STEPS = 2
+_MIXED_PRECISION_LAST_STEPS = 2
+
+# Emitted by ``MixedPrecisionRuntime`` (``cosmos_framework/utils/generator/
+# mixed_precision.py``): the install summary at load time and the per-request
+# per-step precision trace at request end. The trace is parsed and compared
+# exactly — a run that silently ignored the flags would log an all-W8A8 trace
+# (or none at all) and still produce a perfectly valid video.
+_MIXED_PRECISION_INSTALL_LOG = "Mixed precision installed:"
+_MIXED_PRECISION_TRACE_LOG = re.compile(r"MIXED_PRECISION_TRACE steps=([A-Za-z0-9,]+)")
+
 # Audio sanity thresholds for the muxed sound track.
 _RMS_SILENCE_FLOOR = 1e-4  # below this the track is effectively silence
 _PEAK_SANITY_CEIL = 1.5    # decoded float audio should sit within ~[-1, 1]
@@ -354,22 +377,32 @@ def _download_fp8_checkpoint() -> Path:
     nvidia/Cosmos3-Experimental), so a fork PR without the runner secret does not
     go red. Any other failure (a deleted revision, a broken download) still fails
     loudly: a silently-skipping FP8 job would look green while testing nothing.
+
+    ``REQUIRE_FP8=1`` promotes that skip to a hard failure. The CI job whose
+    stated purpose includes FP8 coverage sets it, so a token rotation or a
+    permissions change on the gated repo turns that job red instead of silently
+    dropping every FP8 case while still reporting green.
     """
     from huggingface_hub import snapshot_download
     from huggingface_hub.errors import GatedRepoError, HfHubHTTPError, RepositoryNotFoundError
 
+    def _skip_or_fail_no_access(reason: str) -> None:
+        if os.environ.get("REQUIRE_FP8") == "1":
+            pytest.fail(f"REQUIRE_FP8=1 but the FP8 checkpoint is unreachable: {reason}")
+        pytest.skip(reason)
+
     try:
         repo_root = snapshot_download(
             repo_id=_FP8_REPOSITORY,
             revision=_FP8_REVISION,
             allow_patterns=[f"{_FP8_SUBDIRECTORY}/*"],
         )
     except (GatedRepoError, RepositoryNotFoundError) as error:
-        pytest.skip(f"no access to {_FP8_REPOSITORY} (needs an HF_TOKEN with read access): {error!r}")
+        _skip_or_fail_no_access(f"no access to {_FP8_REPOSITORY} (needs an HF_TOKEN with read access): {error!r}")
     except HfHubHTTPError as error:
         status_code = getattr(error.response, "status_code", None)
         if status_code in (401, 403):
-            pytest.skip(f"no access to {_FP8_REPOSITORY} (HTTP {status_code}): {error!r}")
+            _skip_or_fail_no_access(f"no access to {_FP8_REPOSITORY} (HTTP {status_code}): {error!r}")
         raise
 
     checkpoint_path = Path(repo_root) / _FP8_SUBDIRECTORY
@@ -627,3 +660,76 @@ def test_nano_fp8_inference(tmp_path: Path, layout: str) -> None:
         video = so.parent / "vision.mp4"
         assert video.is_file(), f"FP8 {layout} run produced no vision.mp4 ({so})"
         _assert_video_has_content(video)
+
+    @pytest.mark.level(2)
+    @pytest.mark.gpus(MAX_GPUS)
+    def test_nano_fp8_mixed_precision_inference(tmp_path: Path) -> None:
+        """text2video from the FP8 Nano checkpoint with mixed-precision diffusion steps.
+
+        Same run as the ``sharded`` case of ``test_nano_fp8_inference`` plus the
+        ``--mixed-precision-first-steps`` / ``--mixed-precision-last-steps`` flags,
+        so the first/last 2 of the 10 denoising steps run W8A16 (dequantized E4M3
+        weight + dense GEMM) while the middle 6 keep the TorchAO W8A8 path. The
+        sharded layout is the one that constrains the feature: FSDP-sharded FP8
+        weights support only the default ``mixed_precision_w8a16_cache='none'``
+        (per-step on-the-fly dequant), which is exactly the mode exercised here.
+
+        The pass criterion is the schedule itself, not just a valid video: the
+        ``MIXED_PRECISION_TRACE`` line is parsed and compared exactly against the
+        expected ``2x W8A16 / 6x W8A8 / 2x W8A16`` sequence, so a run where the
+        flags never engaged (all-W8A8 trace, or no trace at all) fails even though
+        its output video would look fine. ``_assert_video_has_content`` then
+        catches the numerically-broken-but-still-running case.
+        """
+        checkpoint_path = _download_fp8_checkpoint()
+        out_dir = tmp_path / "out_fp8_mixed_precision"
+        cmd = [
+            "torchrun",
+            f"--nproc_per_node={MAX_GPUS}",
+            f"--master_port={_free_port()}",
+            "-m",
+            "cosmos_framework.scripts.inference",
+            *_FP8_LAYOUTS["sharded"],
+            "-i",
+            "inputs/omni/t2v.json",
+            "-o",
+            str(out_dir),
+            "--checkpoint-path",
+            str(checkpoint_path),
+            *_FP8_GENERATION_ARGS,
+            f"--mixed-precision-first-steps={_MIXED_PRECISION_FIRST_STEPS}",
+            f"--mixed-precision-last-steps={_MIXED_PRECISION_LAST_STEPS}",
+        ]
+        log = _run(cmd, tmp_path / "inference_fp8_mixed_precision.log")
+
+        # The FP8 path itself still engaged (same guard as test_nano_fp8_inference).
+        swap_match = _FP8_SWAP_LOG.search(log)
+        assert swap_match is not None and int(swap_match.group(1)) > 0, (
+            "no ModelOpt FP8 linear swap in the mixed-precision run; FP8 path never engaged"
+        )
+        # ... and the mixed-precision runtime was installed on top of it.
+        assert _MIXED_PRECISION_INSTALL_LOG in log, (
+            "mixed precision was never installed despite --mixed-precision-first/last-steps"
+        )
+
+        # Exact per-step precision schedule. num_steps is read from
+        # _FP8_GENERATION_ARGS so the expectation cannot drift from the run.
+        (num_steps,) = [int(a.split("=")[1]) for a in _FP8_GENERATION_ARGS if a.startswith("--num-steps=")]
+        expected = (
+            ["W8A16"] * _MIXED_PRECISION_FIRST_STEPS
+            + ["W8A8"] * (num_steps - _MIXED_PRECISION_FIRST_STEPS - _MIXED_PRECISION_LAST_STEPS)
+            + ["W8A16"] * _MIXED_PRECISION_LAST_STEPS
+        )
+        traces = [m.split(",") for m in _MIXED_PRECISION_TRACE_LOG.findall(log)]
+        assert expected in traces, (
+            f"expected a MIXED_PRECISION_TRACE of {'/'.join(expected)}, got traces={traces}"
+        )
+
+        results = sorted(out_dir.rglob("sample_outputs.json"))
+        assert len(results) == 1, f"expected 1 mixed-precision sample_outputs.json, found {[str(p) for p in results]}"
+        so = results[0]
+        args = json.loads(so.read_text()).get("args", {})
+        assert args.get("model_mode") == "text2video", f"expected a text2video sample, got {args.get('model_mode')}"
+        video = so.parent / "vision.mp4"
+        assert video.is_file(), f"FP8 mixed-precision run produced no vision.mp4 ({so})"
+        _assert_video_has_content(video)