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

[Rockdu]

Third of three. Base: #156.

What

• tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py on stage-c-3-gpu-h200 — the runner already had a spare third GPU the SD3 GRPO e2e does not use.
• scripts/run_diffusion_nft_sd3_pickscore.py — --deterministic-mode, matching the three other e2e-backed recipes.
• tests/ci/e2e_metrics_registry.py — a recorded NaN now matches itself.
• The recorded standard.

It runs the recipe's real configuration

PickScore reward, flowgrpo_pickscore, 8 prompts × 8 samples = 64 per rollout, micro-batch 4, 2 colocated train/rollout GPUs + 1 dedicated reward GPU. The only knob cut is --num-rollout, 100 → 4.

An earlier attempt used the recipe's --smoke path and produced nothing worth comparing: 4 samples, all scoring 0 on OCR, zero reward variance, zero advantage, zero gradient. That is what --smoke is for — checking the pipeline runs — and not what an e2e should verify.

The recorded standard

step	grad_norm	nft_pos_loss − nft_neg_loss	reward mean / std
1	nan	0	0.74586 / 0.04408
2	0.024831360206007957	0	0.76954 / 0.06348
3	0.027672801166772842	+2.797e−05	0.74847 / 0.03868
4	0.066750682890415190	+4.323e−05	0.75464 / 0.04986

Why 4 rollouts. --num-steps-per-rollout 1 makes one rollout one optimizer step, and the fp16 grad scaler starts at 65536 against an fp16 ceiling of 65504, so step 1 overflows and is skipped. Until a step lands the weights do not move, the EMA reference stays identical to the policy, and NFT's two loss branches are algebraically one value — which is exactly what steps 1 and 2 show. They separate at step 3. Two rollouts would exercise none of the dual-branch loss, the thing this test exists to guard. Three is the minimum; four leaves a step of margin.

Eight was the first attempt and cost 5 minutes of runner time per PR more than it needed: fixed cost is 283s (image, engines, SD3.5 + PickScore load, first rollout), each further rollout 74s.

The standard is reproducible, and that took #156

Recorded twice, independently, on this branch's 00eb322. The two runs are bit-identical across all 10 metric series × 4 steps.

That is the end-to-end proof for #156. Before its seeding fix, two recordings of the same configuration on an identically-pinned stack diverged at the first training step:

	before #156	after #156
train/nft_loss @1	15.090966796875 vs 14.9947021484375	identical
train/nft_t_mean @1	0.6582821235060692 vs 0.6582821292802692	identical
all 10 × 4	training side diverged throughout	all identical

The rollout side agreed even then — it was the train-pair construction that drew from an unseeded RNG. Finding that is why this PR sat unrecordable until #156 existed.

A recorded NaN could never match itself

_values_match compared with ==, or math.isclose under a tolerance. Neither matches NaN against NaN, so any standard containing one was permanently red.

Not hypothetical: an fp16 run emits grad_norm=nan on the step the grad scaler overflows its init scale, deterministically, and actor.py emits the metric between unscale_ and step — so the value the scaler is about to discard is the one that gets logged. Dropping grad_norm from the metric list instead would have hidden a real property of fp16 NFT training.

Now a NaN matches a NaN and nothing else: a run that stops overflowing still fails, and no tolerance lets a number pass against a recorded NaN. Both directions are asserted.

Metrics

The ones the GRPO e2e structurally 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 — the path #155 changes.

train/nft_r_mean is in the list and is exactly 0.5 at every step, which is what a mean-zero GRPO advantage forces. It guards little; dropping it needs no re-record (the check iterates the test's list, not the standard's), so say the word.

Stack

1. fix(nft): give NFT its own timestep hook instead of sharing process_timestep_as_input #155 — the inverted sigma rescaling
2. fix(nft): seed the sigma shuffle and the corruption noise so NFT runs reproduce #156 — seed NFT's random streams
3. this PR

Checklist

• pre-commit run passes — on the touched files (black/ruff/isort/autoflake)
• Added/updated tests for new behaviour
• Test discovery verified — run_suite.py --hw cuda --suite stage-c-3-gpu-h200 --list-only lists it
• E2E standard recorded on the verifying runner class (["h200","3gpu"]), and reproduced bit-for-bit on a second independent recording
• If launch flags changed, python3 train.py --help still parses — no flag changes
• If a public flag was added, it appears in the CLI reference docs — no new flags
• If an example was added, it has a real walkthrough — no new examples

The recipe trains, not just starts

100 rollouts of the unmodified recipe (no --num-rollout cut, no --smoke) on a 4×H200 devbox, from this branch's code — wandb run. Job succeeded, 100/100 steps, no errors.

every 10th rollout	1	11	21	31	41	51	61	71	81	91
rollout/reward/raw_mean	0.7459	0.8249	0.8900	0.8533	0.8062	0.8418	0.8112	0.8426	0.8428	0.8507
rollout/reward/raw_std	0.0441	0.0628	0.0487	0.0365	0.0822	0.0515	0.0617	0.0404	0.0518	0.0580
train/nft_loss	15.15	19.51	17.31	18.45	20.78	18.58	20.59	20.30	19.37	19.29
train/grad_norm	nan	0.0307	0.0230	0.0430	0.0152	0.0189	0.0361	0.0316	0.0235	0.0229

PickScore reward rises 0.77321 → 0.84485 over the run, +9.3% (first-10 vs last-10 mean). Reward std stays in 0.04–0.08 throughout, so the signal is real spread and not a collapsed batch.

Two things in this long run confirm what the 4-step standard asserts:

• grad_norm is nan at step 1 only — the grad scaler overflows its 65536 init scale once, backs off, and never again. The standard records that same single NaN, which is why _values_match had to learn to match one.
• nft_pos_loss == nft_neg_loss exactly at steps 1, 2 and 5, and differs everywhere else (max |diff| 1.5e−04). Steps 1–2 are the scaler warmup the standard captures; step 5 is a later skipped step. This is the dual-branch behaviour the e2e guards, and it confirms a 2-rollout test would have seen none of it.