Skip to content

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

Merged
Rockdu merged 3 commits into
radixark:fix/rerecord-sd3-ocr-sglang-driftfrom
Rockdu:refactor/require-rollout-sigmas
Aug 3, 2026
Merged

refactor(diffusion): require rollout scheduler sigmas, drop timesteps-derived fallbacks#89
Rockdu merged 3 commits into
radixark:fix/rerecord-sd3-ocr-sglang-driftfrom
Rockdu:refactor/require-rollout-sigmas

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Paired with #88 (the OCR standard re-record)

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

  • 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 force-pushed the refactor/require-rollout-sigmas branch from b2b02a8 to df2e704 Compare August 3, 2026 07:52
@Rockdu
Rockdu changed the base branch from main to fix/rerecord-sd3-ocr-sglang-drift August 3, 2026 07:52
@Rockdu
Rockdu marked this pull request as ready for review August 3, 2026 08:39
@Rockdu
Rockdu merged commit 6720ee4 into radixark:fix/rerecord-sd3-ocr-sglang-drift Aug 3, 2026
27 of 28 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