diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 9c0ee663..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." ) @@ -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). 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..1f54b4ee 100644 --- a/miles/utils/arguments.py +++ b/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 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,