From 14f6e2237ae67d15633bcc1c3664b6b4ec903a69 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 12:35:31 -0700 Subject: [PATCH 1/3] fix(ci): match a recorded NaN against itself in the e2e metric comparison --- tests/ci/e2e_metrics_registry.py | 5 +++++ tests/fast/ci/test_e2e_metrics_registry.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/tests/ci/e2e_metrics_registry.py b/tests/ci/e2e_metrics_registry.py index effd0f31..8c2fe1d9 100644 --- a/tests/ci/e2e_metrics_registry.py +++ b/tests/ci/e2e_metrics_registry.py @@ -105,6 +105,11 @@ def load_series(jsonl_path: str | Path, metrics: list[str]) -> dict[str, list[li def _values_match(got: float, want: float, tol: dict | None) -> bool: + # A recorded NaN is a value like any other -- fp16 runs emit grad_norm=nan on the step the + # grad scaler overflows its init scale, deterministically. Neither == nor isclose matches + # NaN against itself, so without this a standard containing one could never pass. + if math.isnan(got) or math.isnan(want): + return math.isnan(got) and math.isnan(want) if tol is None: return got == want return math.isclose(got, want, rel_tol=tol.get("rtol", 0.0), abs_tol=tol.get("atol", 0.0)) diff --git a/tests/fast/ci/test_e2e_metrics_registry.py b/tests/fast/ci/test_e2e_metrics_registry.py index 7418e880..932c3615 100644 --- a/tests/fast/ci/test_e2e_metrics_registry.py +++ b/tests/fast/ci/test_e2e_metrics_registry.py @@ -58,6 +58,24 @@ def test_strict_mismatch_fails_but_tolerance_passes(sandbox, monkeypatch): reg.check_or_update("test_foo.py", drifted, ["m"], tolerances={"m": {"atol": 1e-3}}) +def test_recorded_nan_matches_itself_but_not_a_number(sandbox, monkeypatch): + # fp16 runs emit grad_norm=nan on the step the scaler overflows its init scale. + nan_run = sandbox / "nan.jsonl" + finite_run = sandbox / "finite.jsonl" + _write_jsonl(nan_run, [{"step": 1, "m": float("nan")}, {"step": 2, "m": 0.25}]) + _write_jsonl(finite_run, [{"step": 1, "m": 0.5}, {"step": 2, "m": 0.25}]) + monkeypatch.setenv("MILES_E2E_METRICS_UPDATE", "1") + reg.check_or_update("test_foo.py", nan_run, ["m"]) + monkeypatch.delenv("MILES_E2E_METRICS_UPDATE") + reg.check_or_update("test_foo.py", nan_run, ["m"]) + # A run that stops overflowing is a real change, not a match. + with pytest.raises(AssertionError, match="strict"): + reg.check_or_update("test_foo.py", finite_run, ["m"]) + # Nor does a tolerance let a number pass against a recorded NaN. + with pytest.raises(AssertionError, match="tol"): + reg.check_or_update("test_foo.py", finite_run, ["m"], tolerances={"m": {"atol": 1e9}}) + + def test_series_shape_mismatches_fail(sandbox, monkeypatch): std = sandbox / "std.jsonl" _write_jsonl(std, [{"step": 1, "m": 0.5}, {"step": 2, "m": 0.25}]) From 00eb32246d1bc8182383f4c02376a1303cbb6f0e Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 14:43:41 -0700 Subject: [PATCH 2/3] test(e2e): add the DiffusionNFT pickscore e2e on stage-c-3-gpu-h200 --- scripts/run_diffusion_nft_sd3_pickscore.py | 1 + .../e2e/short/test_sd3_nft_pickscore_3xGPU.py | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+) create mode 100644 tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py diff --git a/scripts/run_diffusion_nft_sd3_pickscore.py b/scripts/run_diffusion_nft_sd3_pickscore.py index fa97dca4..c4947aa4 100644 --- a/scripts/run_diffusion_nft_sd3_pickscore.py +++ b/scripts/run_diffusion_nft_sd3_pickscore.py @@ -130,6 +130,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--rollout-num-gpus-per-engine 1 " f"--num-gpus-per-node {2 if args.smoke else 3} " "--colocate " + "--deterministic-mode " ) U.execute_train( diff --git a/tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py b/tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py new file mode 100644 index 00000000..c8a3d382 --- /dev/null +++ b/tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py @@ -0,0 +1,37 @@ +"""E2E: SD3.5-medium DiffusionNFT with PickScore, 3-GPU (2 colocated FSDP DP=2 + sglang +rollout engines, 1 dedicated reward GPU) — runs the example script's real configuration, +not a reduced one, and checks its metric series against the registered standard +(tests/ci/fixtures/e2e_standards/). Runs with --deterministic-mode, so every metric is +compared strictly, bit for bit. + +Only --num-rollout is cut down, 100 -> 4, which is the shortest run that still reaches the +behaviour worth guarding. Step 1 overflows the fp16 grad scaler's 65536 init scale and is +skipped, so the weights do not move; step 2 is the first that lands, so through it the EMA +reference is still identical to the policy and NFT's two loss branches are algebraically +one value. They separate at step 3. Two rollouts would exercise none of that. + +The NFT-side metrics are what the GRPO e2e cannot cover: NFT has no log-prob ratio, its +loss is the two-branch x0-MSE, and nft_t_mean tracks the sigma grid the prepare hook feeds +the DiT. +""" + +from tests.ci.e2e_metrics_registry import register_e2e_ci + +register_e2e_ci( + est_time=900, + suite="stage-c-3-gpu-h200", + script="scripts/run_diffusion_nft_sd3_pickscore.py", + args=["--num-rollout", "4", "--cuda-visible-devices", "0,1,2"], + metrics=[ + "rollout/reward/raw_num_samples", + "rollout/reward/raw_mean", + "rollout/reward/raw_median", + "rollout/reward/raw_std", + "train/grad_norm", + "train/nft_loss", + "train/nft_pos_loss", + "train/nft_neg_loss", + "train/nft_r_mean", + "train/nft_t_mean", + ], +) From ed90e75c31be4e03e945e9faf6fbafd9096a2dd6 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 12 Aug 2026 02:03:23 -0700 Subject: [PATCH 3/3] test(e2e): register the DiffusionNFT standard --- .../test_sd3_nft_pickscore_3xGPU.json | 188 ++++++++++++++++++ 1 file changed, 188 insertions(+) create mode 100644 tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json diff --git a/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json b/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json new file mode 100644 index 00000000..2e2038f1 --- /dev/null +++ b/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json @@ -0,0 +1,188 @@ +{ + "meta": { + "commit": "00eb32246d1bc8182383f4c02376a1303cbb6f0e", + "source": "test_sd3_nft_pickscore_3xGPU.py" + }, + "metrics": { + "rollout/reward/raw_mean": [ + [ + 0, + 0.74586021900177 + ], + [ + 1, + 0.769538402557373 + ], + [ + 2, + 0.748466968536377 + ], + [ + 3, + 0.7546428442001343 + ] + ], + "rollout/reward/raw_median": [ + [ + 0, + 0.7397080659866333 + ], + [ + 1, + 0.7668601274490356 + ], + [ + 2, + 0.744907796382904 + ], + [ + 3, + 0.7552561163902283 + ] + ], + "rollout/reward/raw_num_samples": [ + [ + 0, + 64.0 + ], + [ + 1, + 64.0 + ], + [ + 2, + 64.0 + ], + [ + 3, + 64.0 + ] + ], + "rollout/reward/raw_std": [ + [ + 0, + 0.04408378154039383 + ], + [ + 1, + 0.06347957998514175 + ], + [ + 2, + 0.03868388384580612 + ], + [ + 3, + 0.04985775798559189 + ] + ], + "train/grad_norm": [ + [ + 1.0, + NaN + ], + [ + 2.0, + 0.024831360206007957 + ], + [ + 3.0, + 0.027672801166772842 + ], + [ + 4.0, + 0.06675068289041519 + ] + ], + "train/nft_loss": [ + [ + 1.0, + 15.153754340277779 + ], + [ + 2.0, + 16.450520833333332 + ], + [ + 3.0, + 17.195638020833332 + ], + [ + 4.0, + 19.235731336805557 + ] + ], + "train/nft_neg_loss": [ + [ + 1.0, + 0.3367631700303819 + ], + [ + 2.0, + 0.365570068359375 + ], + [ + 3.0, + 0.3821360270182292 + ], + [ + 4.0, + 0.4274359809027778 + ] + ], + "train/nft_pos_loss": [ + [ + 1.0, + 0.3367631700303819 + ], + [ + 2.0, + 0.365570068359375 + ], + [ + 3.0, + 0.38216400146484375 + ], + [ + 4.0, + 0.42747921413845485 + ] + ], + "train/nft_r_mean": [ + [ + 1.0, + 0.5000000074505806 + ], + [ + 2.0, + 0.5000000053809749 + ], + [ + 3.0, + 0.5000000260770321 + ], + [ + 4.0, + 0.4999999875823657 + ] + ], + "train/nft_t_mean": [ + [ + 1.0, + 0.7304325223796897 + ], + [ + 2.0, + 0.7304325221727291 + ], + [ + 3.0, + 0.7304325236214532 + ], + [ + 4.0, + 0.7304325207240052 + ] + ] + } +}