fix(nft): give NFT its own timestep hook instead of sharing process_timestep_as_input

[Rockdu]

What

• Add process_sigma_as_timesteps_input(sigmas, *, num_train_timesteps) to TrainPipelineConfig: the base scales up to the scheduler range, qwen_image passes the sigma through.
• prepare_nft_batch calls it instead of pre-multiplying into process_timestep_as_input.
• Drop the now-dead needs_timestep_scaling left on ltx.py, and the wan2_2.py comment that outlived the attribute it annotated.

Why

#125 folded three prepare hooks onto one process_timestep_as_input. NFT does not belong on it: its input is not a trajectory timestep.

NFT's pairs carry a sigma read straight off scheduler.sigmas (data_conversion_hub/nft.py:72), so σ ∈ [0, 1], while flow-GRPO and SFT pass a timestep in the scheduler range. The rescaling therefore runs the opposite direction, which is exactly what the old needs_timestep_scaling branch in nft.py encoded — inverted relative to its two siblings:

	flow_grpo / sft	nft
True (qwen_image)	t / N	t unchanged
False (sd3, wan2_2, ltx)	t unchanged	t * N

Sharing one hook required NFT to pre-multiply by num_train_timesteps so the shared hook had a scheduler-range value to work on. For qwen_image that hook then divides it straight back out, and (σ * 1000) / 1000 is not the identity in float32 — it lands a ULP off on ~2% of sigmas (measured over 200k samples). That is the exact rewrite #125 set out to prevent; its own test says "any rewrite of the expression drifts ULPs".

Behaviour

σ is the NFT pair sigma, N the scheduler range.

family	pre-#125	main today	this PR
sd3, wan2_2, ltx	σ * N	σ * N	σ * N
qwen_image	σ	(σ * N) / 1000	σ

Only qwen_image + NFT is affected; every family recovers its pre-#125 tensor bit-for-bit. The identity families were never wrong, which the tests assert explicitly — reverting nft.py alone leaves test_sd3_style_family_gets_the_scheduler_range green and fails only the qwen and wiring cases.

Why two hooks rather than one signature

The two constants are different facts that happen to share a value, and #125's hardcoded 1000 is right:

• σ ↔ scheduler-range timestep uses num_train_timesteps — a scheduler property (scheduling_flow_match_euler_discrete.py:148,239).
• scheduler-range timestep → qwen DiT input uses the literal 1000 — a model property; sglang-d's DiT divides by the same literal (runtime/models/dits/qwen_image.py:1536) and diffusers leaves that division to its caller.

Reading num_train_timesteps in process_timestep_as_input (the pre-#125 behaviour) would diverge from the rollout under any other scheduler range. Noted in a comment so neither hook gets "fixed" into reading the other.

Files

• configs/train_pipeline_config.py — hook added, base scales by num_train_timesteps
• configs/qwen_image.py — passes the sigma through
• loss_hub/nft.py — calls the sigma hook, no pre-multiply
• configs/ltx.py, configs/wan2_2.py — refactor(fsdp): upgrade needs_timestep_scaling to process_timestep_as_input in TrainPipelineConfig #125 leftovers
• tests/.../test_train_pipeline_config_registry.py — TestProcessSigmaAsTimestepsInput
• tests/.../test_loss_hub_nft.py — TestPrepareNftBatch

The bit-exactness guards use torch.equal, not allclose: the buggy composition passes allclose. 0.8474337458610535 is a sigma the 1000-round-trip does not return unchanged.

Stack

1. this PR — the inverted rescaling
2. fix(nft): seed the sigma shuffle and the corruption noise so NFT runs reproduce #156 — seed NFT's random streams, so its runs reproduce (found while trying to record an e2e standard for this)
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 matches clean main; the 4 failures there (test_hybrid_shard_mesh, test_lora_args, test_metric_buffer_dist) 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