Skip to content

refactor(diffusion): require rollout scheduler sigmas, drop timesteps-derived fallbacks - #92

Merged
Rockdu merged 3 commits into
radixark:mainfrom
Rockdu:refactor/require-rollout-sigmas
Aug 3, 2026
Merged

refactor(diffusion): require rollout scheduler sigmas, drop timesteps-derived fallbacks#92
Rockdu merged 3 commits into
radixark:mainfrom
Rockdu:refactor/require-rollout-sigmas

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Re-lands #89: it was accidentally merged into the already-merged #88 branch instead of main (the stacked base was not auto-retargeted because the branch was not deleted). Same three commits, rebased onto current main.

What

  • Training now consumes the rollout engine's scheduler.sigmas snapshot verbatim; the timesteps / num_train_timesteps reconstruction fallbacks are removed (scheduler_meta_from_rollout, NFT converter, Flow-GRPO converter). Missing sigmas is a hard ValueError, making sglang 585a7d05e (#32683) the minimum rollout engine.
  • Both converters now share scheduler_meta_from_samples, which builds the batch scheduler meta from sample 0 and rejects any sample whose timesteps/sigmas differ — the NFT converter previously used sample 0's meta without checking the rest.

Why

The fallback was not equivalent to the snapshot: traj.timesteps is [T+1] (terminal 0.0 appended with the final x0 latent), so cat([timesteps / N, zeros(1)]) produced 12 entries for a 10-step schedule vs the snapshot's 11, and the division does not round-trip sigma * 1000 in fp32 (rel. diff up to 1.42e-8, measured in CI run 30778705251). Two paths producing almost the same array is exactly how stage-c-3-gpu-h200 silently drifted when upstream started supplying sigmas; failing loudly removes that bug class.

Validation

  • Numerical no-op by construction: sglang has populated dit_trajectory.sigmas unconditionally since 585a7d05e, so the deleted branches were already dead. Prior CI controls agree: pinning pre-sigmas sglang 3312645a3 reproduces the old standard 8/8 (run 30770528080), and nulling dit_trajectory.sigmas on current sglang reproduces it 8/8 (run 30772902308). No e2e standard re-record needed.
  • Smoke test on 8x H200 (sglang main b8109b5d + this branch): all 5 recipes ran 2 rollouts end-to-end — SD3-OCR GRPO, SD3 NFT (nft_num_timesteps=9, nft_t_mean=0.7304325, matching the sigmas-snapshot theory value), LTX-2.3 GRPO, Qwen-Image Flow-GRPO-aligned, Wan2.2-A14B (dual-DiT per-component metrics present). Healthy log_prob_mean_abs_diff on all: ~1e-4 (SD3), ~1e-6 (LTX), ~4e-5 (Qwen-Image), ~1e-5 (Wan2.2).
  • tests/fast locally: 136 passed (excluding test_metric_buffer_dist / test_hybrid_shard_mesh, which need a Linux dist/GPU environment).

Not in this PR (tracked separately): value-threshold trim for resolve_nft_sigmas (requires re-recording the NFT e2e standard), the degenerate terminal pair in next_timesteps, and the upstream [T] annotation fix (sgl-project/sglang#33332).

Files

  • miles/utils/train_data_utils.pyscheduler_meta_from_rollout requires scheduler_sigmas; new shared scheduler_meta_from_samples.
  • miles/backends/fsdp_utils/actor.py — drop the now-unused num_train_timesteps plumbing.
  • miles/ray/data_conversion_hub/flow_grpo.py, miles/ray/data_conversion_hub/nft.py — use the shared helper; fallbacks removed.
  • miles/utils/types.pyDiTTrajectory.sigmas comment updated.
  • tests/fast/utils/test_grouping_parity.py, tests/fast/backends/fsdp_utils/test_loss_hub_nft.py — sigmas-required and mismatched-meta tests (registered stage-a-cpu).

Checklist

  • pre-commit run --all-files passes — not run locally
  • Added/updated tests for new behaviour
  • pytest -x is green — tests/fast locally, see Validation
  • If launch flags changed, python3 train.py --help still parses — N/A, no flag changes
  • If a public flag was added, it appears in the CLI reference docs — N/A
  • If an example was added, it has a real walkthrough — N/A

@Rockdu
Rockdu merged commit 90ce94d into radixark:main Aug 3, 2026
15 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