diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d0af5e24..545dbe3e 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,7 +1,8 @@ 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 @@ -94,28 +95,29 @@ 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) + init_context = self._get_init_weight_context_manager() 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 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) if args.use_lora: - model = apply_lora(model, args, self.train_pipeline_config) + 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) - model.to(torch.cuda.current_device()) - - self.train_pipeline_config.preprocess_model_before_fsdp(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, @@ -123,6 +125,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 +190,35 @@ 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 cpu_init_weights + + @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)} @@ -677,13 +711,17 @@ 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 +) -> 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 @@ -698,14 +736,79 @@ 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 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 + + 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, + 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_component``), 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 = { + "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): 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 38fb8e20..b8f8e47b 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -196,6 +196,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 a2946b16..1d563cde 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,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..c5ed662c 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 @@ -16,11 +16,13 @@ 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__) @@ -43,13 +45,18 @@ def _enable_deterministic_flash_attention(self, name: str) -> None: ) @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.""" + + @abc.abstractmethod + def load_scheduler(self, args) -> Any: + """Return the pipeline's training scheduler.""" def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None: """Turn on grad checkpointing; default = the diffusers protocol method.""" @@ -80,31 +87,75 @@ 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) + kwargs = { + "subfolder": component, + "torch_dtype": master_dtype, + "low_cpu_mem_usage": dist.get_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 dist.get_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 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: + 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: 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):