Skip to content

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

Draft
Rockdu wants to merge 1 commit into
radixark:mainfrom
Rockdu:tpc-decouple-flying-wires
Draft

refactor(diffusion): move rollout patch group selection to a CLI arg#97
Rockdu wants to merge 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) to select patch groups explicitly at launch; the LTX launch script now 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 script/e2e passes the flag explicitly, and --apply-sgld-monkey-patches is untouched.

Files

  • miles/utils/arguments.py — add --rollout-patch-group, merge it into args.rollout_patch_groups, validate groups against the registry; drop the cfg_cls.rollout_patch_group injection
  • 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; document the sde_timestep_divisor unit convention
  • miles/backends/fsdp_utils/configs/ltx.py — drop rollout_patch_group = "ltx"
  • scripts/run-diffusion-grpo-ltx23-sglang.sh — pass --rollout-patch-group ltx
  • 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

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