fix(nft): fix deterministic mode and add the SD3.5 DiffusionNFT e2e

[Rockdu]

What

• Seed the global RNG in RolloutManager.__init__, so DiffusionNFT's timestep shuffle reproduces across runs.
• Run scripts/run-diffusion-nft-sd3-pickscore.sh under --deterministic-mode, matching the other recorded recipes.
• Add tests/e2e/short/test_sd3_pickscore_nft_2xGPU.py to the existing stage-c-3-gpu-h200 suite (2 train GPUs colocated with 2 sglang rollout engines + 1 PickScore GPU, NUM_ROLLOUT=2), plus its recorded standard.

Why

The SD3.5 DiffusionNFT recipe was not bitwise reproducible, so no e2e standard could be recorded for it. Two runs of 3 rollouts under --deterministic-mode agreed on only 7 of 29 metric series.

miles/ray/data_conversion_hub/nft.py:79 shuffles each sample's timesteps with torch.randperm(...) and no generator, i.e. from the global RNG. That code runs inside the @ray.remote class RolloutManager, which never seeded torch — only the FSDP train actor does (actor.py:72). So the permutation came from default entropy and differed every run. --deterministic-mode cannot cover this: _enable_deterministic_training only sets the cuDNN/use_deterministic_algorithms knobs (it does no seeding at all) and is only invoked in the train actor.

Seeding is preferable to passing --no-diffusion-nft-shuffle-timesteps: the RNG still advances per rollout, so shuffling keeps varying across rollouts as intended while reproducing across runs.

This does not invalidate the existing standards

The obvious risk with seeding a shared actor is that already-recorded e2e standards start drifting and need re-recording. They do not, on three independent grounds:

1. nft.py:79 is the only unguarded global-torch-RNG draw in the whole package — grep -rE "torch\.(rand|randn|randint|randperm|normal)\(" miles/ | grep -v "generator=" returns that one line. step_strategy_hub.py uses explicit local generators (np.random.default_rng(seed), torch.Generator().manual_seed(...)), and the remaining random/np.random call sites (data order, port picking, dp_rank) are untouched by torch.manual_seed.
2. That line is gated on loss_type == "nft" (arguments.py:1476), so Flow-GRPO recipes never execute it, and flow_grpo.py itself makes no RNG draws at all.
3. Empirically: tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py was run on this branch and passed bit-for-bit against its committed standard — [e2e-metrics] PASSED: 8 metric series match test_sd3_ocr_grpo_2xGPU.json.

Two things for the reviewer to note

• The fix changes the NFT training trajectory — the permutation is now seed-determined rather than random. Nothing is invalidated (no NFT standard existed before this PR), but earlier ad-hoc NFT numbers are no longer comparable.
• train/grad_norm is deliberately absent from the metric list. fp16 forward enables ShardedGradScaler, whose initial scale (65536) overflows on NFT's much larger loss (~15 vs Flow-GRPO's tiny policy loss), so step 1 logs nan, the optimizer step is skipped and the scale halves. Since the verifier compares with got == want and nan == nan is False, including it would fail on every run — and a tolerance would not help, as math.isclose(nan, nan) is also False. The Flow-GRPO standards have clean floats at every step, which is why they can include it.

DiffusionNFT is a dual-policy x0-MSE objective with no importance ratio, so the Flow-GRPO log_prob_* series are never emitted by nft_loss_formula and cannot be used here. nft_pos_loss vs nft_neg_loss is the NFT analogue of the log-prob drift check: it measures how far the trained policy has moved from the EMA reference.

Validation

Determinism. h200, 3 GPUs, NUM_ROLLOUT=3, --deterministic-mode, full metrics JSONL compared with exact float equality (perf/* wall-clock keys excluded).

	Identical series	Mismatched
Before the fix	7 / 29	22
After the fix	29 / 29	0

The pre-fix divergence signature, which the fix removes:

Metric	run A	run B
rollout/reward/raw_mean @0	0.74586021900177	0.74586021900177
rollout/reward/raw_mean @1	0.769538402557373	0.769538402557373
rollout/reward/raw_mean @2	0.7332153916358948	0.7405926585197449
train/nft_loss @1	14.920654296875	14.985009765625
train/nft_t_mean @1	0.6582821259275079	0.6582821285352111

nft_t_mean differing at ~1e-9 with nft_num_timesteps identical is the fingerprint of reordering rather than different values. Rollouts 0 and 1 stayed bit-identical because step 1's update is skipped by the fp16 scaler, so the weights were unchanged going into rollout 1; divergence only surfaced at rollout 2.

Determinism on the CI runner. The standard was recorded three times on ["h200", "3gpu"] — the same runner class the test verifies on — and all three agree on all 12 series, including across the final code edit (2b9b6f8 → 773c686).

Curve health. Same recipe at NUM_ROLLOUT=100:

Metric	mean(first 20)	mean(last 20)	delta
rollout/reward/raw_mean	0.79100	0.84665	+0.0557

Reward buckets rise then plateau (0.770 0.812 0.826 0.819 0.835 0.839 0.837 0.847 0.850 0.844); the median rises 0.767 → 0.850. The largest single-rollout drop is -0.0714 at idx 32 (nft_loss spiking 17.2 → 23.4) and it recovers on the very next rollout to 0.828 / 16.9; all five worst drops are single-rollout and recover, with no sustained decline. The nft_pos_loss/nft_neg_loss gap grows from ~1e-5 to ~5e-5 over the run — a trend rather than flat, but 0.014% of the branch loss, smooth, and expected as the policy drifts from the lagging EMA reference (--ema-decay 0.001).

Cost. 3 rollouts end-to-end in 7m27s, of which ~5.5 min is engine/model load. est_time=600 for NUM_ROLLOUT=2, taking stage-c-3-gpu-h200 to 1800s total against the 60-minute job timeout.

Files

File	Role
miles/ray/rollout.py	Seed the global RNG in RolloutManager.__init__ (the fix).
scripts/run-diffusion-nft-sd3-pickscore.sh	Add --deterministic-mode.
tests/e2e/short/test_sd3_pickscore_nft_2xGPU.py	New e2e registration for stage-c-3-gpu-h200.
tests/ci/fixtures/e2e_standards/test_sd3_pickscore_nft_2xGPU.json	Recorded standard, 12 series. Only this file is added; the two existing standards are untouched.
tests/fast/backends/fsdp_utils/test_loss_hub_nft.py	CPU test that the seeded shuffle reproduces; already registered to stage-a-cpu.

Checklist

• pre-commit run --all-files passes — run on the touched files only (pre-commit run --files ...), all hooks pass
• Added/updated tests for new behaviour — CPU reproducibility test + the GPU e2e and its standard
• pytest -x is green — scope: tests/fast/backends/fsdp_utils/test_loss_hub_nft.py, 7 passed; the full fast suite was not run locally
• If launch flags changed, python3 train.py --help still parses — no new flags; the recipe uses the existing --deterministic-mode
• If a public flag was added, it appears in the CLI reference docs — n/a, no new flag
• If an example was added, it has a real walkthrough — n/a, existing recipe modified, not a new example