Skip to content

refactor(diffusion): move rollout patch group selection to a CLI arg - #97

Merged
Rockdu merged 1 commit into
radixark:mainfrom
Rockdu:tpc-decouple-flying-wires
Aug 8, 2026
Merged

Rockdu merged 1 commit into
radixark:mainfrom
Rockdu:tpc-decouple-flying-wires

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

What

  • Remove TrainPipelineConfig.rollout_patch_group: family config classes no longer inject rollout monkey-patch groups behind the user's back.
  • Add --rollout-patch-group (comma-separated) as the single way to select patch groups at launch; it also replaces --apply-sgld-monkey-patches (now --rollout-patch-group sgld). The LTX launch script passes --rollout-patch-group ltx.
  • Validate every selected group against the @register_rollout_patch_group registry in miles_validate_args, so an unknown group fails at arg validation instead of inside the sglang-d scheduler child at engine startup.

Why

rollout_patch_group was the last direct config→rollout coupling: a train-side config class decided which monkey patches the inference engine applies, with the knowledge crossing two processes (train arg validation → MILES_ROLLOUT_PATCH_GROUPS env → sglang-d scheduler) and no validation until deep inside the engine. Selecting patch groups in launch arguments keeps arg validation to validation only (no inference from config classes) and makes the rollout-parity patching visible in the launch command. Behavior is unchanged for the in-tree recipes: the LTX and flowgrpo-aligned scripts pass the flag explicitly. --apply-sgld-monkey-patches is removed (breaking for out-of-tree launch commands): keeping two selection mechanisms for the same registry would leave the old implicit path alive.

Files

  • miles/utils/arguments.py — add --rollout-patch-group, resolve it into args.rollout_patch_groups, validate groups against the registry; drop the cfg_cls.rollout_patch_group injection and --apply-sgld-monkey-patches
  • miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py — add validate_rollout_patch_groups(); reuse it in apply_env_selected_rollout_patches
  • miles/backends/fsdp_utils/configs/train_pipeline_config.py — remove the rollout_patch_group field
  • miles/backends/fsdp_utils/configs/ltx.py — drop rollout_patch_group = "ltx"
  • scripts/run-diffusion-grpo-ltx23-sglang.sh, scripts/run-diffusion-grpo-pickscore-5gpu-flowgrpo-aligned.sh — pass --rollout-patch-group ltx / sgld
  • tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py — cover the validator (registered in stage-a-cpu)

Checklist

  • pre-commit run --all-files passes
  • Added/updated tests for new behaviour
  • pytest -x is green — tests/fast: 147 passed; test_hybrid_shard_mesh / test_metric_buffer_dist fail identically on clean main in this sandbox (no local torch.distributed rendezvous), unrelated to this change
  • If launch flags changed, python3 train_diffusion.py --help still parses (new flag renders; run with a stubbed sglang, CPU-only env)
  • If a public flag was added, it appears in the CLI reference docs — repo has no CLI reference docs; N/A
  • If an example was added, it has a real walkthrough — N/A, no example added

@Rockdu
Rockdu marked this pull request as ready for review August 6, 2026 22:40
@Rockdu
Rockdu force-pushed the tpc-decouple-flying-wires branch 2 times, most recently from d766c6f to 985fd5b Compare August 8, 2026 21:38
@Rockdu
Rockdu force-pushed the tpc-decouple-flying-wires branch from 985fd5b to 39cc33e Compare August 8, 2026 21:48
@Rockdu
Rockdu merged commit 3035666 into radixark:main Aug 8, 2026
17 checks passed
@Rockdu
Rockdu deleted the tpc-decouple-flying-wires branch August 8, 2026 22:38
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