From 60da0a38d3ffeedb03941c8293482d571493ccf2 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 10 Aug 2026 11:12:13 -0700 Subject: [PATCH 1/2] refactor(args): fold --diffusion-model into --hf-checkpoint Two flags named the same repo id from opposite ends. --hf-checkpoint fed the diffusers train-side loader (components, scheduler, model_index.json) and the eval dataset cache key; --diffusion-model fed the sglang-d engine's model_path, LTX's checkpoint resolution, and the model-family lookup. Five of the six recipes passed the identical value to both, and the config field the family lookup reads has always been called hf_ckpt_name_patterns -- the split was accretion, not design. miles has only --hf-checkpoint. Keep --hf-checkpoint, delete --diffusion-model, and point its three readers at the surviving flag. One value for one model means the engine and the trainer can no longer disagree about which pipeline is being trained, and the family can no longer be resolved from a different name than the weights. --hf-checkpoint is now required. It defaulted to None while --diffusion-model defaulted to stable-diffusion-3.5-medium, so a run that omitted both got an SD3.5 engine and family with a None training checkpoint -- it crashed in DiffusersModelBackend anyway, just later and less clearly. The three `if getattr(args, "diffusion_model", None):` guards in miles_validate_args were using that non-empty default as an "is this a diffusion run" sentinel. They are unconditional here: miles_validate_args has exactly one caller, parse_args, whose parser always registers these arguments. Renaming underneath them would have silently flipped them off whenever --hf-checkpoint was absent, skipping the TrainPipelineConfig resolution and validate_sp_args, and pinning sp_size to 1 so dp_size and the global-batch-size divisibility check came out wrong. They are now plain statements. The LTX recipe's `--hf-checkpoint gpt2` is dropped rather than carried over. LTX loads through MilesModelBackend, which resolves the DiT from the model reference and never touches hf_checkpoint; the only read on that path is the eval dataset cache key, and the recipe sets no eval flags at all. gpt2 also matches no family's hf_ckpt_name_patterns, which is precisely why the family lookup had to read the other flag. The recipe now passes --hf-checkpoint Lightricks/LTX-2.3. Verified by capturing the generated train_args of all six recipes plus the smoke and debug-alignment variants (7 in total): the only change is the removal of --diffusion-model and its duplicate value, and for LTX the removal of the dead gpt2, with Lightricks/LTX-2.3 surviving on --hf-checkpoint. Co-Authored-By: Claude Opus 5 --- .../configs/train_pipeline_config.py | 2 +- .../backends/fsdp_utils/models/ltx/loading.py | 14 ++-- .../sglang_diffusion_engine.py | 2 +- miles/utils/arguments.py | 83 +++++++++---------- scripts/run_diffusion_grpo_ltx23_sglang.py | 7 +- ...on_grpo_pickscore_5gpu_flowgrpo_aligned.py | 5 +- scripts/run_diffusion_grpo_sd3_ocr_sglang.py | 2 +- ...run_diffusion_grpo_wan22_pickscore_5gpu.py | 5 +- scripts/run_diffusion_nft_sd3_pickscore.py | 5 +- scripts/run_diffusion_sft_wan22.py | 2 +- tests/fast/utils/test_lora_args.py | 2 +- 11 files changed, 54 insertions(+), 75 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 9c0ee663..cbbad1cb 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -79,7 +79,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). diff --git a/miles/backends/fsdp_utils/models/ltx/loading.py b/miles/backends/fsdp_utils/models/ltx/loading.py index 2e131f3f..dbfefb34 100644 --- a/miles/backends/fsdp_utils/models/ltx/loading.py +++ b/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", diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 80386caa..2f9cdcb8 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/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, diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7a5426ff..2e9e25fc 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -185,11 +185,10 @@ 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 (hence the TrainPipelineConfig) is resolved from its name. Required." ), ) parser.add_argument( @@ -206,12 +205,6 @@ def add_rollout_arguments(parser): "and `truncated`." ), ) - parser.add_argument( - "--diffusion-model", - type=str, - default="stabilityai/stable-diffusion-3.5-medium", - help="HuggingFace model id for diffusion rollout.", - ) parser.add_argument( "--train-pipeline-config-path", type=str, @@ -1473,34 +1466,36 @@ 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.") - 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 - 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) + from miles.utils.misc import load_function + + 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 + 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.hf_checkpoint) + 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 +1508,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 +1623,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 +1689,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 diff --git a/scripts/run_diffusion_grpo_ltx23_sglang.py b/scripts/run_diffusion_grpo_ltx23_sglang.py index ee8f3a64..25081024 100644 --- a/scripts/run_diffusion_grpo_ltx23_sglang.py +++ b/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 " diff --git a/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py b/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py index df867237..a878e01f 100644 --- a/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py +++ b/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 " diff --git a/scripts/run_diffusion_grpo_sd3_ocr_sglang.py b/scripts/run_diffusion_grpo_sd3_ocr_sglang.py index 329bc776..33c71f6d 100644 --- a/scripts/run_diffusion_grpo_sd3_ocr_sglang.py +++ b/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 " diff --git a/scripts/run_diffusion_grpo_wan22_pickscore_5gpu.py b/scripts/run_diffusion_grpo_wan22_pickscore_5gpu.py index 1171bf9f..c3c3f133 100644 --- a/scripts/run_diffusion_grpo_wan22_pickscore_5gpu.py +++ b/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 " diff --git a/scripts/run_diffusion_nft_sd3_pickscore.py b/scripts/run_diffusion_nft_sd3_pickscore.py index c3b49fda..fa97dca4 100644 --- a/scripts/run_diffusion_nft_sd3_pickscore.py +++ b/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 " diff --git a/scripts/run_diffusion_sft_wan22.py b/scripts/run_diffusion_sft_wan22.py index 42dffd2c..29fd0461 100644 --- a/scripts/run_diffusion_sft_wan22.py +++ b/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: diff --git a/tests/fast/utils/test_lora_args.py b/tests/fast/utils/test_lora_args.py index 2fc931e3..c0434ff4 100644 --- a/tests/fast/utils/test_lora_args.py +++ b/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, From c3e3acd02c431549af1453ff1cb6761863b98ac5 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 10 Aug 2026 12:11:34 -0700 Subject: [PATCH 2/2] feat(args): add --diffusion-model-family so your own weights can name their family Folding --diffusion-model into --hf-checkpoint left one path narrower than it was. Before, you could point --hf-checkpoint at a local directory and --diffusion-model at the hub id, and the family lookup read the hub id. Now there is one value, so a checkpoint whose path does not happen to contain "ltx", "sd3", "wan2.2" or "qwen-image" cannot be resolved -- which is the normal case for a directory of your own weights. The escape hatch for that already existed but only as MILES_DIFFUSION_MODEL_FAMILY, an environment variable, so it never appears in --help. Promote it to a flag: --diffusion-model-family ltx --hf-checkpoint /root/models/my-own-dit Precedence is flag, then the env var, then name matching. --train-pipeline-config-path remains the answer for a family that is not registered at all; passing both now raises instead of silently letting the config path win. The value is normalized to stripped lowercase on args, because downstream readers compare it exactly -- encoder_hub.get_encoder tests `family == "wan2_2"`, so "WAN2_2" would have missed. This mirrors what miles does for the same problem from the other side. miles never pattern-matches the checkpoint name: it derives the model identity from the checkpoint's own config, `type(load_hf_config(args.hf_checkpoint)).__name__.lower()`, with --model-name as the explicit CLI override. Deriving from the pipeline class in model_index.json would be the closer analogue here, but LTX's official repo ships no model_index.json -- it loads a single-file DiT through MilesModelBackend -- so the one family that most needs the hook is the one that path cannot serve. Name matching with an explicit override is what works for all four families. Co-Authored-By: Claude Opus 5 --- .../configs/train_pipeline_config.py | 3 ++- miles/utils/arguments.py | 23 +++++++++++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index cbbad1cb..790db46c 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/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." ) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2e9e25fc..1f54b4ee 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -188,7 +188,8 @@ def add_rollout_arguments(parser): "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 (hence the TrainPipelineConfig) is resolved from its name. Required." + "family is matched from its name unless --diffusion-model-family says otherwise. " + "Required." ), ) parser.add_argument( @@ -205,6 +206,17 @@ def add_rollout_arguments(parser): "and `truncated`." ), ) + parser.add_argument( + "--diffusion-model-family", + type=str, + 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", type=str, @@ -1472,6 +1484,8 @@ def miles_validate_args(args): 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 @@ -1481,7 +1495,12 @@ def miles_validate_args(args): resolve_diffusion_model_family, ) - args.diffusion_model_family = resolve_diffusion_model_family(args.hf_checkpoint) + if args.diffusion_model_family is None: + args.diffusion_model_family = resolve_diffusion_model_family(args.hf_checkpoint) + else: + # 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: