From 110c7a8281e65ec85fb8bed9200cb2e1a3a01235 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 15 Jul 2026 06:35:29 +0000 Subject: [PATCH 01/10] feat(fsdp): meta-init + rank0-streamed weight loading (--fsdp-load-mode stream) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every rank used to materialize the full pipeline on CPU (DiffusionPipeline.from_pretrained), move the whole unsharded model to GPU, and only then fully_shard — peak GPU = full model (54 GB for one Wan2.2-A14B expert in fp32 master dtype), N×full disk reads, and no way to honor a meta init context (diffusers' pipeline loader crashes under one). New default path (--fsdp-load-mode stream), per component: build on meta via init_empty_weights(include_buffers=False) # buffers real -> fully_shard on meta (costs nothing) -> to_empty: allocate only this rank's shards; carry non-persistent buffers (Wan rope tables) across the transition -> rank0 streams safetensors file-by-file, set_model_state_dict broadcast_from_rank0 scatters into the shards -> LoRA adapters build on meta too (low_cpu_mem_usage) and initialize after materialization; data-aware inits (pissa/olora) are rejected -> family postprocess hook runs the before-fsdp hook after materialization, where qwen-image's rope parity rebuild has real CUDA tensors instead of silently no-oping on meta --fsdp-load-mode legacy keeps the old path for custom pipelines. Wan2.2-A14B (both experts, fp32, 2×H200): init 285s -> 71s, peak GPU 80.5 -> 54.0 GB (= the sharded params themselves); only rank0 reads the checkpoint. Streamed weights verified bitwise-equal to from_pretrained (params + buffers) at 1.3B scale and in the new 2-GPU CI test. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 117 +++++++++++-- miles/backends/fsdp_utils/arguments.py | 6 + .../configs/train_pipeline_config.py | 8 + miles/backends/fsdp_utils/model_backend.py | 165 +++++++++++++++++- .../test_fsdp_stream_load_equivalence.py | 160 +++++++++++++++++ .../backends/fsdp_utils/test_stream_load.py | 164 +++++++++++++++++ 6 files changed, 601 insertions(+), 19 deletions(-) create mode 100644 tests/fast-gpu/test_fsdp_stream_load_equivalence.py create mode 100644 tests/fast/backends/fsdp_utils/test_stream_load.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index c86f07c5..29ad6770 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -108,9 +108,19 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.train_pipeline_config = load_function(args.train_pipeline_config_path)() self.model_backend = load_function(args.model_backend_path)(self.train_pipeline_config) - raw_models, self.scheduler = self.model_backend.load_models_and_scheduler( - args, master_dtype=self._master_dtype - ) + + if args.fsdp_load_mode not in ("stream", "legacy"): + raise ValueError(f"--fsdp-load-mode must be 'stream' or 'legacy', got {args.fsdp_load_mode!r}") + stream_load = args.fsdp_load_mode == "stream" + if stream_load: + # params on meta: no rank ever holds a full unsharded copy + raw_models, self.scheduler = self.model_backend.build_models_and_scheduler( + args, master_dtype=self._master_dtype + ) + else: + raw_models, self.scheduler = self.model_backend.load_models_and_scheduler( + args, master_dtype=self._master_dtype + ) self.models: dict[str, torch.nn.Module] = {} for component, model in raw_models.items(): @@ -118,25 +128,47 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if args.fsdp_attention_backend is not None: self.model_backend.set_attention_backend(model, args.fsdp_attention_backend) + key_map = None if args.use_lora: - model = apply_lora(model, args, self.train_pipeline_config) + model = apply_lora(model, args, self.train_pipeline_config, on_meta=stream_load) + if stream_load: + key_map = peft_checkpoint_key_map(model) model.train() if args.gradient_checkpointing: self.model_backend.enable_gradient_checkpointing(model) - model.to(torch.cuda.current_device()) - - self.train_pipeline_config.preprocess_model_before_fsdp(model) - - model = apply_fsdp2( - model, - mesh=self.parallel_state.dp_mesh, - cpu_offload=self.args.fsdp_cpu_offload, - args=self.args, - no_split_modules=self.model_backend.fsdp_no_split_modules(model), - ) + if stream_load: + # shard while still on meta (costs nothing), materialize only + # this rank's shards, then fill weights file-by-file from rank0. + model = apply_fsdp2( + model, + mesh=self.parallel_state.dp_mesh, + cpu_offload=self.args.fsdp_cpu_offload, + args=self.args, + no_split_modules=self.model_backend.fsdp_no_split_modules(model), + ) + device = torch.device("cpu") if self.args.fsdp_cpu_offload else torch.cuda.current_device() + materialize_sharded_model(model, device) + self.model_backend.stream_load_weights( + model, component, args, master_dtype=self._master_dtype, key_map=key_map + ) + if args.use_lora: + reset_lora_adapters(model, args.diffusion_init_lora_weight) + self.train_pipeline_config.postprocess_model_after_materialize(model) + else: + model.to(torch.cuda.current_device()) + + self.train_pipeline_config.preprocess_model_before_fsdp(model) + + model = apply_fsdp2( + model, + mesh=self.parallel_state.dp_mesh, + cpu_offload=self.args.fsdp_cpu_offload, + args=self.args, + no_split_modules=self.model_backend.fsdp_no_split_modules(model), + ) self.models[component] = model # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -692,13 +724,16 @@ def _resolve_dtype(name: str) -> torch.dtype: return {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[name] -def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> None: +def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config, on_meta: bool = False) -> None: """Apply PEFT LoRA to the model. Args: model: The model to apply LoRA to. args: Arguments containing LoRA settings. train_pipeline_config: The train pipeline config. + on_meta: Base model params are on meta (--fsdp-load-mode stream); + create the adapters on meta too and defer their init to + reset_lora_adapters, after materialization. """ from peft import LoraConfig, get_peft_model @@ -707,6 +742,12 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - init_lora_weight = args.diffusion_init_lora_weight if init_lora_weight == "kaiming-uniform": init_lora_weight = True # namely kaiming-uniform + if on_meta and init_lora_weight not in (True, "gaussian"): + raise ValueError( + f"--diffusion-init-lora-weight {args.diffusion_init_lora_weight!r} is data-aware " + f"(needs real base weights at wrap time) and incompatible with --fsdp-load-mode " + f"stream; run with --fsdp-load-mode legacy" + ) model = get_peft_model( model, LoraConfig( @@ -715,12 +756,56 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - target_modules=targets, init_lora_weights=init_lora_weight, ), + low_cpu_mem_usage=on_meta, ) if dist.get_rank() == 0: model.print_trainable_parameters() return model +def reset_lora_adapters(model: torch.nn.Module, init_lora_weight) -> None: + """(Re)initialize LoRA A/B in place. Under --fsdp-load-mode stream the + adapters were created on meta, so PEFT's own init never ran on real + memory; to_empty left them uninitialized. nn.init works on the sharded + DTensor params directly (the RNG tracker keeps ranks coordinated).""" + from peft.tuners.lora import LoraLayer + + if init_lora_weight == "kaiming-uniform": + init_lora_weight = True + for module in model.modules(): + if isinstance(module, LoraLayer): + for adapter_name in module.lora_A: + module.reset_lora_parameters(adapter_name, init_lora_weight) + + +def peft_checkpoint_key_map(model: torch.nn.Module): + """Checkpoint keys are un-wrapped DiT names; after get_peft_model the live + FQNs gain ``base_model.model.`` and ``.base_layer``. Build the inverse of + the stripping done when pushing weights to sglang-d + (diffusion_update_weight_utils), so streamed checkpoints land on the + wrapped names.""" + mapping = {} + for name, _ in model.named_parameters(): + if "lora_" in name: + continue # adapter params don't come from the checkpoint + checkpoint_key = name.replace(".base_layer", "") + checkpoint_key = checkpoint_key.removeprefix("base_model.model.") + mapping[checkpoint_key] = name + return lambda key: mapping.get(key, key) + + +def materialize_sharded_model(model: torch.nn.Module, device) -> None: + """to_empty allocates this rank's shards but leaves every buffer + uninitialized. Persistent buffers are refilled by the checkpoint load + that follows; non-persistent ones (e.g. Wan's rope freq tables) exist + only as __init__'s output, so carry them across the transition.""" + real_buffers = {name: buffer.detach().clone() for name, buffer in model.named_buffers() if not buffer.is_meta} + model.to_empty(device=device) + for name, buffer in model.named_buffers(): + if name in real_buffers: + buffer.copy_(real_buffers[name]) + + def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index 6184d39e..1656cc9e 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/miles/backends/fsdp_utils/arguments.py @@ -45,6 +45,12 @@ class FSDPArgs: fp16: bool = False # FSDP configuration + # Weight init/loading strategy. "stream": build on meta, shard first, then + # rank0 streams safetensors files and broadcasts into the shards — peak GPU + # = shard size, only rank0 touches disk. "legacy": every rank materializes + # the full model on CPU and moves it to GPU before sharding — peak GPU = + # full model; escape hatch for custom pipelines / data-aware LoRA inits. + fsdp_load_mode: str = "stream" fsdp_state_dict_cpu_offload: bool = True # If True, offload full state dict to CPU during collection. fsdp_cpu_offload: bool = ( False # If True, offload parameters, gradients, and optimizer states to CPU (optimizer runs on CPU) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 04079032..79e49fa7 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -179,3 +179,11 @@ def cfg_combine( @abc.abstractmethod def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: """Preprocess the model before FSDP.""" + + def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: + """Called under ``--fsdp-load-mode stream`` once the sharded model is + materialized and pretrained weights are loaded — the earliest point + that path has real tensors on device. Default: run the before-fsdp + hook here, since its usual job (device-sensitive cache rebuilds, e.g. + qwen-image rope parity) would silently no-op on a meta-device model.""" + self.preprocess_model_before_fsdp(model) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 9bb8e500..459fae27 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -1,10 +1,15 @@ """Model backend: owns model-side behavior for the FSDP trainer. Selected via ``--model-backend-path`` (miles custom-function style); the -family config declares the default. Three concerns, all properties of the +family config declares the default. The concerns are all properties of the concrete modeling rather than of the training loop: - - ``load_models_and_scheduler``: checkpoint -> ``({component: model}, scheduler)`` + - ``build_models_and_scheduler``: checkpoint -> ``({component: model}, scheduler)`` + with parameters on the meta device (``--fsdp-load-mode stream``) + - ``stream_load_weights``: fill one already-sharded component's weights, + reading the checkpoint on rank 0 only + - ``load_models_and_scheduler``: legacy path — full weights on every rank + (``--fsdp-load-mode legacy``) - ``enable_gradient_checkpointing``: how this model turns on grad ckpt - ``fsdp_no_split_modules``: which block classes FSDP wraps @@ -15,11 +20,21 @@ from __future__ import annotations import abc +import json +import logging +import os from typing import Any +from collections.abc import Callable import torch +import torch.distributed as dist from diffusers import DiffusionPipeline +logger = logging.getLogger(__name__) + +_SAFETENSORS_INDEX = "diffusion_pytorch_model.safetensors.index.json" +_SAFETENSORS_SINGLE = "diffusion_pytorch_model.safetensors" + class ModelBackend(abc.ABC): def __init__(self, train_pipeline_config): @@ -32,7 +47,37 @@ def load_models_and_scheduler( *, master_dtype: torch.dtype, ) -> tuple[dict[str, torch.nn.Module], Any]: - """Return ``({component: model}, scheduler)`` on CPU.""" + """Legacy path: return ``({component: model}, scheduler)`` fully + materialized on CPU, on every rank.""" + + def build_models_and_scheduler( + self, + args, + *, + master_dtype: torch.dtype, + ) -> tuple[dict[str, torch.nn.Module], Any]: + """Return ``({component: model}, scheduler)`` with component + parameters on the meta device (buffers real, computed by __init__). + Weights are filled in later via ``stream_load_weights``, after FSDP + sharding, so no rank ever holds a full unsharded copy.""" + raise NotImplementedError( + f"{type(self).__name__} does not implement meta-device init; run with --fsdp-load-mode legacy" + ) + + def stream_load_weights( + self, + model: torch.nn.Module, + component: str, + args, + *, + master_dtype: torch.dtype, + key_map: Callable[[str], str] | None = None, + ) -> None: + """Fill ``model``'s weights from the checkpoint. ``model`` is already + FSDP-sharded and materialized; rank 0 reads, everyone receives.""" + raise NotImplementedError( + f"{type(self).__name__} does not implement streaming load; run with --fsdp-load-mode legacy" + ) def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None: """Turn on grad checkpointing; default = the diffusers protocol method.""" @@ -75,3 +120,117 @@ def load_models_and_scheduler( scheduler = pipeline.scheduler del pipeline return raw_models, scheduler + + def build_models_and_scheduler( + self, + args, + *, + master_dtype: torch.dtype, + ) -> tuple[dict[str, torch.nn.Module], Any]: + import diffusers + from accelerate import init_empty_weights + + root = self._resolve_checkpoint_dir(args) + with open(os.path.join(root, "model_index.json")) as f: + model_index = json.load(f) + + raw_models: dict[str, torch.nn.Module] = {} + for component in args.update_weight_target_modules: + if component not in model_index: + raise ValueError( + f"--update-weight-target-module: pipeline {args.hf_checkpoint} " f"has no component '{component}'" + ) + library, cls_name = model_index[component] + if library != "diffusers": + raise ValueError( + f"component '{component}' comes from library '{library}'; only diffusers " + f"components support meta init — run with --fsdp-load-mode legacy" + ) + cls = getattr(diffusers, cls_name, None) + if cls is None: + raise ValueError( + f"component class '{cls_name}' not found in diffusers (remote code?); " + f"run with --fsdp-load-mode legacy" + ) + config = cls.load_config(root, subfolder=component) + # include_buffers=False: parameters land on meta (zero memory), but + # buffers and plain tensor attributes are computed for real by + # __init__ (Wan's rope freq tables, Qwen's pos/neg freqs). This is + # the same context diffusers' own from_pretrained builds under, so + # buffer values match the legacy path bit-exactly. + with init_empty_weights(include_buffers=False): + model = cls.from_config(config) + raw_models[component] = model.to(master_dtype) + + _, scheduler_cls_name = model_index["scheduler"] + scheduler = getattr(diffusers, scheduler_cls_name).from_pretrained(root, subfolder="scheduler") + return raw_models, scheduler + + def stream_load_weights( + self, + model: torch.nn.Module, + component: str, + args, + *, + master_dtype: torch.dtype, + key_map: Callable[[str], str] | None = None, + ) -> None: + from safetensors.torch import load_file + from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict + + component_dir = os.path.join(self._resolve_checkpoint_dir(args), component) + index_path = os.path.join(component_dir, _SAFETENSORS_INDEX) + rank = dist.get_rank() + if rank == 0: + if os.path.exists(index_path): + with open(index_path) as f: + files = sorted(set(json.load(f)["weight_map"].values())) + elif os.path.exists(os.path.join(component_dir, _SAFETENSORS_SINGLE)): + files = [_SAFETENSORS_SINGLE] + else: + raise FileNotFoundError(f"no safetensors checkpoint under {component_dir}") + else: + files = None + # rank0 discovered the file list; everyone must loop the same number of + # times since set_model_state_dict is collective. + obj = [files] + dist.broadcast_object_list(obj, src=0) + files = obj[0] + + options = StateDictOptions(full_state_dict=True, broadcast_from_rank0=True, strict=False) + loaded_keys: set[str] = set() + for fname in files: + partial: dict[str, torch.Tensor] = {} + if rank == 0: + # load_file mmaps; peak CPU ~= one file (+ a dtype-cast copy) + for key, tensor in load_file(os.path.join(component_dir, fname)).items(): + partial[key_map(key) if key_map is not None else key] = tensor.to(master_dtype) + set_model_state_dict(model, partial, options=options) + key_obj = [set(partial.keys())] + dist.broadcast_object_list(key_obj, src=0) + loaded_keys |= key_obj[0] + del partial + + # strict=False enables the per-file partial loads above, so missing + # weights would pass silently — verify full coverage ourselves. + expected = {name for name, _ in model.named_parameters() if "lora_" not in name} + missing = expected - loaded_keys + if missing: + raise RuntimeError( + f"{component}: {len(missing)} parameters not covered by checkpoint, " f"e.g. {sorted(missing)[:3]}" + ) + logger.info("[FSDP] %s: streamed %d files, %d tensors", component, len(files), len(loaded_keys)) + + @staticmethod + def _resolve_checkpoint_dir(args) -> str: + ckpt = args.hf_checkpoint + if os.path.isdir(ckpt): + return ckpt + from huggingface_hub import snapshot_download + + # every rank needs the configs/index; weight files only on rank 0, + # where stream_load_weights reads them. + patterns = ["model_index.json", "*/*.json"] + if not dist.is_initialized() or dist.get_rank() == 0: + patterns += [f"{component}/*" for component in args.update_weight_target_modules] + return snapshot_download(ckpt, allow_patterns=patterns) diff --git a/tests/fast-gpu/test_fsdp_stream_load_equivalence.py b/tests/fast-gpu/test_fsdp_stream_load_equivalence.py new file mode 100644 index 00000000..080d753e --- /dev/null +++ b/tests/fast-gpu/test_fsdp_stream_load_equivalence.py @@ -0,0 +1,160 @@ +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=180, + suite="stage-b-3-gpu-h200", + labels=[], +) + +import json +import os +from argparse import Namespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +WORLD_SIZE = 2 + +_TINY_WAN_CONFIG = { + "_class_name": "WanTransformer3DModel", + "attention_head_dim": 8, + "num_attention_heads": 2, + "num_layers": 2, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 16, + "ffn_dim": 32, + "patch_size": [1, 2, 2], + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "eps": 1e-6, +} + + +def _make_checkpoint(root: str) -> None: + """Tiny diffusers-layout Wan checkpoint with a SHARDED transformer, so the + stream loader's multi-file path is exercised.""" + from diffusers import WanTransformer3DModel + + torch.manual_seed(7) + model = WanTransformer3DModel.from_config(_TINY_WAN_CONFIG) + model.save_pretrained(os.path.join(root, "transformer"), max_shard_size="20KB") + index = os.path.join(root, "transformer", "diffusion_pytorch_model.safetensors.index.json") + assert os.path.exists(index), "tiny checkpoint must be sharded for this test" + with open(os.path.join(root, "model_index.json"), "w") as f: + json.dump( + { + "_class_name": "WanPipeline", + "scheduler": ["diffusers", "UniPCMultistepScheduler"], + "transformer": ["diffusers", "WanTransformer3DModel"], + }, + f, + ) + os.makedirs(os.path.join(root, "scheduler"), exist_ok=True) + with open(os.path.join(root, "scheduler", "scheduler_config.json"), "w") as f: + json.dump({"_class_name": "UniPCMultistepScheduler", "num_train_timesteps": 1000}, f) + + +def _fsdp_args(**overrides) -> Namespace: + base = dict( + hf_checkpoint=None, + update_weight_target_modules=["transformer"], + diffusion_forward_dtype="bf16", + fsdp_reduce_dtype="fp32", + gradient_checkpointing=False, + lora_target_modules=["to_q", "to_k", "to_v"], + diffusion_init_lora_weight="gaussian", + lora_rank=4, + lora_alpha=4, + ) + base.update(overrides) + return Namespace(**base) + + +def _stream_load(ckpt: str, use_lora: bool): + from miles.backends.fsdp_utils.actor import ( + apply_fsdp2, + apply_lora, + materialize_sharded_model, + peft_checkpoint_key_map, + reset_lora_adapters, + ) + from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend + + args = _fsdp_args(hf_checkpoint=ckpt) + backend = DiffusersModelBackend(None) + raw_models, _ = backend.build_models_and_scheduler(args, master_dtype=torch.float32) + model = raw_models["transformer"] + key_map = None + if use_lora: + model = apply_lora(model, args, None, on_meta=True) + key_map = peft_checkpoint_key_map(model) + model = apply_fsdp2(model, args=args) + materialize_sharded_model(model, torch.cuda.current_device()) + backend.stream_load_weights(model, "transformer", args, master_dtype=torch.float32, key_map=key_map) + if use_lora: + reset_lora_adapters(model, args.diffusion_init_lora_weight) + return model + + +def _worker(rank: int, ckpt: str, use_lora: bool): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29531" + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=WORLD_SIZE) + try: + from diffusers import WanTransformer3DModel + + model = _stream_load(ckpt, use_lora) + + # reference: plain full from_pretrained — what the legacy path + # materializes per component before sharding + ref = WanTransformer3DModel.from_pretrained(ckpt, subfolder="transformer", torch_dtype=torch.float32) + ref_sd = ref.state_dict() + + checked = 0 + for name, param in model.state_dict().items(): + if "lora_" in name: + continue + ref_key = name.replace(".base_layer", "").removeprefix("base_model.model.") + full = param.full_tensor() if hasattr(param, "full_tensor") else param + assert torch.equal(full.cpu(), ref_sd[ref_key]), f"mismatch: {name}" + checked += 1 + assert checked == len(ref_sd), f"covered {checked}/{len(ref_sd)} reference tensors" + + # non-persistent rope tables are not in any state_dict: they must + # survive to_empty via the snapshot in materialize_sharded_model + buffers = dict(model.named_buffers()) + ref_buffers = dict(ref.named_buffers()) + for name, ref_buffer in ref_buffers.items(): + live = next(b for n, b in buffers.items() if n.endswith(name)) + assert torch.equal(live.cpu(), ref_buffer), f"buffer mismatch: {name}" + + if use_lora: + from peft.tuners.lora import LoraLayer + + lora_layers = [m for m in model.modules() if isinstance(m, LoraLayer)] + assert lora_layers, "no LoRA layers injected" + for layer in lora_layers: + a = layer.lora_A["default"].weight + b = layer.lora_B["default"].weight + a = a.full_tensor() if hasattr(a, "full_tensor") else a + b = b.full_tensor() if hasattr(b, "full_tensor") else b + assert torch.isfinite(a).all() and a.abs().sum() > 0, "lora_A not initialized" + assert (b == 0).all(), "lora_B must start at zero" + finally: + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize("use_lora", [False, True], ids=["base", "lora"]) +def test_stream_load_matches_from_pretrained(tmp_path, use_lora): + if torch.cuda.device_count() < WORLD_SIZE: + pytest.skip(f"needs {WORLD_SIZE} GPUs") + ckpt = str(tmp_path / "ckpt") + os.makedirs(ckpt) + _make_checkpoint(ckpt) + mp.spawn(_worker, args=(ckpt, use_lora), nprocs=WORLD_SIZE, join=True) diff --git a/tests/fast/backends/fsdp_utils/test_stream_load.py b/tests/fast/backends/fsdp_utils/test_stream_load.py new file mode 100644 index 00000000..49063b04 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_stream_load.py @@ -0,0 +1,164 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) + +import json +from argparse import Namespace + +import pytest +import torch + +from miles.backends.fsdp_utils.actor import ( + apply_lora, + materialize_sharded_model, + peft_checkpoint_key_map, + reset_lora_adapters, +) +from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend, ModelBackend + + +def _tiny_wan_checkpoint(tmp_path): + """A minimal diffusers-layout checkpoint dir: model_index.json + a 1-block + WanTransformer3DModel config + scheduler config. No weight files — enough + for build_models_and_scheduler, which must not read any.""" + config = { + "_class_name": "WanTransformer3DModel", + "attention_head_dim": 8, + "num_attention_heads": 2, + "num_layers": 1, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 16, + "ffn_dim": 32, + "patch_size": [1, 2, 2], + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "eps": 1e-6, + } + model_index = { + "_class_name": "WanPipeline", + "scheduler": ["diffusers", "UniPCMultistepScheduler"], + "transformer": ["diffusers", "WanTransformer3DModel"], + "text_encoder": ["transformers", "UMT5EncoderModel"], + } + scheduler_config = {"_class_name": "UniPCMultistepScheduler", "num_train_timesteps": 1000} + (tmp_path / "transformer").mkdir() + (tmp_path / "scheduler").mkdir() + (tmp_path / "model_index.json").write_text(json.dumps(model_index)) + (tmp_path / "transformer" / "config.json").write_text(json.dumps(config)) + (tmp_path / "scheduler" / "scheduler_config.json").write_text(json.dumps(scheduler_config)) + return tmp_path + + +class TestBuildModelsAndScheduler: + # The core contract: params on meta (no weight I/O), buffers real (computed + # by __init__) — Wan's non-persistent rope tables are the regression case. + def test_params_meta_buffers_real(self, tmp_path): + ckpt = _tiny_wan_checkpoint(tmp_path) + args = Namespace(hf_checkpoint=str(ckpt), update_weight_target_modules=["transformer"]) + raw_models, scheduler = DiffusersModelBackend(None).build_models_and_scheduler( + args, master_dtype=torch.float32 + ) + model = raw_models["transformer"] + assert all(p.is_meta for p in model.parameters()) + buffers = dict(model.named_buffers()) + assert "rope.freqs_cos" in buffers and not buffers["rope.freqs_cos"].is_meta + assert type(scheduler).__name__ == "UniPCMultistepScheduler" + + def test_missing_component_raises(self, tmp_path): + ckpt = _tiny_wan_checkpoint(tmp_path) + args = Namespace(hf_checkpoint=str(ckpt), update_weight_target_modules=["vae"]) + with pytest.raises(ValueError, match="has no component"): + DiffusersModelBackend(None).build_models_and_scheduler(args, master_dtype=torch.float32) + + # transformers-library components (text_encoder) can't meta-init through + # diffusers class resolution — must point at the legacy escape hatch. + def test_non_diffusers_component_raises(self, tmp_path): + ckpt = _tiny_wan_checkpoint(tmp_path) + args = Namespace(hf_checkpoint=str(ckpt), update_weight_target_modules=["text_encoder"]) + with pytest.raises(ValueError, match="legacy"): + DiffusersModelBackend(None).build_models_and_scheduler(args, master_dtype=torch.float32) + + def test_custom_backend_without_meta_support_points_to_legacy(self): + class _LegacyOnly(ModelBackend): + def load_models_and_scheduler(self, args, *, master_dtype): + raise NotImplementedError + + with pytest.raises(NotImplementedError, match="fsdp-load-mode legacy"): + _LegacyOnly(None).build_models_and_scheduler(None, master_dtype=torch.float32) + + +class TestMaterializeShardedModel: + # to_empty wipes buffers; non-persistent ones aren't in any checkpoint, so + # materialize must carry them over (Wan rope). Persistent ones may be + # garbage afterwards — the stream load refills them. + def test_non_persistent_buffers_survive(self): + class _M(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + self.register_buffer("freqs", torch.arange(4, dtype=torch.float32), persistent=False) + + from accelerate import init_empty_weights + + with init_empty_weights(include_buffers=False): + m = _M() + assert m.linear.weight.is_meta and not m.freqs.is_meta + materialize_sharded_model(m, torch.device("cpu")) + assert not m.linear.weight.is_meta + torch.testing.assert_close(m.freqs, torch.arange(4, dtype=torch.float32)) + + +class _TinyDit(torch.nn.Module): + def __init__(self): + super().__init__() + self.to_q = torch.nn.Linear(8, 8, bias=False) + self.norm = torch.nn.LayerNorm(8) + + +def _lora_args(init_weight="gaussian"): + return Namespace( + lora_target_modules=["to_q"], + diffusion_init_lora_weight=init_weight, + lora_rank=4, + lora_alpha=4, + ) + + +class TestLoraStreamPath: + # Key map must invert the FQN wrapping so raw checkpoint keys land on the + # peft-wrapped names — same string contract as diffusion_update_weight_utils. + def test_key_map_inverts_peft_wrapping(self, monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "get_rank", lambda: 1) + with torch.device("meta"): + model = _TinyDit() + model = apply_lora(model, _lora_args(), None, on_meta=True) + key_map = peft_checkpoint_key_map(model) + assert key_map("to_q.weight") == "base_model.model.to_q.base_layer.weight" + assert key_map("norm.weight") == "base_model.model.norm.weight" + assert key_map("unknown.weight") == "unknown.weight" # passthrough + + def test_adapters_created_on_meta_then_reset(self, monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "get_rank", lambda: 1) + with torch.device("meta"): + model = _TinyDit() + model = apply_lora(model, _lora_args(), None, on_meta=True) + lora_layer = model.base_model.model.to_q + assert lora_layer.lora_A["default"].weight.is_meta + model.to_empty(device="cpu") + reset_lora_adapters(model, "gaussian") + assert torch.isfinite(lora_layer.lora_A["default"].weight).all() + assert (lora_layer.lora_B["default"].weight == 0).all() + + # pissa/olora-style inits read the base weights at wrap time; on meta they + # would silently produce garbage — must be an explicit error. + def test_data_aware_init_rejected_on_meta(self): + with torch.device("meta"): + model = _TinyDit() + with pytest.raises(ValueError, match="legacy"): + apply_lora(model, _lora_args(init_weight="pissa"), None, on_meta=True) From 0686e072786b719a22debba838bfe928247358ad Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 17 Jul 2026 00:09:27 +0000 Subject: [PATCH 02/10] refactor FSDP component loading --- miles/backends/fsdp_utils/actor.py | 183 ++++++------- miles/backends/fsdp_utils/arguments.py | 6 - .../backends/fsdp_utils/configs/qwen_image.py | 10 +- miles/backends/fsdp_utils/configs/sd3.py | 3 - .../configs/train_pipeline_config.py | 11 +- miles/backends/fsdp_utils/configs/wan2_2.py | 3 - miles/backends/fsdp_utils/model_backend.py | 257 ++++-------------- .../test_fsdp_stream_load_equivalence.py | 160 ----------- .../backends/fsdp_utils/test_stream_load.py | 164 ----------- 9 files changed, 140 insertions(+), 657 deletions(-) delete mode 100644 tests/fast-gpu/test_fsdp_stream_load_equivalence.py delete mode 100644 tests/fast/backends/fsdp_utils/test_stream_load.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 29ad6770..09322a5c 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,8 +1,9 @@ import functools import logging +import warnings from argparse import Namespace from collections import defaultdict -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext import ray import torch @@ -109,66 +110,43 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.train_pipeline_config = load_function(args.train_pipeline_config_path)() self.model_backend = load_function(args.model_backend_path)(self.train_pipeline_config) - if args.fsdp_load_mode not in ("stream", "legacy"): - raise ValueError(f"--fsdp-load-mode must be 'stream' or 'legacy', got {args.fsdp_load_mode!r}") - stream_load = args.fsdp_load_mode == "stream" - if stream_load: - # params on meta: no rank ever holds a full unsharded copy - raw_models, self.scheduler = self.model_backend.build_models_and_scheduler( - args, master_dtype=self._master_dtype - ) - else: - raw_models, self.scheduler = self.model_backend.load_models_and_scheduler( - args, master_dtype=self._master_dtype - ) + self.scheduler = self.model_backend.load_scheduler(args) + init_context = self._get_init_weight_context_manager() self.models: dict[str, torch.nn.Module] = {} - for component, model in raw_models.items(): - # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap + for component in args.update_weight_target_modules: + with init_context(): + model = self.model_backend.load_component(component, args, master_dtype=self._master_dtype) if args.fsdp_attention_backend is not None: self.model_backend.set_attention_backend(model, args.fsdp_attention_backend) - key_map = None if args.use_lora: - model = apply_lora(model, args, self.train_pipeline_config, on_meta=stream_load) - if stream_load: - key_map = peft_checkpoint_key_map(model) + model = apply_lora(model, args, self.train_pipeline_config, on_meta=dist.get_rank() != 0) model.train() if args.gradient_checkpointing: self.model_backend.enable_gradient_checkpointing(model) - if stream_load: - # shard while still on meta (costs nothing), materialize only - # this rank's shards, then fill weights file-by-file from rank0. - model = apply_fsdp2( - model, - mesh=self.parallel_state.dp_mesh, - cpu_offload=self.args.fsdp_cpu_offload, - args=self.args, - no_split_modules=self.model_backend.fsdp_no_split_modules(model), - ) - device = torch.device("cpu") if self.args.fsdp_cpu_offload else torch.cuda.current_device() - materialize_sharded_model(model, device) - self.model_backend.stream_load_weights( - model, component, args, master_dtype=self._master_dtype, key_map=key_map - ) - if args.use_lora: - reset_lora_adapters(model, args.diffusion_init_lora_weight) - self.train_pipeline_config.postprocess_model_after_materialize(model) - else: - model.to(torch.cuda.current_device()) - - self.train_pipeline_config.preprocess_model_before_fsdp(model) - - model = apply_fsdp2( - model, - mesh=self.parallel_state.dp_mesh, - cpu_offload=self.args.fsdp_cpu_offload, - args=self.args, - no_split_modules=self.model_backend.fsdp_no_split_modules(model), - ) + if dist.get_rank() != 0 and any(not parameter.is_meta for parameter in model.parameters()): + raise RuntimeError(f"{component} did not honor meta initialization") + sync_model_dtypes(model) + full_state = model.state_dict() + model = apply_fsdp2( + model, + mesh=self.parallel_state.dp_mesh, + cpu_offload=self.args.fsdp_cpu_offload, + args=self.args, + no_split_modules=self.model_backend.fsdp_no_split_modules(model), + ) + device = ( + torch.device("cpu") + if self.args.fsdp_cpu_offload + else torch.device("cuda", torch.cuda.current_device()) + ) + load_sharded_model(model, full_state, device) + del full_state + self.train_pipeline_config.postprocess_model_after_materialize(model) self.models[component] = model # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -232,6 +210,23 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty return self.args.start_rollout_id + def _get_init_weight_context_manager(self): + from accelerate import init_empty_weights + + if dist.get_rank() == 0: + return lambda: torch.device("cpu") + + @contextmanager + def meta_init_weights(): + with init_empty_weights(include_buffers=False), warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"for .*: copying from a non-meta parameter in the checkpoint to a meta parameter.*", + ) + yield + + return meta_init_weights + def _get_parallel_config(self) -> dict: return {"dp_size": getattr(self.parallel_state, "dp_size", 1)} @@ -724,37 +719,22 @@ def _resolve_dtype(name: str) -> torch.dtype: return {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[name] -def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config, on_meta: bool = False) -> None: - """Apply PEFT LoRA to the model. - - Args: - model: The model to apply LoRA to. - args: Arguments containing LoRA settings. - train_pipeline_config: The train pipeline config. - on_meta: Base model params are on meta (--fsdp-load-mode stream); - create the adapters on meta too and defer their init to - reset_lora_adapters, after materialization. - """ +def apply_lora( + model: torch.nn.Module, args: Namespace, train_pipeline_config, on_meta: bool = False +) -> torch.nn.Module: from peft import LoraConfig, get_peft_model - # Per-model fallback when --lora-target-modules is unset (runtime inference: depends on loaded pipeline). targets = args.lora_target_modules or train_pipeline_config.lora_target_modules init_lora_weight = args.diffusion_init_lora_weight if init_lora_weight == "kaiming-uniform": - init_lora_weight = True # namely kaiming-uniform - if on_meta and init_lora_weight not in (True, "gaussian"): - raise ValueError( - f"--diffusion-init-lora-weight {args.diffusion_init_lora_weight!r} is data-aware " - f"(needs real base weights at wrap time) and incompatible with --fsdp-load-mode " - f"stream; run with --fsdp-load-mode legacy" - ) + init_lora_weight = True model = get_peft_model( model, LoraConfig( r=args.lora_rank, lora_alpha=args.lora_alpha, target_modules=targets, - init_lora_weights=init_lora_weight, + init_lora_weights=False if on_meta else init_lora_weight, ), low_cpu_mem_usage=on_meta, ) @@ -763,47 +743,42 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config, o return model -def reset_lora_adapters(model: torch.nn.Module, init_lora_weight) -> None: - """(Re)initialize LoRA A/B in place. Under --fsdp-load-mode stream the - adapters were created on meta, so PEFT's own init never ran on real - memory; to_empty left them uninitialized. nn.init works on the sharded - DTensor params directly (the RNG tracker keeps ranks coordinated).""" - from peft.tuners.lora import LoraLayer +def load_sharded_model(model: torch.nn.Module, full_state: dict, device) -> None: + from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict - if init_lora_weight == "kaiming-uniform": - init_lora_weight = True - for module in model.modules(): - if isinstance(module, LoraLayer): - for adapter_name in module.lora_A: - module.reset_lora_parameters(adapter_name, init_lora_weight) - - -def peft_checkpoint_key_map(model: torch.nn.Module): - """Checkpoint keys are un-wrapped DiT names; after get_peft_model the live - FQNs gain ``base_model.model.`` and ``.base_layer``. Build the inverse of - the stripping done when pushing weights to sglang-d - (diffusion_update_weight_utils), so streamed checkpoints land on the - wrapped names.""" - mapping = {} - for name, _ in model.named_parameters(): - if "lora_" in name: - continue # adapter params don't come from the checkpoint - checkpoint_key = name.replace(".base_layer", "") - checkpoint_key = checkpoint_key.removeprefix("base_model.model.") - mapping[checkpoint_key] = name - return lambda key: mapping.get(key, key) - - -def materialize_sharded_model(model: torch.nn.Module, device) -> None: - """to_empty allocates this rank's shards but leaves every buffer - uninitialized. Persistent buffers are refilled by the checkpoint load - that follows; non-persistent ones (e.g. Wan's rope freq tables) exist - only as __init__'s output, so carry them across the transition.""" real_buffers = {name: buffer.detach().clone() for name, buffer in model.named_buffers() if not buffer.is_meta} model.to_empty(device=device) for name, buffer in model.named_buffers(): if name in real_buffers: buffer.copy_(real_buffers[name]) + set_model_state_dict( + model, + full_state, + options=StateDictOptions( + full_state_dict=True, + cpu_offload=device.type == "cpu", + broadcast_from_rank0=True, + strict=True, + ), + ) + + +def sync_model_dtypes(model: torch.nn.Module) -> None: + dtypes = None + if dist.get_rank() == 0: + dtypes = { + "parameters": {name: parameter.dtype for name, parameter in model.named_parameters()}, + "buffers": {name: buffer.dtype for name, buffer in model.named_buffers()}, + } + objects = [dtypes] + dist.broadcast_object_list(objects, src=0) + dtypes = objects[0] + if dist.get_rank() == 0: + return + for name, parameter in model.named_parameters(): + parameter.data = parameter.data.to(dtypes["parameters"][name]) + for name, buffer in model.named_buffers(): + buffer.data = buffer.data.to(dtypes["buffers"][name]) def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None): diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index 1656cc9e..6184d39e 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/miles/backends/fsdp_utils/arguments.py @@ -45,12 +45,6 @@ class FSDPArgs: fp16: bool = False # FSDP configuration - # Weight init/loading strategy. "stream": build on meta, shard first, then - # rank0 streams safetensors files and broadcasts into the shards — peak GPU - # = shard size, only rank0 touches disk. "legacy": every rank materializes - # the full model on CPU and moves it to GPU before sharding — peak GPU = - # full model; escape hatch for custom pipelines / data-aware LoRA inits. - fsdp_load_mode: str = "stream" fsdp_state_dict_cpu_offload: bool = True # If True, offload full state dict to CPU during collection. fsdp_cpu_offload: bool = ( False # If True, offload parameters, gradients, and optimizer states to CPU (optimizer runs on CPU) diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index a2946b16..8e69855d 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -22,12 +22,7 @@ def _rebuild_pos_embed_freqs_on_cuda(model) -> None: RoPE output differs → every block's output drifts → frozen-weight ``noise_pred`` mean|Δ| ~2e-02. """ - try: - device = next(model.parameters()).device - except StopIteration: - return - if device.type != "cuda": - return + device = torch.device("cuda", torch.cuda.current_device()) for submod in model.modules(): # Match by attribute shape rather than class name so we also # handle ``QwenEmbedLayer3DRope`` and similar variants @@ -190,6 +185,5 @@ def cfg_combine( combined = combined * (pos_norm / combined_norm) return combined - def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: - """Preprocess the model before FSDP.""" + def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: _rebuild_pos_embed_freqs_on_cuda(model) diff --git a/miles/backends/fsdp_utils/configs/sd3.py b/miles/backends/fsdp_utils/configs/sd3.py index 5673e1cf..332e9fa1 100644 --- a/miles/backends/fsdp_utils/configs/sd3.py +++ b/miles/backends/fsdp_utils/configs/sd3.py @@ -70,6 +70,3 @@ def cfg_combine( ) -> torch.Tensor: scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) - - def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: - return None diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 79e49fa7..221acd19 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -176,14 +176,5 @@ def cfg_combine( ) -> torch.Tensor: """Apply classifier-free guidance. Model-specific (e.g. rescale or not).""" - @abc.abstractmethod - def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: - """Preprocess the model before FSDP.""" - def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: - """Called under ``--fsdp-load-mode stream`` once the sharded model is - materialized and pretrained weights are loaded — the earliest point - that path has real tensors on device. Default: run the before-fsdp - hook here, since its usual job (device-sensitive cache rebuilds, e.g. - qwen-image rope parity) would silently no-op on a meta-device model.""" - self.preprocess_model_before_fsdp(model) + return None diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index c570173e..7e227cd2 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -66,6 +66,3 @@ def cfg_combine( ) -> torch.Tensor: scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) - - def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: - return None diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 459fae27..9eb80967 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -1,236 +1,95 @@ -"""Model backend: owns model-side behavior for the FSDP trainer. - -Selected via ``--model-backend-path`` (miles custom-function style); the -family config declares the default. The concerns are all properties of the -concrete modeling rather than of the training loop: - - - ``build_models_and_scheduler``: checkpoint -> ``({component: model}, scheduler)`` - with parameters on the meta device (``--fsdp-load-mode stream``) - - ``stream_load_weights``: fill one already-sharded component's weights, - reading the checkpoint on rank 0 only - - ``load_models_and_scheduler``: legacy path — full weights on every rank - (``--fsdp-load-mode legacy``) - - ``enable_gradient_checkpointing``: how this model turns on grad ckpt - - ``fsdp_no_split_modules``: which block classes FSDP wraps - -Defaults implement the diffusers protocol (see ``models/__init__.py``); a -native model overrides methods here instead of retrofitting its instances. -""" - from __future__ import annotations -import abc -import json -import logging -import os +import importlib from typing import Any -from collections.abc import Callable import torch import torch.distributed as dist from diffusers import DiffusionPipeline -logger = logging.getLogger(__name__) -_SAFETENSORS_INDEX = "diffusion_pytorch_model.safetensors.index.json" -_SAFETENSORS_SINGLE = "diffusion_pytorch_model.safetensors" - - -class ModelBackend(abc.ABC): +class ModelBackend: def __init__(self, train_pipeline_config): self.config = train_pipeline_config - @abc.abstractmethod - def load_models_and_scheduler( - self, - args, - *, - master_dtype: torch.dtype, - ) -> tuple[dict[str, torch.nn.Module], Any]: - """Legacy path: return ``({component: model}, scheduler)`` fully - materialized on CPU, on every rank.""" - - def build_models_and_scheduler( - self, - args, - *, - master_dtype: torch.dtype, - ) -> tuple[dict[str, torch.nn.Module], Any]: - """Return ``({component: model}, scheduler)`` with component - parameters on the meta device (buffers real, computed by __init__). - Weights are filled in later via ``stream_load_weights``, after FSDP - sharding, so no rank ever holds a full unsharded copy.""" - raise NotImplementedError( - f"{type(self).__name__} does not implement meta-device init; run with --fsdp-load-mode legacy" - ) - - def stream_load_weights( + def load_component( self, - model: torch.nn.Module, component: str, args, *, master_dtype: torch.dtype, - key_map: Callable[[str], str] | None = None, - ) -> None: - """Fill ``model``'s weights from the checkpoint. ``model`` is already - FSDP-sharded and materialized; rank 0 reads, everyone receives.""" - raise NotImplementedError( - f"{type(self).__name__} does not implement streaming load; run with --fsdp-load-mode legacy" - ) + ) -> torch.nn.Module: + raise NotImplementedError + + def load_scheduler(self, args) -> Any: + raise NotImplementedError def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None: - """Turn on grad checkpointing; default = the diffusers protocol method.""" model.enable_gradient_checkpointing() def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]: - """Block class names FSDP wraps; default = the model's own declaration.""" return model._no_split_modules def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None: - """Select the DiT attention backend; default = the diffusers protocol method.""" model.set_attention_backend(backend) class DiffusersModelBackend(ModelBackend): - """Load trainable components from a diffusers pipeline checkpoint.""" - - def load_models_and_scheduler( + def load_component( self, + component: str, args, *, master_dtype: torch.dtype, - ) -> tuple[dict[str, torch.nn.Module], Any]: - pipeline = DiffusionPipeline.from_pretrained( - args.hf_checkpoint, - torch_dtype=master_dtype, - trust_remote_code=True, - text_encoder=None, - vae=None, - tokenizer=None, - ) - raw_models: dict[str, torch.nn.Module] = {} - for component in args.update_weight_target_modules: - sub_model = getattr(pipeline, component, None) - if sub_model is None: - raise ValueError( - f"--update-weight-target-module: pipeline {args.hf_checkpoint} " f"has no component '{component}'" - ) - raw_models[component] = sub_model - scheduler = pipeline.scheduler + ) -> torch.nn.Module: + config = DiffusionPipeline.load_config(args.hf_checkpoint) + if component not in config: + raise ValueError(f"pipeline {args.hf_checkpoint} has no component {component!r}") + + pipeline = self._load_pipeline(args, config, component, master_dtype) + model = getattr(pipeline, component) del pipeline - return raw_models, scheduler + return model - def build_models_and_scheduler( - self, - args, - *, - master_dtype: torch.dtype, - ) -> tuple[dict[str, torch.nn.Module], Any]: - import diffusers - from accelerate import init_empty_weights - - root = self._resolve_checkpoint_dir(args) - with open(os.path.join(root, "model_index.json")) as f: - model_index = json.load(f) - - raw_models: dict[str, torch.nn.Module] = {} - for component in args.update_weight_target_modules: - if component not in model_index: - raise ValueError( - f"--update-weight-target-module: pipeline {args.hf_checkpoint} " f"has no component '{component}'" - ) - library, cls_name = model_index[component] - if library != "diffusers": - raise ValueError( - f"component '{component}' comes from library '{library}'; only diffusers " - f"components support meta init — run with --fsdp-load-mode legacy" - ) - cls = getattr(diffusers, cls_name, None) - if cls is None: - raise ValueError( - f"component class '{cls_name}' not found in diffusers (remote code?); " - f"run with --fsdp-load-mode legacy" - ) - config = cls.load_config(root, subfolder=component) - # include_buffers=False: parameters land on meta (zero memory), but - # buffers and plain tensor attributes are computed for real by - # __init__ (Wan's rope freq tables, Qwen's pos/neg freqs). This is - # the same context diffusers' own from_pretrained builds under, so - # buffer values match the legacy path bit-exactly. - with init_empty_weights(include_buffers=False): - model = cls.from_config(config) - raw_models[component] = model.to(master_dtype) - - _, scheduler_cls_name = model_index["scheduler"] - scheduler = getattr(diffusers, scheduler_cls_name).from_pretrained(root, subfolder="scheduler") - return raw_models, scheduler - - def stream_load_weights( - self, - model: torch.nn.Module, - component: str, - args, - *, - master_dtype: torch.dtype, - key_map: Callable[[str], str] | None = None, - ) -> None: - from safetensors.torch import load_file - from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict - - component_dir = os.path.join(self._resolve_checkpoint_dir(args), component) - index_path = os.path.join(component_dir, _SAFETENSORS_INDEX) - rank = dist.get_rank() - if rank == 0: - if os.path.exists(index_path): - with open(index_path) as f: - files = sorted(set(json.load(f)["weight_map"].values())) - elif os.path.exists(os.path.join(component_dir, _SAFETENSORS_SINGLE)): - files = [_SAFETENSORS_SINGLE] - else: - raise FileNotFoundError(f"no safetensors checkpoint under {component_dir}") - else: - files = None - # rank0 discovered the file list; everyone must loop the same number of - # times since set_model_state_dict is collective. - obj = [files] - dist.broadcast_object_list(obj, src=0) - files = obj[0] - - options = StateDictOptions(full_state_dict=True, broadcast_from_rank0=True, strict=False) - loaded_keys: set[str] = set() - for fname in files: - partial: dict[str, torch.Tensor] = {} - if rank == 0: - # load_file mmaps; peak CPU ~= one file (+ a dtype-cast copy) - for key, tensor in load_file(os.path.join(component_dir, fname)).items(): - partial[key_map(key) if key_map is not None else key] = tensor.to(master_dtype) - set_model_state_dict(model, partial, options=options) - key_obj = [set(partial.keys())] - dist.broadcast_object_list(key_obj, src=0) - loaded_keys |= key_obj[0] - del partial - - # strict=False enables the per-file partial loads above, so missing - # weights would pass silently — verify full coverage ourselves. - expected = {name for name, _ in model.named_parameters() if "lora_" not in name} - missing = expected - loaded_keys - if missing: - raise RuntimeError( - f"{component}: {len(missing)} parameters not covered by checkpoint, " f"e.g. {sorted(missing)[:3]}" - ) - logger.info("[FSDP] %s: streamed %d files, %d tensors", component, len(files), len(loaded_keys)) + def load_scheduler(self, args) -> Any: + config = DiffusionPipeline.load_config(args.hf_checkpoint) + pipeline = self._load_pipeline(args, config, "scheduler") + scheduler = pipeline.scheduler + del pipeline + return scheduler + + @classmethod + def _load_pipeline(cls, args, config: dict, component: str, master_dtype=None): + skipped = { + name: None for name, value in config.items() if name != component and isinstance(value, (list, tuple)) + } + kwargs = {"trust_remote_code": True, "low_cpu_mem_usage": dist.get_rank() == 0, **skipped} + if master_dtype is not None: + kwargs["torch_dtype"] = master_dtype + + model_cls = cls._component_class(config.get(component)) + keep_in_fp32 = getattr(model_cls, "_keep_in_fp32_modules", None) + if dist.get_rank() != 0 and keep_in_fp32 is not None: + # Diffusers otherwise forces low_cpu_mem_usage and materializes meta ranks. + model_cls._keep_in_fp32_modules = None + try: + return DiffusionPipeline.from_pretrained(args.hf_checkpoint, **kwargs) + finally: + if dist.get_rank() != 0 and keep_in_fp32 is not None: + model_cls._keep_in_fp32_modules = keep_in_fp32 @staticmethod - def _resolve_checkpoint_dir(args) -> str: - ckpt = args.hf_checkpoint - if os.path.isdir(ckpt): - return ckpt - from huggingface_hub import snapshot_download - - # every rank needs the configs/index; weight files only on rank 0, - # where stream_load_weights reads them. - patterns = ["model_index.json", "*/*.json"] - if not dist.is_initialized() or dist.get_rank() == 0: - patterns += [f"{component}/*" for component in args.update_weight_target_modules] - return snapshot_download(ckpt, allow_patterns=patterns) + def _component_class(spec): + if not isinstance(spec, (list, tuple)) or len(spec) != 2: + return None + library, class_name = spec + if not library or not class_name: + return None + try: + module = importlib.import_module(library) + except ImportError: + try: + module = importlib.import_module(f"diffusers.pipelines.{library}") + except ImportError: + return None + return getattr(module, class_name, None) diff --git a/tests/fast-gpu/test_fsdp_stream_load_equivalence.py b/tests/fast-gpu/test_fsdp_stream_load_equivalence.py deleted file mode 100644 index 080d753e..00000000 --- a/tests/fast-gpu/test_fsdp_stream_load_equivalence.py +++ /dev/null @@ -1,160 +0,0 @@ -from tests.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=180, - suite="stage-b-3-gpu-h200", - labels=[], -) - -import json -import os -from argparse import Namespace - -import pytest -import torch -import torch.distributed as dist -import torch.multiprocessing as mp - -WORLD_SIZE = 2 - -_TINY_WAN_CONFIG = { - "_class_name": "WanTransformer3DModel", - "attention_head_dim": 8, - "num_attention_heads": 2, - "num_layers": 2, - "in_channels": 4, - "out_channels": 4, - "text_dim": 16, - "freq_dim": 16, - "ffn_dim": 32, - "patch_size": [1, 2, 2], - "cross_attn_norm": True, - "qk_norm": "rms_norm_across_heads", - "eps": 1e-6, -} - - -def _make_checkpoint(root: str) -> None: - """Tiny diffusers-layout Wan checkpoint with a SHARDED transformer, so the - stream loader's multi-file path is exercised.""" - from diffusers import WanTransformer3DModel - - torch.manual_seed(7) - model = WanTransformer3DModel.from_config(_TINY_WAN_CONFIG) - model.save_pretrained(os.path.join(root, "transformer"), max_shard_size="20KB") - index = os.path.join(root, "transformer", "diffusion_pytorch_model.safetensors.index.json") - assert os.path.exists(index), "tiny checkpoint must be sharded for this test" - with open(os.path.join(root, "model_index.json"), "w") as f: - json.dump( - { - "_class_name": "WanPipeline", - "scheduler": ["diffusers", "UniPCMultistepScheduler"], - "transformer": ["diffusers", "WanTransformer3DModel"], - }, - f, - ) - os.makedirs(os.path.join(root, "scheduler"), exist_ok=True) - with open(os.path.join(root, "scheduler", "scheduler_config.json"), "w") as f: - json.dump({"_class_name": "UniPCMultistepScheduler", "num_train_timesteps": 1000}, f) - - -def _fsdp_args(**overrides) -> Namespace: - base = dict( - hf_checkpoint=None, - update_weight_target_modules=["transformer"], - diffusion_forward_dtype="bf16", - fsdp_reduce_dtype="fp32", - gradient_checkpointing=False, - lora_target_modules=["to_q", "to_k", "to_v"], - diffusion_init_lora_weight="gaussian", - lora_rank=4, - lora_alpha=4, - ) - base.update(overrides) - return Namespace(**base) - - -def _stream_load(ckpt: str, use_lora: bool): - from miles.backends.fsdp_utils.actor import ( - apply_fsdp2, - apply_lora, - materialize_sharded_model, - peft_checkpoint_key_map, - reset_lora_adapters, - ) - from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend - - args = _fsdp_args(hf_checkpoint=ckpt) - backend = DiffusersModelBackend(None) - raw_models, _ = backend.build_models_and_scheduler(args, master_dtype=torch.float32) - model = raw_models["transformer"] - key_map = None - if use_lora: - model = apply_lora(model, args, None, on_meta=True) - key_map = peft_checkpoint_key_map(model) - model = apply_fsdp2(model, args=args) - materialize_sharded_model(model, torch.cuda.current_device()) - backend.stream_load_weights(model, "transformer", args, master_dtype=torch.float32, key_map=key_map) - if use_lora: - reset_lora_adapters(model, args.diffusion_init_lora_weight) - return model - - -def _worker(rank: int, ckpt: str, use_lora: bool): - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = "29531" - torch.cuda.set_device(rank) - dist.init_process_group("nccl", rank=rank, world_size=WORLD_SIZE) - try: - from diffusers import WanTransformer3DModel - - model = _stream_load(ckpt, use_lora) - - # reference: plain full from_pretrained — what the legacy path - # materializes per component before sharding - ref = WanTransformer3DModel.from_pretrained(ckpt, subfolder="transformer", torch_dtype=torch.float32) - ref_sd = ref.state_dict() - - checked = 0 - for name, param in model.state_dict().items(): - if "lora_" in name: - continue - ref_key = name.replace(".base_layer", "").removeprefix("base_model.model.") - full = param.full_tensor() if hasattr(param, "full_tensor") else param - assert torch.equal(full.cpu(), ref_sd[ref_key]), f"mismatch: {name}" - checked += 1 - assert checked == len(ref_sd), f"covered {checked}/{len(ref_sd)} reference tensors" - - # non-persistent rope tables are not in any state_dict: they must - # survive to_empty via the snapshot in materialize_sharded_model - buffers = dict(model.named_buffers()) - ref_buffers = dict(ref.named_buffers()) - for name, ref_buffer in ref_buffers.items(): - live = next(b for n, b in buffers.items() if n.endswith(name)) - assert torch.equal(live.cpu(), ref_buffer), f"buffer mismatch: {name}" - - if use_lora: - from peft.tuners.lora import LoraLayer - - lora_layers = [m for m in model.modules() if isinstance(m, LoraLayer)] - assert lora_layers, "no LoRA layers injected" - for layer in lora_layers: - a = layer.lora_A["default"].weight - b = layer.lora_B["default"].weight - a = a.full_tensor() if hasattr(a, "full_tensor") else a - b = b.full_tensor() if hasattr(b, "full_tensor") else b - assert torch.isfinite(a).all() and a.abs().sum() > 0, "lora_A not initialized" - assert (b == 0).all(), "lora_B must start at zero" - finally: - dist.barrier() - dist.destroy_process_group() - - -@pytest.mark.parametrize("use_lora", [False, True], ids=["base", "lora"]) -def test_stream_load_matches_from_pretrained(tmp_path, use_lora): - if torch.cuda.device_count() < WORLD_SIZE: - pytest.skip(f"needs {WORLD_SIZE} GPUs") - ckpt = str(tmp_path / "ckpt") - os.makedirs(ckpt) - _make_checkpoint(ckpt) - mp.spawn(_worker, args=(ckpt, use_lora), nprocs=WORLD_SIZE, join=True) diff --git a/tests/fast/backends/fsdp_utils/test_stream_load.py b/tests/fast/backends/fsdp_utils/test_stream_load.py deleted file mode 100644 index 49063b04..00000000 --- a/tests/fast/backends/fsdp_utils/test_stream_load.py +++ /dev/null @@ -1,164 +0,0 @@ -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) - -import json -from argparse import Namespace - -import pytest -import torch - -from miles.backends.fsdp_utils.actor import ( - apply_lora, - materialize_sharded_model, - peft_checkpoint_key_map, - reset_lora_adapters, -) -from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend, ModelBackend - - -def _tiny_wan_checkpoint(tmp_path): - """A minimal diffusers-layout checkpoint dir: model_index.json + a 1-block - WanTransformer3DModel config + scheduler config. No weight files — enough - for build_models_and_scheduler, which must not read any.""" - config = { - "_class_name": "WanTransformer3DModel", - "attention_head_dim": 8, - "num_attention_heads": 2, - "num_layers": 1, - "in_channels": 4, - "out_channels": 4, - "text_dim": 16, - "freq_dim": 16, - "ffn_dim": 32, - "patch_size": [1, 2, 2], - "cross_attn_norm": True, - "qk_norm": "rms_norm_across_heads", - "eps": 1e-6, - } - model_index = { - "_class_name": "WanPipeline", - "scheduler": ["diffusers", "UniPCMultistepScheduler"], - "transformer": ["diffusers", "WanTransformer3DModel"], - "text_encoder": ["transformers", "UMT5EncoderModel"], - } - scheduler_config = {"_class_name": "UniPCMultistepScheduler", "num_train_timesteps": 1000} - (tmp_path / "transformer").mkdir() - (tmp_path / "scheduler").mkdir() - (tmp_path / "model_index.json").write_text(json.dumps(model_index)) - (tmp_path / "transformer" / "config.json").write_text(json.dumps(config)) - (tmp_path / "scheduler" / "scheduler_config.json").write_text(json.dumps(scheduler_config)) - return tmp_path - - -class TestBuildModelsAndScheduler: - # The core contract: params on meta (no weight I/O), buffers real (computed - # by __init__) — Wan's non-persistent rope tables are the regression case. - def test_params_meta_buffers_real(self, tmp_path): - ckpt = _tiny_wan_checkpoint(tmp_path) - args = Namespace(hf_checkpoint=str(ckpt), update_weight_target_modules=["transformer"]) - raw_models, scheduler = DiffusersModelBackend(None).build_models_and_scheduler( - args, master_dtype=torch.float32 - ) - model = raw_models["transformer"] - assert all(p.is_meta for p in model.parameters()) - buffers = dict(model.named_buffers()) - assert "rope.freqs_cos" in buffers and not buffers["rope.freqs_cos"].is_meta - assert type(scheduler).__name__ == "UniPCMultistepScheduler" - - def test_missing_component_raises(self, tmp_path): - ckpt = _tiny_wan_checkpoint(tmp_path) - args = Namespace(hf_checkpoint=str(ckpt), update_weight_target_modules=["vae"]) - with pytest.raises(ValueError, match="has no component"): - DiffusersModelBackend(None).build_models_and_scheduler(args, master_dtype=torch.float32) - - # transformers-library components (text_encoder) can't meta-init through - # diffusers class resolution — must point at the legacy escape hatch. - def test_non_diffusers_component_raises(self, tmp_path): - ckpt = _tiny_wan_checkpoint(tmp_path) - args = Namespace(hf_checkpoint=str(ckpt), update_weight_target_modules=["text_encoder"]) - with pytest.raises(ValueError, match="legacy"): - DiffusersModelBackend(None).build_models_and_scheduler(args, master_dtype=torch.float32) - - def test_custom_backend_without_meta_support_points_to_legacy(self): - class _LegacyOnly(ModelBackend): - def load_models_and_scheduler(self, args, *, master_dtype): - raise NotImplementedError - - with pytest.raises(NotImplementedError, match="fsdp-load-mode legacy"): - _LegacyOnly(None).build_models_and_scheduler(None, master_dtype=torch.float32) - - -class TestMaterializeShardedModel: - # to_empty wipes buffers; non-persistent ones aren't in any checkpoint, so - # materialize must carry them over (Wan rope). Persistent ones may be - # garbage afterwards — the stream load refills them. - def test_non_persistent_buffers_survive(self): - class _M(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(4, 4) - self.register_buffer("freqs", torch.arange(4, dtype=torch.float32), persistent=False) - - from accelerate import init_empty_weights - - with init_empty_weights(include_buffers=False): - m = _M() - assert m.linear.weight.is_meta and not m.freqs.is_meta - materialize_sharded_model(m, torch.device("cpu")) - assert not m.linear.weight.is_meta - torch.testing.assert_close(m.freqs, torch.arange(4, dtype=torch.float32)) - - -class _TinyDit(torch.nn.Module): - def __init__(self): - super().__init__() - self.to_q = torch.nn.Linear(8, 8, bias=False) - self.norm = torch.nn.LayerNorm(8) - - -def _lora_args(init_weight="gaussian"): - return Namespace( - lora_target_modules=["to_q"], - diffusion_init_lora_weight=init_weight, - lora_rank=4, - lora_alpha=4, - ) - - -class TestLoraStreamPath: - # Key map must invert the FQN wrapping so raw checkpoint keys land on the - # peft-wrapped names — same string contract as diffusion_update_weight_utils. - def test_key_map_inverts_peft_wrapping(self, monkeypatch): - import torch.distributed as dist - - monkeypatch.setattr(dist, "get_rank", lambda: 1) - with torch.device("meta"): - model = _TinyDit() - model = apply_lora(model, _lora_args(), None, on_meta=True) - key_map = peft_checkpoint_key_map(model) - assert key_map("to_q.weight") == "base_model.model.to_q.base_layer.weight" - assert key_map("norm.weight") == "base_model.model.norm.weight" - assert key_map("unknown.weight") == "unknown.weight" # passthrough - - def test_adapters_created_on_meta_then_reset(self, monkeypatch): - import torch.distributed as dist - - monkeypatch.setattr(dist, "get_rank", lambda: 1) - with torch.device("meta"): - model = _TinyDit() - model = apply_lora(model, _lora_args(), None, on_meta=True) - lora_layer = model.base_model.model.to_q - assert lora_layer.lora_A["default"].weight.is_meta - model.to_empty(device="cpu") - reset_lora_adapters(model, "gaussian") - assert torch.isfinite(lora_layer.lora_A["default"].weight).all() - assert (lora_layer.lora_B["default"].weight == 0).all() - - # pissa/olora-style inits read the base weights at wrap time; on meta they - # would silently produce garbage — must be an explicit error. - def test_data_aware_init_rejected_on_meta(self): - with torch.device("meta"): - model = _TinyDit() - with pytest.raises(ValueError, match="legacy"): - apply_lora(model, _lora_args(init_weight="pissa"), None, on_meta=True) From aef68a2a5335ec9d18803efd379f444e61f2906a Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 17 Jul 2026 01:12:18 +0000 Subject: [PATCH 03/10] clarify LTX weight materialization --- miles/backends/fsdp_utils/model_backend.py | 16 +++++++++++++++- miles/backends/fsdp_utils/models/ltx2.py | 6 +++--- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 3ef8f637..f6f82977 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -1,3 +1,17 @@ +"""Model backend: owns model-side behavior for the FSDP trainer. + +Selected via ``--model-backend-path`` (miles custom-function style); the +family config declares the default. Three concerns, all properties of the +concrete modeling rather than of the training loop: + + - ``load_component`` / ``load_scheduler``: checkpoint -> model components and scheduler + - ``enable_gradient_checkpointing``: how this model turns on grad ckpt + - ``fsdp_no_split_modules``: which block classes FSDP wraps + +Defaults implement the diffusers protocol (see ``models/__init__.py``); a +native model overrides methods here instead of retrofitting its instances. +""" + from __future__ import annotations import functools @@ -171,7 +185,7 @@ def load_component( checkpoint, device="cpu", dtype=master_dtype, - load_weights=dist.get_rank() == 0, + materialize_weights=dist.get_rank() == 0, ) def load_scheduler(self, args) -> Any: diff --git a/miles/backends/fsdp_utils/models/ltx2.py b/miles/backends/fsdp_utils/models/ltx2.py index 43c58cc6..b535d63a 100644 --- a/miles/backends/fsdp_utils/models/ltx2.py +++ b/miles/backends/fsdp_utils/models/ltx2.py @@ -85,7 +85,7 @@ def load_ltx_transformer_for_train( *, device: str = "cpu", dtype: Any = None, - load_weights: bool = True, + materialize_weights: bool = True, ): """Load LTX DiT for FSDP train from materialized diffusers or comfy safetensors. @@ -111,7 +111,7 @@ def load_ltx_transformer_for_train( if _is_materialized_diffusers_checkpoint(checkpoint): config = _read_materialized_transformer_config(checkpoint) meta_model = create_meta_model(LTXModelConfigurator, config, ()) - if not load_weights: + if not materialize_weights: return meta_model.to(dtype=dtype) loader = SafetensorsModelStateDictLoader() sd = load_state_dict( @@ -136,7 +136,7 @@ def load_ltx_transformer_for_train( model_class_configurator=LTXModelConfigurator, model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, ) - if not load_weights: + if not materialize_weights: return builder.meta_model(builder.model_config(), builder.module_ops).to(dtype=dtype) return builder.build(device=torch_device, dtype=dtype) From 5f4f8aad0abcb2881d5785888864b2275b53a697 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 17 Jul 2026 01:24:26 +0000 Subject: [PATCH 04/10] restore deterministic attention note --- miles/backends/fsdp_utils/actor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f2b3b876..10491ae3 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -93,6 +93,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.train_pipeline_config.configure(args) self.model_backend = load_function(args.model_backend_path)(self.train_pipeline_config) if args.deterministic_mode: + # flash-attn is opaque to torch's determinism flag; backends patch their own dispatch. self.model_backend.enable_deterministic_attention(args.fsdp_attention_backend) self.scheduler = self.model_backend.load_scheduler(args) init_context = self._get_init_weight_context_manager() From 39ff3cf9d6b0d9df24e258faa09cbf6cd4bdc3ee Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 17 Jul 2026 02:40:02 +0000 Subject: [PATCH 05/10] restore comments dropped in the loading refactor Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 13 ++++++++++++- miles/backends/fsdp_utils/configs/qwen_image.py | 1 + .../fsdp_utils/configs/train_pipeline_config.py | 1 + miles/backends/fsdp_utils/model_backend.py | 6 ++++++ 4 files changed, 20 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 10491ae3..b8fe0b4e 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -100,6 +100,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.models: dict[str, torch.nn.Module] = {} for component in args.update_weight_target_modules: + # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap with init_context(): model = self.model_backend.load_component(component, args, master_dtype=self._master_dtype) if args.fsdp_attention_backend is not None: @@ -706,12 +707,22 @@ def _resolve_dtype(name: str) -> torch.dtype: def apply_lora( model: torch.nn.Module, args: Namespace, train_pipeline_config, on_meta: bool = False ) -> torch.nn.Module: + """Apply PEFT LoRA to the model. + + Args: + model: The model to apply LoRA to. + args: Arguments containing LoRA settings. + train_pipeline_config: The train pipeline config. + on_meta: Build the adapter on the meta device without initializing + weights (they arrive later via broadcast from rank 0). + """ from peft import LoraConfig, get_peft_model + # Per-model fallback when --lora-target-modules is unset (runtime inference: depends on loaded pipeline). targets = args.lora_target_modules or train_pipeline_config.lora_target_modules init_lora_weight = args.diffusion_init_lora_weight if init_lora_weight == "kaiming-uniform": - init_lora_weight = True + init_lora_weight = True # namely kaiming-uniform model = get_peft_model( model, LoraConfig( diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 8e69855d..1d563cde 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -186,4 +186,5 @@ def cfg_combine( return combined def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: + """Postprocess the model after FSDP wrap + weight materialization.""" _rebuild_pos_embed_freqs_on_cuda(model) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index d29d0bad..a553735c 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -185,4 +185,5 @@ def cfg_combine( """Apply classifier-free guidance. Model-specific (e.g. rescale or not).""" def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: + """Postprocess the model after FSDP wrap + weight materialization (default: no-op).""" return None diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index f6f82977..032a9a96 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -50,15 +50,19 @@ def load_component( *, master_dtype: torch.dtype, ) -> torch.nn.Module: + """Return the ``component`` model on CPU; must honor an ambient meta-device init context.""" raise NotImplementedError def load_scheduler(self, args) -> Any: + """Return the pipeline's training scheduler.""" raise NotImplementedError def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None: + """Turn on grad checkpointing; default = the diffusers protocol method.""" model.enable_gradient_checkpointing() def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]: + """Block class names FSDP wraps; default = the model's own declaration.""" return model._no_split_modules def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None: @@ -200,6 +204,8 @@ def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]: return ["BasicAVTransformerBlock"] def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None: + # ltx_core selects attention via AttentionFunction (not diffusers' set_attention_backend(str)); + # map the flag and reuse ltx_core's own module op to swap it on every Attention submodule. from ltx_core.loader.attention_ops import set_attention_module_op from ltx_core.model.transformer.attention import AttentionFunction, MaskedAttentionFunction From a0d932f187ee264eaea42596bcb0b41b4b1fd9f4 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 17 Jul 2026 02:50:33 +0000 Subject: [PATCH 06/10] align FSDP materialization with miles' verl-derived flow rank0 keeps its real shards through .to(cuda) instead of to_empty + re-fill; every buffer is broadcast from rank0 after set_model_state_dict (covers non-persistent buffers absent from any state_dict, instead of trusting recomputed-from-config values); under fsdp_cpu_offload the load happens on GPU and buffers stay there, since CPUOffloadPolicy manages params only. Add short specs to the three loading helpers. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 61 +++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 14 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index b8fe0b4e..ac4cbd81 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -125,12 +125,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), ) - device = ( - torch.device("cpu") - if self.args.fsdp_cpu_offload - else torch.device("cuda", torch.cuda.current_device()) - ) - load_sharded_model(model, full_state, device) + load_sharded_model(model, full_state, cpu_offload=self.args.fsdp_cpu_offload) del full_state self.train_pipeline_config.postprocess_model_after_materialize(model) self.models[component] = model @@ -196,10 +191,22 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty return self.args.start_rollout_id def _get_init_weight_context_manager(self): + """Init context: rank0 builds on CPU with real weights; other ranks build on + the meta device (zero allocation) and are materialized by ``load_sharded_model``. + + Mirrors miles' ``FSDPTrainRayActor._get_init_weight_context_manager`` (itself + from verl), except: ``include_buffers=False`` is explicit so buffers are + computed for real in ``__init__`` (non-persistent ones exist in no checkpoint), + and diffusers' per-tensor meta-copy warnings are suppressed. No + ``tie_word_embeddings`` escape hatch — DiT components don't tie embeddings. + """ from accelerate import init_empty_weights + def cpu_init_weights(): + return torch.device("cpu") + if dist.get_rank() == 0: - return lambda: torch.device("cpu") + return cpu_init_weights @contextmanager def meta_init_weights(): @@ -738,27 +745,53 @@ def apply_lora( return model -def load_sharded_model(model: torch.nn.Module, full_state: dict, device) -> None: +def load_sharded_model(model: torch.nn.Module, full_state: dict, cpu_offload: bool) -> None: + """Materialize an FSDP2-wrapped model by broadcasting rank0's full state dict. + + Rank0 moves its real shards to GPU; other ranks allocate empty shards and are + filled via ``set_model_state_dict(broadcast_from_rank0=True)``; buffers are then + broadcast from rank0. Mirrors miles' ``FSDPTrainRayActor._fsdp2_load_full_state_dict`` + (itself from verl/utils/fsdp_utils.py::fsdp2_load_full_state_dict). + """ from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict - real_buffers = {name: buffer.detach().clone() for name, buffer in model.named_buffers() if not buffer.is_meta} - model.to_empty(device=device) - for name, buffer in model.named_buffers(): - if name in real_buffers: - buffer.copy_(real_buffers[name]) + if dist.get_rank() == 0: + # Rank 0 was sharded on real CPU weights; move them (and real buffers) along. + model.to(device=torch.cuda.current_device(), non_blocking=True) + else: + # to_empty creates tensors on device without initializing memory. + model.to_empty(device=torch.cuda.current_device()) + set_model_state_dict( model, full_state, options=StateDictOptions( full_state_dict=True, - cpu_offload=device.type == "cpu", + cpu_offload=cpu_offload, broadcast_from_rank0=True, strict=True, ), ) + # set_model_state_dict only covers state_dict entries; non-persistent buffers + # (e.g. Wan's rope tables) exist in no state_dict and were wiped by to_empty + # on non-rank0 ranks — take rank0's real values for every buffer. + for _name, buffer in model.named_buffers(): + dist.broadcast(buffer, src=0) + + if cpu_offload: + model.to("cpu", non_blocking=True) + # CPUOffloadPolicy manages params only; buffers must live on GPU for forward. + for buffer in model.buffers(): + buffer.data = buffer.data.to(torch.cuda.current_device()) def sync_model_dtypes(model: torch.nn.Module) -> None: + """Cast non-rank0 (meta) params/buffers to rank0's per-tensor dtypes. + + Meta ranks load with ``_keep_in_fp32_modules`` disabled (see + ``DiffusersModelBackend._load_pipeline``), so their dtypes can drift from rank0's + real load; sharding and rank0-broadcast require identical dtypes on every rank. + """ dtypes = None if dist.get_rank() == 0: dtypes = { From 5b77215ed3e5781084c6b405349e65559b8b1f31 Mon Sep 17 00:00:00 2001 From: zhihengy Date: Fri, 17 Jul 2026 21:31:24 +0000 Subject: [PATCH 07/10] restore abc.ABC on ModelBackend Keep the loading interface (load_component/load_scheduler) enforced at class-definition time rather than first call; the conditionally-required attention hooks stay NotImplementedError as before. Test-local subclasses stub the abstract loaders themselves (also drops the stale load_models_and_scheduler stubs left from the old interface). Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/model_backend.py | 7 ++++--- tests/fast/backends/fsdp_utils/test_model_backend.py | 10 ++++++++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 032a9a96..614c5aeb 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -14,6 +14,7 @@ from __future__ import annotations +import abc import functools import importlib import inspect @@ -27,7 +28,7 @@ logger = logging.getLogger(__name__) -class ModelBackend: +class ModelBackend(abc.ABC): def __init__(self, train_pipeline_config): self.config = train_pipeline_config @@ -43,6 +44,7 @@ def _enable_deterministic_flash_attention(self, name: str) -> None: f"{name!r}; use a native/math attention backend under deterministic mode." ) + @abc.abstractmethod def load_component( self, component: str, @@ -51,11 +53,10 @@ def load_component( master_dtype: torch.dtype, ) -> torch.nn.Module: """Return the ``component`` model on CPU; must honor an ambient meta-device init context.""" - raise NotImplementedError + @abc.abstractmethod def load_scheduler(self, args) -> Any: """Return the pipeline's training scheduler.""" - raise NotImplementedError def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None: """Turn on grad checkpointing; default = the diffusers protocol method.""" diff --git a/tests/fast/backends/fsdp_utils/test_model_backend.py b/tests/fast/backends/fsdp_utils/test_model_backend.py index bb10f074..1fbb6e95 100644 --- a/tests/fast/backends/fsdp_utils/test_model_backend.py +++ b/tests/fast/backends/fsdp_utils/test_model_backend.py @@ -30,7 +30,10 @@ def test_subclass_can_override(self): seen = {} class _CustomBackend(ModelBackend): - def load_models_and_scheduler(self, args, *, master_dtype): + def load_component(self, component, args, *, master_dtype): + raise NotImplementedError + + def load_scheduler(self, args): raise NotImplementedError def set_attention_backend(self, model, backend): @@ -49,7 +52,10 @@ class _NoAttnModel(torch.nn.Module): pass # no set_attention_backend, like a native ltx_core transformer class _OptOutBackend(ModelBackend): - def load_models_and_scheduler(self, args, *, master_dtype): + def load_component(self, component, args, *, master_dtype): + raise NotImplementedError + + def load_scheduler(self, args): raise NotImplementedError def set_attention_backend(self, model, backend): From 671b9b438f8982bae587b852fb09d4ba51fb1517 Mon Sep 17 00:00:00 2001 From: zhihengy Date: Fri, 17 Jul 2026 21:33:38 +0000 Subject: [PATCH 08/10] untrack test_model_backend.py Co-Authored-By: Claude Fable 5 --- .../backends/fsdp_utils/test_model_backend.py | 64 ------------------- 1 file changed, 64 deletions(-) delete mode 100644 tests/fast/backends/fsdp_utils/test_model_backend.py diff --git a/tests/fast/backends/fsdp_utils/test_model_backend.py b/tests/fast/backends/fsdp_utils/test_model_backend.py deleted file mode 100644 index 1fbb6e95..00000000 --- a/tests/fast/backends/fsdp_utils/test_model_backend.py +++ /dev/null @@ -1,64 +0,0 @@ -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) - -import torch - -from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend, ModelBackend - - -class _RecordingModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.selected = None - - def set_attention_backend(self, backend): # diffusers protocol method - self.selected = backend - - -class TestSetAttentionBackend: - # Default hook delegates to the model's own diffusers protocol method. - def test_diffusers_default_delegates_to_model(self): - model = _RecordingModel() - DiffusersModelBackend(None).set_attention_backend(model, "flash") - assert model.selected == "flash" - - # A custom backend overrides the hook (non-diffusers models select attention a - # different way, or opt out) — the actor no longer calls model.set_attention_backend - # directly, so a model without that method never crashes. - def test_subclass_can_override(self): - seen = {} - - class _CustomBackend(ModelBackend): - def load_component(self, component, args, *, master_dtype): - raise NotImplementedError - - def load_scheduler(self, args): - raise NotImplementedError - - def set_attention_backend(self, model, backend): - seen["backend"] = backend - - model = _RecordingModel() - _CustomBackend(None).set_attention_backend(model, "fa3") - assert seen["backend"] == "fa3" - assert model.selected is None # diffusers default path not taken - - # The concrete crash the refactor removes: a model without set_attention_backend - # (e.g. ltx_core transformers) must not raise when a backend that opts out of the - # diffusers protocol handles it — the actor never calls the model method directly. - def test_model_without_method_does_not_crash(self): - class _NoAttnModel(torch.nn.Module): - pass # no set_attention_backend, like a native ltx_core transformer - - class _OptOutBackend(ModelBackend): - def load_component(self, component, args, *, master_dtype): - raise NotImplementedError - - def load_scheduler(self, args): - raise NotImplementedError - - def set_attention_backend(self, model, backend): - pass # this family selects attention its own way - - _OptOutBackend(None).set_attention_backend(_NoAttnModel(), "fa3") # no AttributeError From ce85e1c4df9a4bce71866b9976829abb2bd85c3a Mon Sep 17 00:00:00 2001 From: zhihengy Date: Fri, 17 Jul 2026 21:36:45 +0000 Subject: [PATCH 09/10] Revert "untrack test_model_backend.py" This reverts commit 671b9b41; the CI test is still wanted. Co-Authored-By: Claude Fable 5 --- .../backends/fsdp_utils/test_model_backend.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 tests/fast/backends/fsdp_utils/test_model_backend.py diff --git a/tests/fast/backends/fsdp_utils/test_model_backend.py b/tests/fast/backends/fsdp_utils/test_model_backend.py new file mode 100644 index 00000000..1fbb6e95 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_model_backend.py @@ -0,0 +1,64 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) + +import torch + +from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend, ModelBackend + + +class _RecordingModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.selected = None + + def set_attention_backend(self, backend): # diffusers protocol method + self.selected = backend + + +class TestSetAttentionBackend: + # Default hook delegates to the model's own diffusers protocol method. + def test_diffusers_default_delegates_to_model(self): + model = _RecordingModel() + DiffusersModelBackend(None).set_attention_backend(model, "flash") + assert model.selected == "flash" + + # A custom backend overrides the hook (non-diffusers models select attention a + # different way, or opt out) — the actor no longer calls model.set_attention_backend + # directly, so a model without that method never crashes. + def test_subclass_can_override(self): + seen = {} + + class _CustomBackend(ModelBackend): + def load_component(self, component, args, *, master_dtype): + raise NotImplementedError + + def load_scheduler(self, args): + raise NotImplementedError + + def set_attention_backend(self, model, backend): + seen["backend"] = backend + + model = _RecordingModel() + _CustomBackend(None).set_attention_backend(model, "fa3") + assert seen["backend"] == "fa3" + assert model.selected is None # diffusers default path not taken + + # The concrete crash the refactor removes: a model without set_attention_backend + # (e.g. ltx_core transformers) must not raise when a backend that opts out of the + # diffusers protocol handles it — the actor never calls the model method directly. + def test_model_without_method_does_not_crash(self): + class _NoAttnModel(torch.nn.Module): + pass # no set_attention_backend, like a native ltx_core transformer + + class _OptOutBackend(ModelBackend): + def load_component(self, component, args, *, master_dtype): + raise NotImplementedError + + def load_scheduler(self, args): + raise NotImplementedError + + def set_attention_backend(self, model, backend): + pass # this family selects attention its own way + + _OptOutBackend(None).set_attention_backend(_NoAttnModel(), "fa3") # no AttributeError From e4dc7d95cd621769236ec3173748134af3a58d4a Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 17 Jul 2026 21:45:48 +0000 Subject: [PATCH 10/10] load diffusers components individually, not via pipeline with None siblings WanPipeline declares transformer/transformer_2 optional with a None default, which drops them from expected_modules: the None passed to DiffusionPipeline.from_pretrained is silently ignored and both 14B experts load from disk on every rank -- during load_scheduler too -- and on non-rank0 (low_cpu_mem_usage=False) trip diffusers' _keep_in_fp32_modules guard, crashing any multi-rank init. Resolve the component class from model_index.json and from_pretrained(subfolder=...) it directly instead. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 2 +- miles/backends/fsdp_utils/model_backend.py | 62 +++++++++++++--------- 2 files changed, 38 insertions(+), 26 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index ac4cbd81..545dbe3e 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -789,7 +789,7 @@ def sync_model_dtypes(model: torch.nn.Module) -> None: """Cast non-rank0 (meta) params/buffers to rank0's per-tensor dtypes. Meta ranks load with ``_keep_in_fp32_modules`` disabled (see - ``DiffusersModelBackend._load_pipeline``), so their dtypes can drift from rank0's + ``DiffusersModelBackend.load_component``), so their dtypes can drift from rank0's real load; sharding and rank0-broadcast require identical dtypes on every rank. """ dtypes = None diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 614c5aeb..c5ed662c 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -94,41 +94,53 @@ def load_component( *, master_dtype: torch.dtype, ) -> torch.nn.Module: - config = DiffusionPipeline.load_config(args.hf_checkpoint) - if component not in config: - raise ValueError(f"pipeline {args.hf_checkpoint} has no component {component!r}") - - pipeline = self._load_pipeline(args, config, component, master_dtype) - model = getattr(pipeline, component) - del pipeline - return model - - def load_scheduler(self, args) -> Any: - config = DiffusionPipeline.load_config(args.hf_checkpoint) - pipeline = self._load_pipeline(args, config, "scheduler") - scheduler = pipeline.scheduler - del pipeline - return scheduler - - @classmethod - def _load_pipeline(cls, args, config: dict, component: str, master_dtype=None): - skipped = { - name: None for name, value in config.items() if name != component and isinstance(value, (list, tuple)) + model_cls = self._resolve_component_class(args, component) + kwargs = { + "subfolder": component, + "torch_dtype": master_dtype, + "low_cpu_mem_usage": dist.get_rank() == 0, } - kwargs = {"trust_remote_code": True, "low_cpu_mem_usage": dist.get_rank() == 0, **skipped} - if master_dtype is not None: - kwargs["torch_dtype"] = master_dtype - model_cls = cls._component_class(config.get(component)) + # Non-rank0 loads with low_cpu_mem_usage=False so the ambient meta-device + # context keeps params on meta; diffusers forbids that combination when the + # class pins modules to fp32, so disable the pin for the duration (dtypes are + # re-synced from rank0 afterwards, see ``sync_model_dtypes``). keep_in_fp32 = getattr(model_cls, "_keep_in_fp32_modules", None) if dist.get_rank() != 0 and keep_in_fp32 is not None: model_cls._keep_in_fp32_modules = None try: - return DiffusionPipeline.from_pretrained(args.hf_checkpoint, **kwargs) + return model_cls.from_pretrained(args.hf_checkpoint, **kwargs) finally: if dist.get_rank() != 0 and keep_in_fp32 is not None: model_cls._keep_in_fp32_modules = keep_in_fp32 + def load_scheduler(self, args) -> Any: + scheduler_cls = self._resolve_component_class(args, "scheduler") + return scheduler_cls.from_pretrained(args.hf_checkpoint, subfolder="scheduler") + + @classmethod + def _resolve_component_class(cls, args, component: str): + """Resolve ``component``'s class from ``model_index.json``. + + Components load individually via ``cls.from_pretrained(subfolder=...)`` rather + than through ``DiffusionPipeline.from_pretrained`` with the siblings passed as + ``None``: pipelines that declare a component optional with a ``None`` default + (e.g. ``WanPipeline.transformer``/``transformer_2``) drop it from + ``expected_modules``, so the ``None`` is silently ignored and the sibling is + loaded from disk anyway — on every rank, and with ``low_cpu_mem_usage=False`` + it also trips diffusers' ``_keep_in_fp32_modules`` guard. + """ + config = DiffusionPipeline.load_config(args.hf_checkpoint) + if component not in config: + raise ValueError(f"pipeline {args.hf_checkpoint} has no component {component!r}") + component_cls = cls._component_class(config[component]) + if component_cls is None: + raise ValueError( + f"cannot resolve the class for component {component!r} of {args.hf_checkpoint} " + f"from spec {config[component]!r}; remote-code components are not supported" + ) + return component_cls + @staticmethod def _component_class(spec): if not isinstance(spec, (list, tuple)) or len(spec) != 2: