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

miles/backends/fsdp_utils/configs/ltx.py

@@ -19,7 +19,6 @@ class LTXTrainPipelineConfig(TrainPipelineConfig):
     supports_cfg_training = False
     # Rollout stores σ×1000 in trajectory timesteps; ltx_core AdaLN uses σ∈[0,1].
     sde_timestep_divisor = 1000.0
-    rollout_patch_group = "ltx"
     hf_ckpt_name_patterns = ("ltx",)
     model_backend_path = "miles.backends.fsdp_utils.model_backend.MilesModelBackend"
     model_package = "miles.backends.fsdp_utils.models.ltx"

miles/backends/fsdp_utils/configs/train_pipeline_config.py

@@ -83,8 +83,6 @@ class TrainPipelineConfig(abc.ABC):
     # Case-insensitive substrings matched against the checkpoint name (--diffusion-model).
     hf_ckpt_name_patterns: tuple[str, ...] = ()
     supports_cfg_training: bool = True
-    # Rollout parity patch group applied by the engine (see monkey_patches; None = none).
-    rollout_patch_group: str | None = None
     # Model-boundary input dtypes (see input_dtype_policy); families opt into casts explicitly.
     input_dtype_policy: dict = {"latents": None, "cond": None, "timestep": None}
     # Default component paths (miles custom-function style); CLI args override.

miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py

@@ -66,10 +66,19 @@ def apply_ltx2_rollout_patches() -> None:
     patch_ltx2_disable_av_cross.apply()
 
 
+def validate_rollout_patch_groups(names: list[str]) -> None:
+    """Reject group names with no registered applier; shared by arg validation and env selection."""
+    unknown = [name for name in names if name not in _ROLLOUT_PATCH_APPLIERS]
+    if unknown:
+        raise ValueError(
+            f"Unknown rollout patch group(s) {unknown}; known: {list(_ROLLOUT_PATCH_APPLIERS)}. "
+            "Each group must be registered here via @register_rollout_patch_group."
+        )
+
+
 def apply_env_selected_rollout_patches() -> None:
     """Apply every group named in the env list (runs in the scheduler grandchild)."""
-    for name in filter(None, os.environ.get(ROLLOUT_PATCH_GROUPS_ENV, "").split(",")):
-        applier = _ROLLOUT_PATCH_APPLIERS.get(name)
-        if applier is None:
-            raise ValueError(f"Unknown rollout patch group {name!r}; known: {list(_ROLLOUT_PATCH_APPLIERS)}")
-        applier()
+    names = [name for name in os.environ.get(ROLLOUT_PATCH_GROUPS_ENV, "").split(",") if name]
+    validate_rollout_patch_groups(names)
+    for name in names:
+        _ROLLOUT_PATCH_APPLIERS[name]()

miles/utils/arguments.py

@@ -436,15 +436,14 @@ def add_rollout_arguments(parser):
                 help="Set rollout_log_prob_no_const=true on POST /rollout/generate.",
             )
             parser.add_argument(
-                "--apply-sgld-monkey-patches",
-                action="store_true",
-                default=False,
+                "--rollout-patch-group",
+                type=str,
+                default=None,
                 help=(
-                    "Apply miles.backends.sglang_diffusion_utils.monkey_patches at "
-                    "sglang-d startup so its DiT forward is bit-exact with diffusers' "
-                    "implementation. Makes rollout (sglang-d path) and training-side "
-                    "log-prob agree on noise_pred down to bf16 ULPs. Small perf hit on "
-                    "the rollout engine."
+                    "Comma-separated rollout patch groups applied at sglang-d startup so its "
+                    "forward is numerically aligned with the training side, e.g. 'sgld' "
+                    "(diffusers op parity, small rollout perf hit) or 'ltx' "
+                    "(see sglang_diffusion_utils/monkey_patches)."
                 ),
             )
             parser.add_argument(
@@ -1557,7 +1556,7 @@ def miles_validate_args(args):
     if args.diffusion_log_image_interval < 1:
         raise ValueError(f"diffusion_log_image_interval must be >= 1, got {args.diffusion_log_image_interval}")
 
-    args.rollout_patch_groups = ["sgld"] if args.apply_sgld_monkey_patches else []
+    args.rollout_patch_groups = [name for name in (args.rollout_patch_group or "").split(",") if name]
 
     if getattr(args, "diffusion_model", None):
         from miles.utils.misc import load_function
@@ -1577,8 +1576,6 @@ def miles_validate_args(args):
             args.train_pipeline_config_path = f"{cfg_cls.__module__}.{cfg_cls.__qualname__}"
         if args.model_backend_path is None:
             args.model_backend_path = cfg_cls.model_backend_path
-        if cfg_cls.rollout_patch_group:
-            args.rollout_patch_groups.append(cfg_cls.rollout_patch_group)
         if not cfg_cls.supports_cfg_training and (
             args.diffusion_guidance_scale != 1.0 or args.diffusion_negative_prompt is not None
         ):
@@ -1590,6 +1587,11 @@ def miles_validate_args(args):
         if args.use_lora and args.lora_target_modules is None:
             args.lora_target_modules = list(cfg_cls.lora_target_modules)
 
+    if args.rollout_patch_groups:
+        from miles.backends.sglang_diffusion_utils.monkey_patches import validate_rollout_patch_groups
+
+        validate_rollout_patch_groups(args.rollout_patch_groups)
+
     if args.lora_ipc_weight_sync:
         if not args.use_lora:
             raise ValueError("--lora-ipc-weight-sync requires --use-lora")

scripts/run-diffusion-grpo-ltx23-sglang.sh

@@ -35,6 +35,7 @@ fi
   --deterministic-mode \
   --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \
   --diffusion-model Lightricks/LTX-2.3 \
+  --rollout-patch-group ltx \
   --hf-checkpoint gpt2 \
   --prompt-data "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" \
   --input-key input \

scripts/run-diffusion-grpo-pickscore-5gpu-flowgrpo-aligned.sh

@@ -102,7 +102,7 @@ hf download --repo-type dataset rockdu/miles-diffusion-datasets \
   --diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window \
   --diffusion-num-sde-steps 2 \
   --diffusion-sde-window-range 3,5 \
-  --apply-sgld-monkey-patches \
+  --rollout-patch-group sgld \
   --diffusion-height 512 \
   --diffusion-width 512 \
   --save "${SAVE_DIR}" \

tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py

@@ -44,3 +44,13 @@ def test_unknown_group_fails_loud(self, monkeypatch):
     def test_builtin_group_registered(self):
         # The decorator ran at import time for the in-repo group.
         assert "sgld" in mp._ROLLOUT_PATCH_APPLIERS
+
+
+class TestValidateRolloutPatchGroups:
+    # The arg-validation entry point behind --rollout-patch-group:
+    #   --rollout-patch-group "sgld,ltx"  ──► registered appliers ──► pass
+    #   --rollout-patch-group "sgld,bogus" ─► "bogus" unregistered ─► ValueError
+    def test_known_pass_unknown_raises(self):
+        mp.validate_rollout_patch_groups(["sgld", "ltx"])
+        with pytest.raises(ValueError, match="Unknown rollout patch group"):
+            mp.validate_rollout_patch_groups(["sgld", "bogus"])