refactor(fsdp)(4.2/5): add per-parameter dtype patch

[Rockdu]

What

Adds a source-guarded PyTorch 2.11 monkeypatch for exact per-parameter FSDP2 mixed-precision policies. The public module dispatches by the installed Torch base version and lazily loads the version-specific implementation, so unsupported versions never import PyTorch 2.11 internals. This PR does not include Miles actor or parallel-plan integration.

The distributed suite covers real full-size Diffusers blocks, exact output/gradient parity, all-gathered dtype and shape checks, FQN validation, mixed-dtype communication, and adversarial dtype/shape combinations.

Most of the code changes are due to CI ported from torch to ensure FSDP functions normally.

API (PyTorch 2.11 only)

import torch
from torch.distributed.fsdp import fully_shard

from miles.backends.fsdp_utils.fsdp_param_dtype_patch import (
    ParamDtypeMixedPrecisionPolicy,
    apply_param_dtype_map_patch,
)

apply_param_dtype_map_patch()
policy = ParamDtypeMixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    output_dtype=None,
    cast_forward_inputs=True,
    param_dtype_map={
        "attn1.to_q.weight": torch.float32,
        "attn1.to_q.bias": torch.float32,
    },
)
fully_shard(model, mp_policy=policy)

For a single FSDP wrap spanning multiple modules:

layer_norms = [block.norm for block in model.blocks]
layer_norm_policy = ParamDtypeMixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    cast_forward_inputs=False,
    param_dtype_map={
        "weight": torch.float32,
        "bias": torch.float32,
    },
)

# One grouped wrap manages both LayerNorms. The exact `weight` and `bias`
# FQNs are broadcast to every matching module in this call.
fully_shard(layer_norms, mp_policy=layer_norm_policy)

param_dtype_map is a sparse mapping of exact FQNs, scoped to the modules managed by that fully_shard call. Unlisted parameters use param_dtype. When one call groups multiple modules, an exact FQN is broadcast to every matching managed parameter (for example, weight across multiple LayerNorms). Multiple aliases of one shared parameter may specify the same dtype; conflicting alias dtypes are rejected. Wildcards are intentionally not part of the patch API; Miles' model integration compiles patterns to exact FQNs separately. Mixed trainable parameter dtypes require an explicit common reduce_dtype.

Compatibility

• Supports the PyTorch 2.11 base version, independent of the CUDA wheel suffix.
• Dispatches before loading private Torch internals; unsupported versions fail with a clear error.
• Keeps the complete source-guarded implementation in _fsdp_param_dtype_patch_2_11.py, leaving room for separate 2.12/2.13 implementations.
• Pins the Torch dependency stack to torch==2.11.0, torchaudio==2.11.0, and torchvision==0.26.0 so bare CPU CI matches the production GPU image.
• Verifies every patched upstream function by source hash before installation.
• Preserves the standard MixedPrecisionPolicy and empty-map paths.
• Tested mapped floating dtypes: float8_e4m3fn, float8_e4m3fnuz, float8_e5m2, float8_e5m2fnuz, float16, bfloat16, float32, and float64.
• float8_e8m0fnu and packed float4_e2m1fn_x2 are intentionally unsupported because PyTorch 2.11 does not implement the required _foreach_copy_ cast path. The patch does not add a slow per-tensor fallback.
• Integral, boolean, and complex tensors may coexist and participate in forward as ignored frozen parameters; upstream FSDP2 cannot manage non-floating nn.Parameter objects because it reconstructs sharded parameters with gradients enabled.

Performance guard

Wan2.2 was run with identical settings and no per-parameter dtype overrides, changing only whether the patch was installed:

• Unpatched: 70.8656s per stable actor training step
• Patched: 70.8985s per stable actor training step
• Delta: +0.0329s (+0.046%), within run-to-run noise

Test plan

• Full repository pre-commit suite.
• Version-dispatch unit tests and successful lazy patch application on torch==2.11.0+cu129.
• Selected upstream PyTorch 2.11 mixed-precision, collective, and frozen-parameter regression tests under run-ci-torch.
• Four-rank H200: full-size single Wan2.2 and LTX2.3 blocks under FSDP 1x4 and HSDP 2x2.
• Bitwise output and per-parameter gradient parity across unpatched, patched standard-policy, and all-BF16-map paths.
• Bitwise parity between explicit FP32 child FSDP wrappers inside a BF16 root and one BF16 root wrapper with the same parameters mapped to FP32 by exact FQN.
• Direct resolver unit coverage plus a four-rank two-block topology where one grouped FSDP call wraps both LayerNorms, each block is wrapped separately, gathered LayerNorm parameters are FP32, and outputs/gradients are bitwise-equal to separate FP32 LayerNorm wraps.
• Forward pre-hooks verify every gathered parameter's FQN, dtype, shape, and numel, including dim-0 padding for prime-sized tensors.
• Twelve distributed validation cases: duplicate-FQN broadcast across grouped modules, separate-wrap scoping, shared aliases with equal and conflicting dtypes, unknown FQNs, grouped LayerNorm parity, reduce-dtype requirements, frozen overrides, mixed forward/backward, empty-map delegation, and empty gradients.
• Two adversarial four-rank cases: single-module FSDP and shared-parameter multi-module HSDP with dimensions 7, 97, 101, and 103 across the supported floating dtype matrix plus ignored integer/boolean/complex parameters.