refactor(args): fold --diffusion-model into --hf-checkpoint

miles/backends/fsdp_utils/configs/train_pipeline_config.py

@@ -59,7 +59,8 @@ def resolve_diffusion_model_family(model_ref: str) -> str:
     raise ValueError(
         f"Cannot resolve diffusion model family for '{model_ref}' "
         f"(known families: {list(_REGISTRY)}). "
-        "Set MILES_DIFFUSION_MODEL_FAMILY to override."
+        "Name it with --diffusion-model-family, or point --train-pipeline-config-path at your own "
+        "TrainPipelineConfig if the family is not one of these."
     )
 
 
@@ -79,7 +80,7 @@ class TrainPipelineConfig(abc.ABC):
     model_family: str | None = None
     lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"]
     optimizer_state_allowed_missing: list[str] = []
-    # Case-insensitive substrings matched against the checkpoint name (--diffusion-model).
+    # Case-insensitive substrings matched against the checkpoint name (--hf-checkpoint).
     hf_ckpt_name_patterns: tuple[str, ...] = ()
     supports_cfg_training: bool = True
     # Mirrors serial sgl-d serving; not valid when the rollout engine runs --enable-cfg-parallel (branches split per rank, different combine formula).

miles/backends/fsdp_utils/models/ltx/loading.py

@@ -166,19 +166,19 @@ def resolve_materialized_model_dir(
 
 
 def resolve_transformer_checkpoint(
-    diffusion_model: str | None,
+    hf_checkpoint: str | None,
     *,
     materialize: bool = True,
 ) -> str:
     """Resolve the single-file DiT checkpoint used by FSDP train."""
-    if diffusion_model:
-        path = Path(str(diffusion_model)).expanduser()
+    if hf_checkpoint:
+        path = Path(str(hf_checkpoint)).expanduser()
         if path.is_file() and path.suffix == ".safetensors":
             return str(path)
 
-        if _is_hf_model_id(str(diffusion_model)):
+        if _is_hf_model_id(str(hf_checkpoint)):
             materialized_dir = resolve_materialized_model_dir(
-                str(diffusion_model),
+                str(hf_checkpoint),
                 materialize=materialize,
             )
             if materialized_dir is not None:
@@ -191,7 +191,7 @@ def resolve_transformer_checkpoint(
                 return str(checkpoint)
 
     raise FileNotFoundError(
-        "Could not resolve LTX transformer checkpoint. Pass --diffusion-model "
+        "Could not resolve LTX transformer checkpoint. Pass --hf-checkpoint "
         "Lightricks/LTX-2.3 (recommended) or a .safetensors override."
     )
 
@@ -205,7 +205,7 @@ def load_component(
 ):
     if component != TRAIN_COMPONENT:
         raise ValueError(f"LTX trains the single DiT ({TRAIN_COMPONENT!r}); got {component!r}")
-    checkpoint = resolve_transformer_checkpoint(str(args.diffusion_model))
+    checkpoint = resolve_transformer_checkpoint(str(args.hf_checkpoint))
     return load_transformer_for_train(
         checkpoint,
         device="cpu",

miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py

@@ -315,7 +315,7 @@ def _compute_server_args(args, host, port, nccl_port):
     # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done
     # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg.
     kwargs = {
-        "model_path": args.diffusion_model,
+        "model_path": args.hf_checkpoint,
         "trust_remote_code": True,
         "host": host,
         "port": port,

miles/utils/arguments.py

@@ -185,11 +185,11 @@ def add_rollout_arguments(parser):
                 type=str,
                 default=None,
                 help=(
-                    "The huggingface checkpoint of the trained model. "
-                    "This is used to initialize sglang and also provide the tokenizer. "
-                    "Note that, we will always update the parameters in sglang with that of megatron before training, "
-                    "so you only need to provide a huggingface checkpoint that has the same architecture as the model you want to train. "
-                    "It doesn't necessary need to contain the most up-to-date parameters."
+                    "The diffusers pipeline to train, as a HuggingFace repo id or a local directory. "
+                    "One value serves three readers, so they cannot disagree: the training side loads "
+                    "the components and scheduler from it, the sglang-d engine serves it, and the model "
+                    "family is matched from its name unless --diffusion-model-family says otherwise. "
+                    "Required."
                 ),
             )
             parser.add_argument(
@@ -207,10 +207,15 @@ def add_rollout_arguments(parser):
                 ),
             )
             parser.add_argument(
-                "--diffusion-model",
+                "--diffusion-model-family",
                 type=str,
-                default="stabilityai/stable-diffusion-3.5-medium",
-                help="HuggingFace model id for diffusion rollout.",
+                default=None,
+                help=(
+                    "Registered family key, e.g. sd3, wan2_2, ltx, qwen_image. Default: matched from "
+                    "--hf-checkpoint against each family's name patterns. Pass it when the checkpoint "
+                    "does not carry the family name, which your own local weights usually do not. Use "
+                    "--train-pipeline-config-path instead for a family that is not registered."
+                ),
             )
             parser.add_argument(
                 "--train-pipeline-config-path",
@@ -1473,34 +1478,43 @@ def miles_validate_args(args):
 
     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
+    if not args.hf_checkpoint:
+        raise ValueError("--hf-checkpoint is required: it names the diffusers pipeline to train and to serve.")
+
+    from miles.utils.misc import load_function
+
+    if args.train_pipeline_config_path is not None:
+        if args.diffusion_model_family is not None:
+            raise ValueError("--train-pipeline-config-path and --diffusion-model-family both name a config; pass one.")
+        # Explicit config path IS the identity (custom classes never need registering).
+        cfg_cls = load_function(args.train_pipeline_config_path)
+        args.diffusion_model_family = None
+    else:
+        from miles.backends.fsdp_utils.configs.train_pipeline_config import (
+            get_train_pipeline_config_cls,
+            resolve_diffusion_model_family,
+        )
 
-        if args.train_pipeline_config_path is not None:
-            # Explicit config path IS the identity (custom classes never need registering).
-            cfg_cls = load_function(args.train_pipeline_config_path)
-            args.diffusion_model_family = None
+        if args.diffusion_model_family is None:
+            args.diffusion_model_family = resolve_diffusion_model_family(args.hf_checkpoint)
         else:
-            from miles.backends.fsdp_utils.configs.train_pipeline_config import (
-                get_train_pipeline_config_cls,
-                resolve_diffusion_model_family,
-            )
-
-            args.diffusion_model_family = resolve_diffusion_model_family(args.diffusion_model)
-            cfg_cls = get_train_pipeline_config_cls(args.diffusion_model_family)
-            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 not cfg_cls.supports_cfg_training and (
-            args.diffusion_guidance_scale != 1.0 or args.diffusion_negative_prompt is not None
-        ):
-            raise ValueError(
-                f"{cfg_cls.__name__} trains unguided (supports_cfg_training=False); set "
-                f"--diffusion-guidance-scale 1.0 and drop --diffusion-negative-prompt"
-            )
-        cfg_cls.validate_args(args)
-        if args.use_lora and args.lora_target_modules is None:
-            args.lora_target_modules = list(cfg_cls.lora_target_modules)
+            # Downstream lookups compare this exactly (encoder_hub.get_encoder), so normalize
+            # here rather than at every reader.
+            args.diffusion_model_family = args.diffusion_model_family.strip().lower()
+        cfg_cls = get_train_pipeline_config_cls(args.diffusion_model_family)
+        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 not cfg_cls.supports_cfg_training and (
+        args.diffusion_guidance_scale != 1.0 or args.diffusion_negative_prompt is not None
+    ):
+        raise ValueError(
+            f"{cfg_cls.__name__} trains unguided (supports_cfg_training=False); set "
+            f"--diffusion-guidance-scale 1.0 and drop --diffusion-negative-prompt"
+        )
+    cfg_cls.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
@@ -1513,7 +1527,7 @@ def miles_validate_args(args):
         if not args.lora_target_modules:
             raise ValueError(
                 "--lora-ipc-weight-sync requires LoRA target modules; "
-                "set --diffusion-model (for per-model defaults) or --lora-target-modules."
+                "set --hf-checkpoint (for per-model defaults) or --lora-target-modules."
             )
 
     if not 0.0 <= args.ema_decay_init <= 1.0:
@@ -1628,11 +1642,10 @@ def miles_validate_args(args):
         args.colocate = False
         args.offload_train = args.offload_rollout = False
 
-    if getattr(args, "diffusion_model", None):
-        from miles.backends.fsdp_utils.arguments import validate_hybrid_shard_args, validate_sp_args
+    from miles.backends.fsdp_utils.arguments import validate_hybrid_shard_args, validate_sp_args
 
-        validate_sp_args(args)
-        validate_hybrid_shard_args(args)
+    validate_sp_args(args)
+    validate_hybrid_shard_args(args)
 
     # always true on offload for colocate at the moment.
     if args.colocate:
@@ -1695,8 +1708,7 @@ def miles_validate_args(args):
         args.global_batch_size = derived_gbs
 
     train_world_size = args.actor_num_gpus_per_node * args.actor_num_nodes
-    sp_size = args.sequence_parallel_size if getattr(args, "diffusion_model", None) else 1
-    dp_size = train_world_size // sp_size
+    dp_size = train_world_size // args.sequence_parallel_size
     if args.global_batch_size is not None:
         assert (
             args.global_batch_size % dp_size == 0

scripts/run_diffusion_grpo_ltx23_sglang.py

@@ -4,9 +4,6 @@
 per epoch from candidate steps 0-9. Everything runs bf16 end to end — master, reduce,
 forward and the sgl-d engine — on the sdpa_math attention backend.
 
---hf-checkpoint points at gpt2 on purpose: LTX-2.3 needs no HF tokenizer here, and the flag
-still wants a resolvable repo id.
-
 Video rollouts take minutes per request, so the health checker gets a far longer interval
 and failure budget than the image recipes.
 
@@ -42,9 +39,7 @@ def prepare(args: ScriptArgs) -> str:
 def execute(args: ScriptArgs, data_dir: str) -> None:
     run_name = f"diffusion_grpo_ltx23_pickscore_{U.create_run_id()}"
 
-    ckpt_args = (
-        f"--hf-checkpoint gpt2 --diffusion-model {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 50 "
-    )
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 50 "
 
     rollout_args = (
         "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "

scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py

@@ -43,10 +43,7 @@ def prepare(args: ScriptArgs) -> str:
 def execute(args: ScriptArgs, data_dir: str) -> None:
     run_name = f"diffusion_grpo_pickscore_5gpu_flowgrpo_aligned_{U.create_run_id()}"
 
-    ckpt_args = (
-        f"--hf-checkpoint {MODEL} --diffusion-model {MODEL} "
-        f"--save {args.output_dir}/{run_name}/ckpt --save-interval 10 "
-    )
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 10 "
 
     rollout_args = (
         "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "

scripts/run_diffusion_grpo_sd3_ocr_sglang.py

@@ -44,7 +44,7 @@ def prepare(args: ScriptArgs) -> str:
 def execute(args: ScriptArgs, data_dir: str) -> None:
     run_name = f"diffusion_grpo_sd3_ocr_sglang_{U.create_run_id()}"
 
-    ckpt_args = f"--hf-checkpoint {MODEL} --diffusion-model {MODEL} --save {args.output_dir}/{run_name}/ckpt "
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt "
 
     rollout_args = (
         "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "

scripts/run_diffusion_grpo_wan22_pickscore_5gpu.py

@@ -56,10 +56,7 @@ def prepare(args: ScriptArgs) -> str:
 def execute(args: ScriptArgs, data_dir: str) -> None:
     run_name = f"diffusion_grpo_wan22_pickscore_5gpu_{U.create_run_id()}"
 
-    ckpt_args = (
-        f"--hf-checkpoint {MODEL} --diffusion-model {MODEL} "
-        f"--save {args.output_dir}/{run_name}/ckpt --save-interval 10 "
-    )
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 10 "
 
     rollout_args = (
         "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "

scripts/run_diffusion_nft_sd3_pickscore.py

@@ -49,10 +49,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None:
     run_name = f"diffusion_nft_sd3_pickscore_{U.create_run_id()}"
     num_rollout = args.num_rollout or (1 if args.smoke else 100)
 
-    ckpt_args = (
-        f"--hf-checkpoint {MODEL} --diffusion-model {MODEL} "
-        f"--save {args.output_dir}/{run_name}/ckpt --save-interval 20 "
-    )
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 20 "
 
     rollout_args = (
         "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "

scripts/run_diffusion_sft_wan22.py

@@ -45,7 +45,7 @@ def execute(args: ScriptArgs) -> None:
     run_name = f"diffusion_sft_wan22_{U.create_run_id()}"
 
     ckpt_args = (
-        f"--hf-checkpoint {MODEL} --diffusion-model {MODEL} --sft-encoder-checkpoint {MODEL} "
+        f"--hf-checkpoint {MODEL} --sft-encoder-checkpoint {MODEL} "
         f"--save {args.output_dir}/{run_name}/ckpt --save-interval 20 "
     )
     if args.resume_ckpt:

tests/fast/utils/test_lora_args.py

@@ -9,7 +9,7 @@
 
 def _server_args(**overrides):
     base = dict(
-        diffusion_model="Qwen/Qwen-Image",
+        hf_checkpoint="Qwen/Qwen-Image",
         diffusion_flow_shift=None,
         rollout_num_gpus_per_engine=1,
         sglang_sp_degree=None,