feat(fsdp): meta-init + rank0-streamed weight loading (--fsdp-load-mode stream)

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,35 +95,39 @@ 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,
                 cpu_offload=self.args.fsdp_cpu_offload,
                 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
 

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

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)

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

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

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