Skip to content

fix(nft): seed the sigma shuffle and the corruption noise so NFT runs reproduce - #156

Merged
Rockdu merged 1 commit into
fix/nft-sigma-input-hookfrom
fix/nft-determinism
Aug 12, 2026
Merged

fix(nft): seed the sigma shuffle and the corruption noise so NFT runs reproduce#156
Rockdu merged 1 commit into
fix/nft-sigma-input-hookfrom
fix/nft-determinism

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Second of three. Base: #155.

What

Seed the two random streams DiffusionNFT draws from, the way SFT already does:

  • data_conversion_hub/nft.py — the per-sample sigma permutation
  • loss_hub/nft.py — the corruption noise in prepare_nft_batch
  • miles/utils/hash_utils.pystable_hash, lifted out of loss_hub/sft.py's private _seed so the converter (under miles/ray/) and the loss hub (under miles/backends/) can share it without importing across that boundary

Nothing in miles upstream covers this: its hash helpers are all private and purpose-built (tensor checksums in local_weight_checksum.py, disk_delta.py), so there was nothing to reuse. hash_utils rather than a seed-specific name because the thing being solved is general — hash() is randomised per process by PYTHONHASHSEED, so anything derived from it stops agreeing across ranks and runs. Seeding is the first caller, not the only possible one.

Why

Both streams ran off the global RNG:

sample_sigmas = sigmas[torch.randperm(num_timesteps)]   # data_conversion_hub/nft.py
xt = corrupt(x0, t, sample_noise(x0))                   # loss_hub/nft.py

sample_noise already takes a generator — the plumbing was there and unused.

So NFT training does not reproduce, and --deterministic-mode does not make it reproduce. Two recordings of the same configuration on the same pinned stack (same image, same sglang commit, same miles-d commit, only NUM_ROLLOUT differing) agree bit-for-bit on the rollout side and diverge on the training side from the very first optimizer step:

metric step run A run B
rollout/reward/raw_mean 0 0.74586021900177 0.74586021900177
rollout/reward/raw_mean 1 0.769538402557373 0.769538402557373
train/nft_loss 1 15.090966796875 14.9947021484375
train/nft_t_mean 1 0.6582821235060692 0.6582821292802692

Identical rollouts, 0.6% apart in loss. nft_t_mean differing confirms the sigma assignment itself differed; the rollout agreeing rules out the generation path.

This was found while trying to record an e2e standard for NFT (#157) — the standard could never have held.

Seeding scheme

Same as SFT, and byte-identical to it: stable_hash produces exactly the integers _seed did, so SFT's streams are untouched. Keyed on what makes a stream distinct rather than on draw order, so a stream reproduces even when unrelated code changes how much randomness it consumes first.

stream key
sigma permutation run seed, sample.index
corruption noise run seed, rollout, microbatch, dp rank

sample.index is the data source's global counter (data_source.py:85), which advances across rollouts and is checkpointed, so each sample draws its own permutation and no two rollouts repeat one.

Behaviour

NFT training numbers change. The noise and permutation sequences are different from the unseeded ones — this replaces an arbitrary stream with a fixed one, it does not correct a wrong value. No NFT e2e standard exists yet to re-record. SFT is untouched: stable_hash returns exactly what _seed returned for the same arguments (verified), and its two call sites only change the name they import.

Tests

TestNftDeterminism — the shuffle reproduces across calls, each sample gets its own permutation, the corruption noise reproduces, and a different microbatch still draws different noise. Reverting either seeding line fails exactly one of the first two.

Stack

  1. fix(nft): give NFT its own timestep hook instead of sharing process_timestep_as_input #155 — the inverted sigma rescaling
  2. this PR
  3. test(e2e): add the DiffusionNFT pickscore e2e on stage-c-3-gpu-h200 #157 — the DiffusionNFT e2e, which needs both

Checklist

  • pre-commit run passes — on the touched files (black/ruff/isort/autoflake)
  • Added/updated tests for new behaviour
  • pytest is green — tests/fast 199 passed vs 189 on clean main (the +10 are new); the 4 failures there reproduce on main and are macOS c10d socket issues ⚠️ run locally with sglang stubbed out; CI is the real check.
  • 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

@Rockdu
Rockdu force-pushed the fix/nft-determinism branch from d732e4e to b8b5ac8 Compare August 12, 2026 07:44
@Rockdu
Rockdu merged commit b09c96f into fix/nft-sigma-input-hook Aug 12, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant