diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index f4a2604c..3f78d290 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/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" diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 6e432575..846beeda 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/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. diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 2755e204..e377f75e 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/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]() diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 18ba848a..5a6821de 100644 --- a/miles/utils/arguments.py +++ b/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") diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index f6693627..c06558fd 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/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 \ diff --git a/scripts/run-diffusion-grpo-pickscore-5gpu-flowgrpo-aligned.sh b/scripts/run-diffusion-grpo-pickscore-5gpu-flowgrpo-aligned.sh index 47763473..7a1bdb55 100755 --- a/scripts/run-diffusion-grpo-pickscore-5gpu-flowgrpo-aligned.sh +++ b/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}" \ diff --git a/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py b/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py index 2a174ac0..b8fcd741 100644 --- a/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py +++ b/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"])