refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch

miles/backends/fsdp_utils/actor.py

@@ -39,6 +39,7 @@
 from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch
 from .lr_scheduler import get_lr_scheduler
 from .metrics import new_metric_buffer
+from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str
 from .parallel import create_fsdp_parallel_state
 from .sequence_parallel.plan import apply_sequence_parallel
 
@@ -95,8 +96,8 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
 
         self.prof = TrainProfiler(args)
 
-        self._master_dtype = _resolve_dtype(args.fsdp_master_dtype)
-        self._forward_dtype = _resolve_dtype(args.diffusion_forward_dtype)
+        self._master_dtype = parse_dtype_from_str(args.fsdp_master_dtype)
+        self._forward_dtype = parse_dtype_from_str(args.diffusion_forward_dtype)
 
         from miles.utils.misc import load_function
 
@@ -140,10 +141,10 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             full_state = model.state_dict() if rank == 0 else {}
             model = apply_fsdp2(
                 model,
+                self.model_backend.fsdp_parallel_plan(model),
                 mesh=self.parallel_state.get_mesh("fsdp"),
                 cpu_offload=self.args.fsdp_cpu_offload,
                 args=self.args,
-                no_split_modules=self.model_backend.fsdp_no_split_modules(model),
             )
             checkpoint.broadcast_full_state_to_fsdp(
                 model,
@@ -597,10 +598,6 @@ def move_torch_optimizer(optimizer, device):
     torch.cuda.synchronize()
 
 
-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) -> torch.nn.Module:
     """Apply PEFT LoRA, leaving non-rank0 adapters uninitialized on meta."""
     from peft import LoraConfig, get_peft_model
@@ -626,36 +623,84 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -
     return model
 
 
-def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None):
+def apply_fsdp2(
+    model,
+    parallel_plan,
+    mesh=None,
+    cpu_offload=False,
+    args=None,
+):
+    """Apply FSDP2 per the model's FSDPParallelPlan.
+
+    ``parallel_plan.param_dtype_patterns`` is matched against FQNs from ``model``. Each child
+    ``fully_shard`` call receives exact FQNs relative to that child module, while
+    parameters managed by the root call retain their root-relative FQNs.
+    """
     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
 
     offload_policy = CPUOffloadPolicy() if cpu_offload else None
 
-    layer_cls_to_wrap = no_split_modules if no_split_modules is not None else model._no_split_modules
-    assert len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None
+    layer_cls_to_wrap = parallel_plan.no_split_modules
+    assert layer_cls_to_wrap is not None and len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None
 
     modules = [module for name, module in model.named_modules() if module.__class__.__name__ in layer_cls_to_wrap]
 
-    param_dtype = _resolve_dtype(args.diffusion_forward_dtype)
-    reduce_dtype = _resolve_dtype(args.fsdp_reduce_dtype)
+    param_dtype = parse_dtype_from_str(args.diffusion_forward_dtype)
+    reduce_dtype = parse_dtype_from_str(args.fsdp_reduce_dtype)
+    # A wrap entry may also be a module LIST — fully_shard can group several modules into one wrap
+    # (one shared all-gather); today every wrap holds a single block.
+    param_dtype_maps = compile_param_dtype_maps(
+        model,
+        modules,
+        parallel_plan.param_dtype_patterns,
+        param_dtype,
+    )
+    has_param_dtype_overrides = bool(any(param_dtype_maps.wrap_maps) or param_dtype_maps.root_map)
+    param_dtype_policy_cls = None
+    if has_param_dtype_overrides:
+        from .fsdp_param_dtype_patch import ParamDtypeMixedPrecisionPolicy, apply_param_dtype_map_patch
+
+        apply_param_dtype_map_patch()
+        param_dtype_policy_cls = ParamDtypeMixedPrecisionPolicy
     logger.info(
-        f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, reduce_dtype={reduce_dtype}"
+        f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, "
+        f"param_dtype={param_dtype}, reduce_dtype={reduce_dtype}, "
+        f"param_dtype_overrides={param_dtype_maps.override_count} "
+        f"({param_dtype_maps.override_numel:,} parameters)"
     )
 
     fsdp_kwargs = {
-        # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent.
-        "mp_policy": MixedPrecisionPolicy(
-            param_dtype=param_dtype,
-            reduce_dtype=reduce_dtype,
-            cast_forward_inputs=False,
-        ),
         "offload_policy": offload_policy,
         "mesh": mesh,
     }
 
-    for module in modules:
-        fully_shard(module, **fsdp_kwargs)
+    # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent.
+    def make_mp_policy(param_dtype_map):
+        if param_dtype_map:
+            assert param_dtype_policy_cls is not None
+            return param_dtype_policy_cls(
+                param_dtype=param_dtype,
+                reduce_dtype=reduce_dtype,
+                cast_forward_inputs=False,
+                param_dtype_map=param_dtype_map,
+            )
+        return MixedPrecisionPolicy(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+            cast_forward_inputs=False,
+        )
+
+    for module, wrap_map in zip(modules, param_dtype_maps.wrap_maps, strict=True):
+        fully_shard(
+            module,
+            mp_policy=make_mp_policy(wrap_map),
+            **fsdp_kwargs,
+        )
 
-    fully_shard(model, **fsdp_kwargs)
+    fully_shard(
+        model,
+        mp_policy=make_mp_policy(param_dtype_maps.root_map),
+        **fsdp_kwargs,
+    )
 
     return model

miles/backends/fsdp_utils/configs/train_pipeline_config.py

@@ -26,7 +26,9 @@ def register_train_pipeline_config(family: str):
     """Decorator: register a TrainPipelineConfig subclass under a family key (``sd3``, ``wan``, ...)."""
 
     def wrapper(cls):
-        _REGISTRY[family.lower()] = cls
+        model_family = family.lower()
+        cls.model_family = model_family
+        _REGISTRY[model_family] = cls
         return cls
 
     return wrapper
@@ -74,6 +76,7 @@ def get_train_pipeline_config_cls(family: str) -> type[TrainPipelineConfig]:
 class TrainPipelineConfig(abc.ABC):
     """Base class. Subclass per model family."""
 
+    model_family: str | None = None
     lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"]
     needs_timestep_scaling: bool = True
     optimizer_state_allowed_missing: list[str] = []

miles/backends/fsdp_utils/mixed_precision.py

@@ -0,0 +1,116 @@
+"""Compile root-relative precision rules for module-relative FSDP policies.
+
+The public rules always use FQNs from the model root, while each
+``MixedPrecisionPolicy`` resolves its map from the modules passed to that
+``fully_shard`` call:
+
+    root rule: "blocks.0.norm.weight"
+                         |
+                         v
+    model root ----- compile_param_dtype_maps -----+
+                                                    |
+                         +--------------------------+------------------+
+                         |                                             |
+                         v                                             v
+    fully_shard(model.blocks[0])                    fully_shard(model)
+    {"norm.weight": fp32}                           {"root_scale": fp32}
+
+Two passes. Pass 1 expands the patterns against root FQNs into one dtype per
+matched parameter. Pass 2 walks the wraps in fully_shard call order, claiming
+parameters first-wrap-wins — the same visited-set rule FSDP2 itself applies —
+and re-keys each override to the FQN local to its owning wrap; parameters no
+wrap claims land in ``root_map`` under their root FQN.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+
+_SUPPORTED_DTYPES = {
+    "bf16": torch.bfloat16,
+    "fp16": torch.float16,
+    "fp32": torch.float32,
+}
+
+
+def parse_dtype_from_str(dtype_name: str) -> torch.dtype:
+    try:
+        return _SUPPORTED_DTYPES[dtype_name]
+    except KeyError as error:
+        raise ValueError(f"Unsupported dtype {dtype_name!r}") from error
+
+
+@dataclass(frozen=True)
+class CompiledParamDtypeMaps:
+    wrap_maps: list[dict[str, torch.dtype]]  # parallel to the wraps argument
+    root_map: dict[str, torch.dtype]
+    override_count: int
+    override_numel: int
+
+
+def compile_param_dtype_maps(
+    model: nn.Module,
+    wraps: Sequence[nn.Module | Sequence[nn.Module]],
+    root_fqn_patterns: Mapping[str, str],
+    default_dtype: torch.dtype,
+) -> CompiledParamDtypeMaps:
+    """Each entry of ``wraps`` is one child ``fully_shard`` call, in call order — a single module or,
+    like ``fully_shard`` itself, a list of modules grouped into one wrap.
+
+    Patterns apply in declaration order and a later pattern overrides an earlier one, so a narrow
+    rule can carve a parameter back out of a broad one. Within one wrap the runtime map is keyed by
+    wrap-local FQN, so two member modules may share a local FQN only when they agree on its dtype.
+    """
+    assigned_dtypes: dict[nn.Parameter, torch.dtype] = {}
+    for pattern, dtype_name in root_fqn_patterns.items():
+        dtype = parse_dtype_from_str(dtype_name)
+        matched = False
+        for fqn, param in model.named_parameters(remove_duplicate=False):
+            if fnmatch.fnmatchcase(fqn, pattern):
+                matched = True
+                assigned_dtypes[param] = dtype
+        if not matched:
+            raise ValueError(f"FSDP parameter dtype pattern {pattern!r} did not match any parameter")
+    # An assignment equal to the group default compiles to nothing (this is also what a carve-out is).
+    assigned_dtypes = {param: dtype for param, dtype in assigned_dtypes.items() if dtype != default_dtype}
+
+    claimed: set[nn.Parameter] = set()
+    wrap_maps: list[dict[str, torch.dtype]] = []
+    for wrap in wraps:
+        # The runtime map is keyed by wrap-local FQN, so "no override" (None) must collide too:
+        # a mapped FQN would silently apply to every member module sharing that name.
+        seen: dict[str, torch.dtype | None] = {}
+        for module in (wrap,) if isinstance(wrap, nn.Module) else wrap:
+            for local_fqn, param in module.named_parameters():
+                if param in claimed:
+                    continue
+                claimed.add(param)
+                dtype = assigned_dtypes.get(param)
+                if local_fqn in seen and seen[local_fqn] != dtype:
+                    raise ValueError(
+                        f"two parameters in one fully_shard group share the local FQN {local_fqn!r} "
+                        f"but want different dtypes ({seen[local_fqn]} vs {dtype}, None = no override); "
+                        "not supported — the dtype map is keyed by local FQN, so they cannot be told apart"
+                    )
+                seen[local_fqn] = dtype
+        wrap_maps.append({fqn: dtype for fqn, dtype in seen.items() if dtype is not None})
+
+    root_fqns: dict[nn.Parameter, str] = {}
+    for fqn, param in model.named_parameters(remove_duplicate=False):
+        root_fqns.setdefault(param, fqn)
+    root_map = {root_fqns[param]: dtype for param, dtype in assigned_dtypes.items() if param not in claimed}
+    return CompiledParamDtypeMaps(
+        wrap_maps,
+        root_map,
+        len(assigned_dtypes),
+        sum(param.numel() for param in assigned_dtypes),
+    )
+
+
+__all__ = ["CompiledParamDtypeMaps", "compile_param_dtype_maps", "parse_dtype_from_str"]

miles/backends/fsdp_utils/model_backend.py

@@ -6,7 +6,7 @@
 
   - ``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
+  - ``fsdp_parallel_plan``: FSDP wrapping and parameter precision policy
   - ``sequence_parallel_plan`` / ``install_sequence_parallel_attention``:
     the model's SP declaration and attention integration
 
@@ -21,11 +21,13 @@
 import functools
 import importlib
 import logging
+from dataclasses import replace
 from typing import Any
 
 import torch
 from diffusers import DiffusionPipeline
 
+from .models.parallel_plan import FSDPParallelPlan
 from .sequence_parallel.diffusers_dispatch import install_diffusers_usp_patch
 from .sequence_parallel.plan import MILES_SP_PLAN_ATTR, SequenceParallelPlan
 
@@ -62,7 +64,7 @@ def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]:
+    def fsdp_parallel_plan(self, model: torch.nn.Module) -> FSDPParallelPlan:
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -123,8 +125,8 @@ def load_scheduler(self, args) -> Any:
     def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None:
         self._pkg.modeling.enable_gradient_checkpointing(model)
 
-    def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]:
-        return list(self._pkg.parallel_plan.FSDP_NO_SPLIT_MODULES)
+    def fsdp_parallel_plan(self, model: torch.nn.Module) -> FSDPParallelPlan:
+        return self._pkg.parallel_plan.FSDP_PARALLEL_PLAN
 
     def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None:
         self._pkg.attention.set_attention_backend(model, backend)
@@ -153,11 +155,19 @@ def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None:
     def enable_gradient_checkpointing(self, model: torch.nn.Module) -> None:
         model.enable_gradient_checkpointing()
 
-    def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]:
+    def fsdp_parallel_plan(self, model: torch.nn.Module) -> FSDPParallelPlan:
+        if self.config is not None and self.config.model_family is not None:
+            from .models.diffusers import load_fsdp_parallel_plan
+
+            plan = load_fsdp_parallel_plan(self.config.model_family)
+        else:
+            plan = FSDPParallelPlan()
+        if plan.no_split_modules is not None:
+            return plan
         no_split_modules = getattr(model, "_no_split_modules", None)
         if not no_split_modules:
             raise ValueError(f"{model.__class__.__name__} declares no _no_split_modules for FSDP wrapping")
-        return list(no_split_modules)
+        return replace(plan, no_split_modules=tuple(no_split_modules))
 
     def install_sequence_parallel_attention(self, model: torch.nn.Module, parallel_state) -> None:
         install_diffusers_usp_patch(model, parallel_state)

miles/backends/fsdp_utils/models/__init__.py

@@ -2,8 +2,9 @@
 
 Onboarding a model family:
 
-- **Diffusers checkpoint** (has ``model_index.json``): nothing to do here.
-  ``DiffusersModelBackend`` loads it; write only a ``configs/<family>.py``.
+- **Diffusers checkpoint** (has ``model_index.json``):
+  ``DiffusersModelBackend`` loads it; add
+  ``models/diffusers/<family>/parallel_plan.py`` for its FSDP precision plan.
 
 - **Native modeling** (official repo code, non-diffusers checkpoint): add a
   package ``models/<family>/`` with at least:
@@ -12,7 +13,7 @@
   - ``modeling.py`` — ``load_scheduler``, ``enable_gradient_checkpointing``,
     optional ``flash_attention_entrypoints`` /
     ``required_flash_kernel_label`` for deterministic flash patching
-  - ``parallel_plan.py`` — ``FSDP_NO_SPLIT_MODULES``,
+  - ``parallel_plan.py`` — ``FSDP_PARALLEL_PLAN``,
     ``sequence_parallel_plan``
   - ``attention.py`` — ``set_attention_backend``
 

miles/backends/fsdp_utils/models/diffusers/__init__.py

@@ -0,0 +1,14 @@
+import importlib
+
+from ..parallel_plan import FSDPParallelPlan
+
+
+def load_fsdp_parallel_plan(model_family: str) -> FSDPParallelPlan:
+    module = importlib.import_module(f"{__name__}.{model_family}.parallel_plan")
+    plan = module.FSDP_PARALLEL_PLAN
+    if not isinstance(plan, FSDPParallelPlan):
+        raise TypeError(f"{module.__name__}.FSDP_PARALLEL_PLAN must be an FSDPParallelPlan")
+    return plan
+
+
+__all__ = ["load_fsdp_parallel_plan"]

miles/backends/fsdp_utils/models/diffusers/qwen_image/__init__.py

@@ -0,0 +1 @@
+

miles/backends/fsdp_utils/models/diffusers/qwen_image/parallel_plan.py

@@ -0,0 +1,4 @@
+from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan
+
+
+FSDP_PARALLEL_PLAN = FSDPParallelPlan()

miles/backends/fsdp_utils/models/diffusers/sd3/__init__.py

@@ -0,0 +1 @@
+

miles/backends/fsdp_utils/models/diffusers/sd3/parallel_plan.py

@@ -0,0 +1,4 @@
+from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan
+
+
+FSDP_PARALLEL_PLAN = FSDPParallelPlan()

miles/backends/fsdp_utils/models/diffusers/wan2_2/__init__.py

@@ -0,0 +1 @@
+