Skip to content

fix(rollout): apply rollout patches through sglang's scheduler bootstrap - #284

Merged
Rockdu merged 3 commits into
mainfrom
fix/rollout-patches-bootstrap
Oct 4, 2026
Merged

Rockdu merged 3 commits into
mainfrom
fix/rollout-patches-bootstrap

Conversation

@Rockdu

@Rockdu Rockdu commented Oct 4, 2026

Copy link
Copy Markdown
Collaborator

What

  • Rebind launch_server.bootstrap_scheduler_process, the scheduler spawn target on current sglang-miles, so the wan, qwen_image and ltx rollout patch groups run in the scheduler child again.
  • The wrapper initializes the platform first, then applies the patches, then hands over to sglang's own bootstrap, whose platform init is then a no-op.

Why

miles applied its rollout patches by rebinding launch_server.run_scheduler_process, which spawn pickles by qualname. sglang-miles now spawns worker_bootstrap.bootstrap_scheduler_process instead, so that rebinding replaced a name nothing calls, and no patch group reached the engine. Nothing errors; the engine just runs its own bf16 fused norm and rope instead of the fp32 shapes the trainer uses.

On test_wan22_pickscore_grpo_17xGPU_single_node_4xGPU_proxy this showed up as train/model_output_mean_abs_diff 0.0 -> 0.020 and train/model_output_rel_max 0.0 -> 0.057, plus shifted step-0 rewards. Re-recording that standard would have hidden a train/rollout parity break. A probe with CUDA IPC disabled on the scheduler hop reproduced the same numbers bit for bit, which ruled out the sglang transport fixes running alongside.

Validation

Full PR Test dispatch on #274 + this fix + #283's ring-attention fix + #282 (--continue-on-error), with sglang from Rockdu/sglang:ci/sglang-miles-ipc-host-and-port (sglang-miles plus sgl-project/sglang#42436 and #42470): https://github.com/radixark/miles_diffusion/actions/runs/37197860435

suite result
stage-c-3-gpu-h200 4/4: sd3_nft_pickscore, sd3_ocr_grpo, krea2_nft_ocr, h3_t2va_grpo
stage-c-5-gpu-h200 5/5: wan22, qwenimage, ltx23, cosmos3, attention parity

wan22, qwenimage and ltx23 are the recipes that use a rollout patch group; all three match their standards again.

Files

  • miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py: the scheduler wrapper and its rebinding.

Checklist

  • pre-commit run --all-files passes — not run
  • Added/updated tests for new behaviour — no new test; the wan22, qwenimage and ltx23 e2e standards are what cover it, and they pass in the run above
  • pytest -x is green — not run locally; GPU CI above is green
  • If launch flags changed, python3 train.py --help still parses — n/a
  • 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 added the run-ci-e2e Run e2e metric-regression tests on this PR label Oct 4, 2026
@Rockdu
Rockdu force-pushed the fix/rollout-patches-bootstrap branch from f885d9f to d4b1905 Compare October 4, 2026 19:19
@Rockdu
Rockdu merged commit da1d6a9 into main Oct 4, 2026
18 of 19 checks passed
OnePunchMonk pushed a commit to OnePunchMonk/miles_diffusion that referenced this pull request Oct 5, 2026
OnePunchMonk pushed a commit to OnePunchMonk/miles_diffusion that referenced this pull request Oct 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-e2e Run e2e metric-regression tests on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant