diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d0af5e24..e3c3ba46 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,7 +1,9 @@ import logging +import warnings from argparse import Namespace from collections import defaultdict -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext +from itertools import chain import ray import torch @@ -94,13 +96,14 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty 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) - 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) + rank = dist.get_rank() self.models: dict[str, torch.nn.Module] = {} - for component, model in raw_models.items(): + for component in args.update_weight_target_modules: # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap + with self._init_weight_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) @@ -112,10 +115,10 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty 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) - + if 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() if rank == 0 else {} model = apply_fsdp2( model, mesh=self.parallel_state.dp_mesh, @@ -123,6 +126,9 @@ 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), ) + 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 # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -185,6 +191,24 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty return self.args.start_rollout_id + @contextmanager + def _init_weight_context(self): + """Build real weights on rank0 and allocation-free meta weights elsewhere.""" + if dist.get_rank() == 0: + with torch.device("cpu"): + yield + return + + from accelerate import init_empty_weights + + # Some models compute buffer values during __init__, which cannot run on meta. + 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 + def _get_parallel_config(self) -> dict: return {"dp_size": getattr(self.parallel_state, "dp_size", 1)} @@ -677,16 +701,11 @@ 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: - """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. - """ +def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> torch.nn.Module: + """Apply PEFT LoRA, leaving non-rank0 adapters uninitialized on meta.""" from peft import LoraConfig, get_peft_model + on_meta = dist.get_rank() != 0 # 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 @@ -698,14 +717,60 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - 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, ) if dist.get_rank() == 0: model.print_trainable_parameters() return model +def load_sharded_model(model: torch.nn.Module, full_state: dict, cpu_offload: bool) -> None: + """Materialize FSDP2 shards from rank0's full state dict.""" + from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict + + 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=cpu_offload, + broadcast_from_rank0=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 buffer in model.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: + """Match meta parameter and buffer dtypes to rank0 before sharding.""" + rank = dist.get_rank() + tensors = list(chain(model.parameters(), model.buffers())) + dtypes = [tensor.dtype for tensor in tensors] if rank == 0 else None + objects = [dtypes] + dist.broadcast_object_list(objects, src=0) + if rank != 0: + for tensor, dtype in zip(tensors, objects[0], strict=True): + tensor.data = tensor.data.to(dtype) + + 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/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index 97b23861..9aed83cf 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -198,6 +198,3 @@ def cfg_combine( if scale == 1.0: return noise_pred_pos 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/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 197aac1d..42910804 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 @@ -191,6 +186,6 @@ 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: + """Postprocess the model after FSDP wrap + weight materialization.""" _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 dc3e0c87..a553735c 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -184,6 +184,6 @@ 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: + """Postprocess the model after FSDP wrap + weight materialization (default: no-op).""" + 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 a9303d5c..9f7495d4 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -4,7 +4,7 @@ family config declares the default. Three concerns, all properties of the concrete modeling rather than of the training loop: - - ``load_models_and_scheduler``: checkpoint -> ``({component: model}, scheduler)`` + - ``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 @@ -14,19 +14,20 @@ from __future__ import annotations -import abc import functools +import importlib import inspect import logging from typing import Any import torch +import torch.distributed as dist from diffusers import DiffusionPipeline logger = logging.getLogger(__name__) -class ModelBackend(abc.ABC): +class ModelBackend: def __init__(self, train_pipeline_config): self.config = train_pipeline_config @@ -42,14 +43,19 @@ 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_models_and_scheduler( + def load_component( self, + component: str, args, *, master_dtype: torch.dtype, - ) -> tuple[dict[str, torch.nn.Module], Any]: - """Return ``({component: model}, scheduler)`` on CPU.""" + ) -> 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.""" @@ -80,31 +86,76 @@ def _enable_deterministic_flash_attention(self, name: str) -> None: setattr(ad, fn_name, functools.partial(getattr(ad, fn_name), deterministic=True)) logger.info("Enabled deterministic flash attention backward for: %s", ", ".join(names)) - 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 - del pipeline - return raw_models, scheduler + ) -> torch.nn.Module: + model_cls = self._resolve_component_class(args, component) + rank = dist.get_rank() + kwargs = { + "subfolder": component, + "torch_dtype": master_dtype, + "low_cpu_mem_usage": rank == 0, + } + + # 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 rank != 0 and keep_in_fp32 is not None: + model_cls._keep_in_fp32_modules = None + try: + return model_cls.from_pretrained(args.hf_checkpoint, **kwargs) + finally: + if 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: + 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) class LTXModelBackend(ModelBackend): @@ -132,25 +183,32 @@ def _enable_deterministic_flash_attention(self, name: str) -> None: ) logger.info("Enabled deterministic ltx_core flash attention backward for: %s", ", ".join(patched)) - def load_models_and_scheduler( + def load_component( self, + component: str, args, *, master_dtype: torch.dtype, - ) -> tuple[dict[str, torch.nn.Module], Any]: + ) -> torch.nn.Module: from miles.backends.fsdp_utils.models.ltx2 import ( - build_ltx_train_scheduler, load_ltx_transformer_for_train, resolve_transformer_checkpoint, ) - modules = list(args.update_weight_target_modules) - if modules != ["transformer"]: - raise ValueError(f"LTX trains the single DiT ('transformer'); got {modules}") - # TODO: meta-init on non-rank-0 before multi-node runs (every rank loads the full weights). + if component != "transformer": + raise ValueError(f"LTX trains the single DiT ('transformer'); got {component!r}") checkpoint = resolve_transformer_checkpoint(str(args.diffusion_model)) - model = load_ltx_transformer_for_train(checkpoint, device="cpu", dtype=master_dtype) - return {"transformer": model}, build_ltx_train_scheduler(args) + return load_ltx_transformer_for_train( + checkpoint, + device="cpu", + dtype=master_dtype, + materialize_weights=dist.get_rank() == 0, + ) + + def load_scheduler(self, args) -> Any: + from miles.backends.fsdp_utils.models.ltx2 import build_ltx_train_scheduler + + return build_ltx_train_scheduler(args) def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None: model.set_gradient_checkpointing(True) diff --git a/miles/backends/fsdp_utils/models/ltx2.py b/miles/backends/fsdp_utils/models/ltx2.py index 2b285283..b535d63a 100644 --- a/miles/backends/fsdp_utils/models/ltx2.py +++ b/miles/backends/fsdp_utils/models/ltx2.py @@ -85,6 +85,7 @@ def load_ltx_transformer_for_train( *, device: str = "cpu", dtype: Any = None, + materialize_weights: bool = True, ): """Load LTX DiT for FSDP train from materialized diffusers or comfy safetensors. @@ -110,6 +111,8 @@ 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 materialize_weights: + return meta_model.to(dtype=dtype) loader = SafetensorsModelStateDictLoader() sd = load_state_dict( str(checkpoint), @@ -128,11 +131,14 @@ def load_ltx_transformer_for_train( ) return meta_model.to(torch_device) - return SingleGPUModelBuilder( + builder = SingleGPUModelBuilder( model_path=str(checkpoint), model_class_configurator=LTXModelConfigurator, model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, - ).build(device=torch_device, dtype=dtype) + ) + 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) def ensure_materialized_model(hf_model_id: str) -> Path: