[DO NOT MERGE] refactor(fsdp)(4.1/5): per-parameter weight-precision control by FSDP2 expansion patch

.claude/skills/install-miles-diffusion/install.sh

@@ -10,7 +10,7 @@
 #   SGLANG_BRANCH   sglang branch to check out (default: sglang-diffusion-rollout-test)
 #   SGLANG_COMMIT   sglang commit SHA to pin (default: pinned working SHA below)
 #   CUDA_VER        torch cuda tag (default: 12.9 -> cu129)
-#   TORCH_VER       torch version (default: 2.9.1)
+#   TORCH_VER       torch version (default: 2.11.0)
 #
 # All package versions are pinned. Pins reflect the currently-validated working
 # environment for miles-diffusion + sglang-diffusion + flow_grpo OCR reward.
@@ -27,7 +27,7 @@ set -euo pipefail
 ENV_NAME="${ENV_NAME:-miles-diffusion}"
 PY_VER="${PY_VER:-3.11}"
 CUDA_VER="${CUDA_VER:-12.9}"
-TORCH_VER="${TORCH_VER:-2.9.1}"
+TORCH_VER="${TORCH_VER:-2.11.0}"
 SGLANG_REPO="${SGLANG_REPO:-https://github.com/Rockdu/sglang.git}"
 SGLANG_BRANCH="${SGLANG_BRANCH:-sglang-diffusion-rollout-test}"
 SGLANG_COMMIT="${SGLANG_COMMIT:-0372158dd66bc7cb0740c733bd60047db790ec7d}"

miles/backends/fsdp_utils/actor.py

@@ -38,6 +38,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
 from .parallel import create_fsdp_parallel_state
 from .sequence_parallel.plan import apply_sequence_parallel
 
@@ -137,12 +138,14 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 raise RuntimeError(f"{component} did not honor meta initialization")
             checkpoint.sync_model_dtypes(model)
             full_state = model.state_dict() if rank == 0 else {}
+            fsdp_parallel_plan = self.model_backend.fsdp_parallel_plan(model)
             model = apply_fsdp2(
                 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),
+                no_split_modules=fsdp_parallel_plan.no_split_modules,
+                param_dtype_patterns=fsdp_parallel_plan.param_dtype_patterns,
             )
             checkpoint.broadcast_full_state_to_fsdp(
                 model,
@@ -609,7 +612,14 @@ 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,
+    mesh=None,
+    cpu_offload=False,
+    args=None,
+    no_split_modules=None,
+    param_dtype_patterns=None,
+):
     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
 
     offload_policy = CPUOffloadPolicy() if cpu_offload else None
@@ -621,19 +631,49 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules
 
     param_dtype = _resolve_dtype(args.diffusion_forward_dtype)
     reduce_dtype = _resolve_dtype(args.fsdp_reduce_dtype)
+    param_dtype_maps = compile_param_dtype_maps(
+        model,
+        modules,
+        param_dtype_patterns or {},
+        param_dtype,
+    )
+    has_param_dtype_overrides = bool(
+        param_dtype_maps.module_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 = {
-        "mp_policy": MixedPrecisionPolicy(
-            param_dtype=param_dtype,
-            reduce_dtype=reduce_dtype,
-        ),
         "offload_policy": offload_policy,
         "mesh": mesh,
     }
 
+    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,
+                param_dtype_map=param_dtype_map,
+            )
+        return MixedPrecisionPolicy(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+        )
+
     if args.gradient_checkpointing:
         # MixedPrecisionPolicy does not cast buffers; a buffer above param_dtype
         # makes the ckpt recompute dtype-diverge from the forward and abort.
@@ -644,8 +684,16 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules
                     module.register_buffer(name, buf.to(param_dtype), persistent=persistent)
 
     for module in modules:
-        fully_shard(module, **fsdp_kwargs)
+        fully_shard(
+            module,
+            mp_policy=make_mp_policy(param_dtype_maps.module_maps.get(module)),
+            **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] = []