Skip to content

test(e2e): add the DiffusionNFT pickscore e2e on stage-c-3-gpu-h200 #157

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Rockdu merged 3 commits into from
Aug 12, 2026
Merged

test(e2e): add the DiffusionNFT pickscore e2e on stage-c-3-gpu-h200 #157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
14f6e22
fix(ci): match a recorded NaN against itself in the e2e metric compar…
Rockdu Aug 11, 2026
00eb322
test(e2e): add the DiffusionNFT pickscore e2e on stage-c-3-gpu-h200
Rockdu Aug 11, 2026
ed90e75
test(e2e): register the DiffusionNFT standard
Rockdu Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions scripts/run_diffusion_nft_sd3_pickscore.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 5 additions & 0 deletions tests/ci/e2e_metrics_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
188 changes: 188 additions & 0 deletions tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
]
]
}
}
37 changes: 37 additions & 0 deletions tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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",
],
)
18 changes: 18 additions & 0 deletions tests/fast/ci/test_e2e_metrics_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}])
Expand Down
Loading