From a439e873ab6b2bcc4c179b82c1d095fdf5648548 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 5 Aug 2026 20:11:24 -0700 Subject: [PATCH 1/7] feat(fsdp): preserve selected parameters in FP32 Compile model precision patterns into FSDP-local maps so numerically sensitive Wan parameters stay in FP32 without giving up BF16 compute. Co-authored-by: Cursor --- .../skills/install-miles-diffusion/install.sh | 4 +- .github/workflows/docker-build.yml | 4 + .github/workflows/pr-test.yml | 2 +- docker/Dockerfile | 3 + miles/backends/fsdp_utils/actor.py | 63 +- .../configs/train_pipeline_config.py | 1 + miles/backends/fsdp_utils/configs/wan2_2.py | 5 + .../fsdp_utils/fsdp_param_dtype_patch.py | 591 ++++++++++++++++++ miles/backends/fsdp_utils/mixed_precision.py | 104 +++ .../fsdp_utils/_param_dtype_map_worker.py | 86 +++ .../fsdp_utils/test_param_dtype_map.py | 39 ++ .../fsdp_utils/test_mixed_precision.py | 94 +++ 12 files changed, 985 insertions(+), 11 deletions(-) create mode 100644 miles/backends/fsdp_utils/fsdp_param_dtype_patch.py create mode 100644 miles/backends/fsdp_utils/mixed_precision.py create mode 100644 tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py create mode 100644 tests/fast-gpu/backends/fsdp_utils/test_param_dtype_map.py create mode 100644 tests/fast/backends/fsdp_utils/test_mixed_precision.py diff --git a/.claude/skills/install-miles-diffusion/install.sh b/.claude/skills/install-miles-diffusion/install.sh index 2f24f63d..243ddded 100755 --- a/.claude/skills/install-miles-diffusion/install.sh +++ b/.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}" diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index 8d051ca0..d3adbfa6 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -26,6 +26,10 @@ on: paths: - 'docker/Dockerfile' - 'requirements.txt' + - 'miles/backends/fsdp_utils/actor.py' + - 'miles/backends/fsdp_utils/configs/**' + - 'miles/backends/fsdp_utils/fsdp_param_dtype_patch.py' + - 'miles/backends/fsdp_utils/mixed_precision.py' concurrency: group: docker-build diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 6578a042..be2895ca 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -68,7 +68,7 @@ jobs: # Compare the PR merge commit with its current base. CHANGED_FILES=$(git diff --name-only HEAD^1 HEAD) # Same paths as docker-build.yml's post-merge filter. - if printf '%s\n' "$CHANGED_FILES" | grep -qE '^(docker/Dockerfile|requirements\.txt)$'; then + if printf '%s\n' "$CHANGED_FILES" | grep -qE '^(docker/Dockerfile|requirements\.txt|miles/backends/fsdp_utils/((actor|fsdp_param_dtype_patch|mixed_precision)\.py|configs/.*\.py))$'; then echo "changed=true" >> "$GITHUB_OUTPUT" else echo "changed=false" >> "$GITHUB_OUTPUT" diff --git a/docker/Dockerfile b/docker/Dockerfile index 6cb7b1fe..bd89a798 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,6 +1,9 @@ ARG SGLANG_IMAGE_TAG=v0.5.12-cu129 FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang +ARG TORCH_VERSION=2.11.0+cu129 +RUN python -c "import torch; assert torch.__version__ == '${TORCH_VERSION}', f'expected torch ${TORCH_VERSION}, got {torch.__version__}'" + # ======================================== Setup ============================================= WORKDIR /root/ diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d9430378..2258b9ab 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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 @@ -143,6 +144,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), + param_dtype_patterns=self.train_pipeline_config.fsdp_param_dtype_patterns, ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -609,7 +611,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 +630,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 +683,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 diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index c392ea57..8ffbfe69 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -75,6 +75,7 @@ class TrainPipelineConfig(abc.ABC): """Base class. Subclass per model family.""" lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"] + fsdp_param_dtype_patterns: dict[str, str] = {} needs_timestep_scaling: bool = True optimizer_state_allowed_missing: list[str] = [] # Case-insensitive substrings matched against the checkpoint name (--diffusion-model). diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 7e227cd2..7ebb11fd 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -11,6 +11,11 @@ @register_train_pipeline_config("wan2_2") class Wan2_2TrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("wan2.2", "wan-2.2") + fsdp_param_dtype_patterns = { + "*scale_shift_table": "fp32", + "*time_embedder*": "fp32", + "*.norm2.*": "fp32", + } # High-noise expert ("transformer") handles t >= boundary, low-noise expert # ("transformer_2") the rest. boundary_ratio = 0.875 diff --git a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py new file mode 100644 index 00000000..24223f51 --- /dev/null +++ b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py @@ -0,0 +1,591 @@ +from __future__ import annotations + +import hashlib +import inspect +import math +import types +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from typing import cast + +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.fsdp import MixedPrecisionPolicy as TorchMixedPrecisionPolicy +from torch.distributed.fsdp._fully_shard import _fsdp_collectives, _fsdp_param_group +from torch.distributed.fsdp._fully_shard._fsdp_api import ReduceScatter +from torch.distributed.fsdp._fully_shard._fsdp_collectives import ( + _div_if_needed, + _get_device_handle, + _get_dim0_padded_size, + _get_gradient_divide_factors, + _raise_assert_with_print, + _to_dtype_if_needed, + compiled_autograd_enabled, +) +from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState +from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup +from torch.distributed.tensor import DTensor + + +_EXPECTED_TORCH_VERSION = "2.11.0+cu129" +_PATCH_SENTINEL = "_miles_param_dtype_map_patch_applied" +_SOURCE_HASHES = { + "FSDPParam.init_dtype_attrs": "2cc968770804055cdde959db7cfa47b92a1137badcdaf2e448555741c2fd282c", + "FSDPParamGroup.__init__": "323868f31033bb5696eaf30f278838901c0d6b7e0da031c517177707b716409b", + "_get_param_all_gather_inputs": "a70fbae57b8aa3dc669d01a409d30c74428546bc4d69635a90cffe9380a5785f", + "foreach_reduce": "46bcdaa0df40823e13922359db131cf329f8adf12eefa72ff16ceb41f9650d90", + "foreach_reduce_scatter_copy_in": "559a065467abbaa578f348bd6a8c6478cfc8ee516e602227b795bc3f6727eb22", +} + +_ORIGINAL_PARAM_GROUP_INIT = FSDPParamGroup.__init__ +_ORIGINAL_INIT_DTYPE_ATTRS = FSDPParam.init_dtype_attrs +_ORIGINAL_GET_PARAM_ALL_GATHER_INPUTS = _fsdp_collectives._get_param_all_gather_inputs +_ORIGINAL_FOREACH_REDUCE = _fsdp_collectives.foreach_reduce +_ORIGINAL_FOREACH_REDUCE_SCATTER_COPY_IN = _fsdp_collectives.foreach_reduce_scatter_copy_in + + +@dataclass(frozen=True) +class ParamDtypeMixedPrecisionPolicy(TorchMixedPrecisionPolicy): + param_dtype_map: Mapping[str, torch.dtype] | None = None + + +def _source_hash(fn) -> str: + return hashlib.sha256(inspect.getsource(fn).encode()).hexdigest() + + +def _verify_source(name: str, fn) -> None: + actual = _source_hash(fn) + expected = _SOURCE_HASHES[name] + if actual != expected: + raise RuntimeError( + f"Cannot apply the Miles FSDP patch: source hash for {name} changed " + f"(expected {expected}, got {actual})" + ) + + +def _bind_to_collectives(fn, name: str, *, no_grad: bool): + bound = types.FunctionType( + fn.__code__, + _fsdp_collectives.__dict__, + name, + fn.__defaults__, + fn.__closure__, + ) + bound.__kwdefaults__ = fn.__kwdefaults__ + bound.__annotations__ = fn.__annotations__ + bound.__module__ = _fsdp_collectives.__name__ + return torch.no_grad()(bound) if no_grad else bound + + +def _patched_param_group_init( + self, + params: list[nn.Parameter], + modules: tuple[nn.Module, ...], + mesh_info, + post_forward_mesh_info, + device: torch.device, + shard_placement_fn, + mp_policy: TorchMixedPrecisionPolicy, + offload_policy, +) -> None: + param_dtype_map = ( + mp_policy.param_dtype_map + if isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) + else None + ) + if param_dtype_map is None: + _ORIGINAL_PARAM_GROUP_INIT( + self, + params, + modules, + mesh_info, + post_forward_mesh_info, + device, + shard_placement_fn, + mp_policy, + offload_policy, + ) + return + + # MILES_PATCH_BEGIN: fsdp-param-dtype-map + managed_params = set(params) + fqn_to_param: dict[str, nn.Parameter] = {} + for module in modules: + for fqn, param in module.named_parameters(): + if param not in managed_params: + continue + previous = fqn_to_param.get(fqn) + if previous is not None and previous is not param: + raise ValueError( + f"param_dtype_map FQN {fqn!r} is ambiguous across the fully_shard modules" + ) + fqn_to_param[fqn] = param + unknown_fqns = sorted(set(param_dtype_map).difference(fqn_to_param)) + if unknown_fqns: + raise ValueError( + "param_dtype_map contains FQNs that do not name a parameter managed " + f"by this fully_shard call: {unknown_fqns}" + ) + param_overrides = { + fqn_to_param[fqn]: dtype for fqn, dtype in param_dtype_map.items() + } + effective_dtypes = { + param_overrides.get(param, mp_policy.param_dtype) or param.dtype + for param in params + if param.requires_grad + } + if len(effective_dtypes) > 1 and mp_policy.reduce_dtype is None: + raise ValueError("Mixed parameter dtypes require an explicit reduce_dtype") + # MILES_PATCH_END: fsdp-param-dtype-map + + _ORIGINAL_PARAM_GROUP_INIT( + self, + params, + modules, + mesh_info, + post_forward_mesh_info, + device, + shard_placement_fn, + mp_policy, + offload_policy, + ) + + # MILES_PATCH_BEGIN: fsdp-param-dtype-map + for fsdp_param, param in zip(self.fsdp_params, params, strict=True): + override = param_overrides.get(param) + fsdp_param._param_dtype_override = override + effective_param_dtype = ( + override if override is not None else mp_policy.param_dtype + ) + fsdp_param.mp_policy = replace( + mp_policy, + param_dtype=effective_param_dtype, + param_dtype_map=None, + ) + # MILES_PATCH_END: fsdp-param-dtype-map + + +def _patched_init_dtype_attrs( + self: FSDPParam, + mp_policy: TorchMixedPrecisionPolicy, +) -> None: + if not isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy): + _ORIGINAL_INIT_DTYPE_ATTRS(self, mp_policy) + return + + # MILES_PATCH_BEGIN: fsdp-param-dtype-map + param_dtype = ( + self._param_dtype_override + if self._param_dtype_override is not None + else mp_policy.param_dtype + ) + reduce_dtype = mp_policy.reduce_dtype + self.orig_dtype = self.sharded_param.dtype + if mp_policy.param_dtype_map is None and reduce_dtype == param_dtype: + reduce_dtype = None + if param_dtype == self.orig_dtype: + param_dtype = None + self.param_dtype = param_dtype + self.reduce_dtype = reduce_dtype + # MILES_PATCH_END: fsdp-param-dtype-map + + +# Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. +def _patched_get_param_all_gather_inputs( + fsdp_params: list[FSDPParam], +) -> list[list[torch.Tensor]]: + if compiled_autograd_enabled(): + return [fsdp_param.all_gather_inputs for fsdp_param in fsdp_params] + + # Intentionally try to run a fast-path that bypasses abstractions for the + # common FSDP case of bf16/fp32 mixed precision in order to use foreach + # copy for lower CPU overhead and more efficient copying in eager + def use_foreach_copy(fsdp_param: FSDPParam) -> bool: + return ( + fsdp_param.param_dtype is not None + and not fsdp_param.offload_to_cpu + and not hasattr(fsdp_param._sharded_local_tensor, "fsdp_pre_all_gather") + ) + + param_all_gather_inputs: list[list[torch.Tensor]] = [[] for _ in fsdp_params] + # MILES_PATCH_BEGIN: fsdp-param-dtype-map + foreach_copy_infos: dict[ + torch.dtype, tuple[list[int], list[torch.Tensor], list[int]] + ] = {} + + # 1st pass: for foreach-copy parameters, get inputs and metadata for the + # foreach copy, and for the others, actually get their all-gather inputs + for i, fsdp_param in enumerate(fsdp_params): + if use_foreach_copy(fsdp_param): + param_dtype = cast(torch.dtype, fsdp_param.param_dtype) + indices, inputs, input_numels = foreach_copy_infos.setdefault( + param_dtype, ([], [], []) + ) + indices.append(i) + all_gather_input = ( + fsdp_param._sharded_param_data + if fsdp_param.sharded_state == ShardedState.SHARDED + else cast(torch.Tensor, fsdp_param._sharded_post_forward_param_data) + ) + inputs.append(all_gather_input) + input_numels.append(all_gather_input.numel()) + else: + param_all_gather_inputs[i] = fsdp_param.all_gather_inputs + + # 2nd pass: use foreach copy to compute the remaining all-gather inputs + for param_dtype, (indices, inputs, input_numels) in foreach_copy_infos.items(): + device = fsdp_params[indices[0]].device + flat_foreach_copy_input = torch.empty( + (sum(input_numels),), device=device, dtype=param_dtype + ) + splits = torch.split(flat_foreach_copy_input, input_numels) + torch._foreach_copy_(splits, inputs) + for i, split in zip(indices, splits): + param_all_gather_inputs[i] = [split] + # MILES_PATCH_END: fsdp-param-dtype-map + + return param_all_gather_inputs + + +# Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. +def _patched_foreach_reduce( + fsdp_params: list[FSDPParam], + unsharded_grads: list[torch.Tensor], + reduce_scatter_group: dist.ProcessGroup, + reduce_scatter_stream: torch.Stream, + reduce_scatter_comm: ReduceScatter, + orig_dtype: torch.dtype | None, + reduce_dtype: torch.dtype | None, + device: torch.device, + gradient_divide_factor: float | None, + all_reduce_group: dist.ProcessGroup | None, + all_reduce_stream: torch.Stream, + all_reduce_grads: bool, + partial_reduce_output: torch.Tensor | None, + all_reduce_hook: Callable[[torch.Tensor], None] | None, + force_sum_reduction_for_comms: bool = False, +) -> tuple[ + torch.Tensor, + torch.Event, + torch.Event, + torch.Tensor | None, + torch.Event | None, + torch.Tensor | None, +]: + """ + ``unsharded_grads`` owns the references to the gradients computed by + autograd, so clearing the list frees the gradients. + """ + + # MILES_PATCH_BEGIN: fsdp-param-dtype-map + grad_dtypes = {grad.dtype for grad in unsharded_grads} + if reduce_dtype is None and len(grad_dtypes) != 1: + _raise_assert_with_print( + "FSDP reduce-scatter requires an explicit reduce dtype for mixed " + f"gradient dtypes but got {grad_dtypes}" + ) + reduce_dtype = reduce_dtype or unsharded_grads[0].dtype + # MILES_PATCH_END: fsdp-param-dtype-map + (predivide_factor, postdivide_factor, reduce_scatter_op, all_reduce_op) = ( + _get_gradient_divide_factors( + reduce_scatter_group, + all_reduce_group, + reduce_dtype, + device.type, + gradient_divide_factor, + force_sum_reduction_for_comms, + ) + ) + + if reduce_scatter_group is None: + world_size = 1 + else: + world_size = reduce_scatter_group.size() + device_handle = _get_device_handle(device.type) + current_stream = device_handle.current_stream() + + if world_size > 1: + for i, (fsdp_param, unsharded_grad) in enumerate( + zip(fsdp_params, unsharded_grads) + ): + if (shard_dim := fsdp_param.fsdp_placement.dim) == 0: + continue + if unsharded_grad.size(shard_dim) % world_size != 0: + raise AssertionError( + f"Shard({shard_dim}) requires even sharding: {unsharded_grad.size()=} {world_size=}" + ) + chunks = torch.chunk(unsharded_grad, world_size, dim=shard_dim) + unsharded_grads[i] = torch.cat(chunks, dim=0) + + padded_unsharded_sizes = tuple( + _get_dim0_padded_size(grad.size(), world_size) for grad in unsharded_grads + ) + reduce_scatter_input_numel = sum(s.numel() for s in padded_unsharded_sizes) + reduce_scatter_output_numel = reduce_scatter_input_numel // world_size + reduce_scatter_input = reduce_scatter_comm.allocate( + (reduce_scatter_input_numel,), + dtype=reduce_dtype, + device=device, + ) + + foreach_reduce_scatter_copy_in( + unsharded_grads, reduce_scatter_input, world_size + ) + + # Only after the copy-in finishes can we free the gradients + unsharded_grads.clear() + reduce_scatter_stream.wait_stream(current_stream) + all_reduce_input = None + all_reduce_event = None + + with device_handle.stream(reduce_scatter_stream): + reduce_output = reduce_scatter_comm.allocate( + (reduce_scatter_output_numel,), + dtype=reduce_dtype, + device=device, + ) + _div_if_needed(reduce_scatter_input, predivide_factor) + if world_size > 1: + reduce_scatter_comm( + output_tensor=reduce_output, + input_tensor=reduce_scatter_input, + group=reduce_scatter_group, + op=reduce_scatter_op, + ) + else: + # For single GPU, just copy the input to output (no actual reduce-scatter needed), and + # account for a possible gradient_divide_factor. + if gradient_divide_factor is not None: + reduce_output.copy_( + reduce_scatter_input / gradient_divide_factor + ) + else: + reduce_output.copy_(reduce_scatter_input) + reduce_scatter_event = reduce_scatter_stream.record_event() + post_reduce_stream = reduce_scatter_stream + if all_reduce_group is not None: # HSDP or DDP/replicate + # Accumulations must run in the reduce-scatter stream + if not all_reduce_grads: + if partial_reduce_output is not None: + partial_reduce_output += reduce_output + else: + partial_reduce_output = reduce_output + return ( + reduce_scatter_input, + reduce_scatter_event, + post_reduce_stream.record_event(), + all_reduce_input, + all_reduce_event, + partial_reduce_output, + ) + if partial_reduce_output is not None: + reduce_output += partial_reduce_output + post_reduce_stream = all_reduce_stream + if world_size >= 1: + all_reduce_stream.wait_stream(reduce_scatter_stream) + else: + all_reduce_stream.wait_stream(current_stream) + with device_handle.stream(all_reduce_stream): + dist.all_reduce( + reduce_output, + group=all_reduce_group, + op=all_reduce_op, + ) + all_reduce_input = reduce_output + all_reduce_event = all_reduce_stream.record_event() + # -- END: ops in reduce_scatter stream + + if all_reduce_hook is not None: + # Execute user-specified all reduce hook. + # If native HSDP is used, this is executed after the HSDP all reduce. + # If 1-d FSDP is used, this is executed post reduce-scatter. + post_reduce_stream = all_reduce_stream + all_reduce_stream.wait_stream(reduce_scatter_stream) + with device_handle.stream(all_reduce_stream): + all_reduce_hook(reduce_output) + # -- END: ops post reduce_scatter + + with device_handle.stream(post_reduce_stream): + _div_if_needed(reduce_output, postdivide_factor) + reduce_output = _to_dtype_if_needed(reduce_output, orig_dtype) + # View out and accumulate sharded gradients + flat_grad_offset = 0 # [0, reduce_scatter_output_numel - 1] + for padded_unsharded_size, fsdp_param in zip( + padded_unsharded_sizes, fsdp_params + ): + # Assume even sharding for Shard(i), i > 0; otherwise would require + # copy-out for contiguous strides + new_sharded_grad = torch.as_strided( + reduce_output, + size=fsdp_param.sharded_size, + stride=fsdp_param.contiguous_sharded_stride, + storage_offset=flat_grad_offset, + ) + to_accumulate_grad = fsdp_param.sharded_param.grad is not None + if fsdp_param.offload_to_cpu: + # Only overlap the D2H copy (copying to pinned memory) if not + # accumulating gradients since the CPU add kernel depends on + # the copy result and we cannot run the add as a callback + non_blocking = fsdp_param.pin_memory and not to_accumulate_grad + # Since the GPU sharded gradient is allocated in the RS stream, + # we can free it here by not keeping a ref without waiting for + # the D2H copy since future RS-stream ops run after the copy + new_sharded_grad = new_sharded_grad.to( + torch.device("cpu"), non_blocking=non_blocking + ) + if non_blocking: + # Record an event on which to block the CPU thread to + # ensure that the D2H copy finishes before the optimizer + fsdp_param.grad_offload_event = ( + post_reduce_stream.record_event() + ) + if to_accumulate_grad: + if not isinstance(fsdp_param.sharded_param.grad, DTensor): + raise AssertionError( + f"Expected fsdp_param.sharded_param.grad to be DTensor, got {type(fsdp_param.sharded_param.grad)}" + ) + fsdp_param.sharded_param.grad._local_tensor += new_sharded_grad + else: + new_sharded_dtensor_grad = fsdp_param.to_sharded_dtensor( + new_sharded_grad + ) + fsdp_param.sharded_param.grad = new_sharded_dtensor_grad + if not compiled_autograd_enabled(): + for hook in ( + getattr( + fsdp_param.sharded_param, + "_post_accumulate_grad_hooks", + {}, + ) + or {} + ).values(): + hook(fsdp_param.sharded_param) + padded_sharded_numel = ( + padded_unsharded_size.numel() // world_size + ) + flat_grad_offset += padded_sharded_numel + post_reduce_event = post_reduce_stream.record_event() + # The RS output is allocated in the RS stream and used in the default + # stream (for optimizer). To ensure its memory is not reused for later + # RSs, we do not need extra synchronization since the sharded parameters + # hold refs through the end of backward. + return ( + reduce_scatter_input, + reduce_scatter_event, + post_reduce_event, + all_reduce_input, + all_reduce_event, + None, + ) + + +# Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. +def _patched_foreach_reduce_scatter_copy_in( + unsharded_grads: list[torch.Tensor], + reduce_scatter_input: torch.Tensor, + world_size: int, +) -> None: + reduce_scatter_input = reduce_scatter_input.view(world_size, -1) + # MILES_PATCH_BEGIN: fsdp-param-dtype-map + if len({grad.dtype for grad in unsharded_grads}) == 1: + torch.ops.fsdp.chunk_cat( + unsharded_grads, + dim=0, + num_chunks=world_size, + out=reduce_scatter_input, + ) + return + + copy_infos: dict[ + torch.dtype, tuple[list[torch.Tensor], list[torch.Tensor]] + ] = {} + padding_views: list[torch.Tensor] = [] + output_offset = 0 + for grad in unsharded_grads: + chunk_size = math.ceil(grad.size(0) / world_size) + trailing_numel = grad.numel() // grad.size(0) + padded_chunk_numel = chunk_size * trailing_numel + destinations, sources = copy_infos.setdefault(grad.dtype, ([], [])) + for rank in range(world_size): + destination = reduce_scatter_input[rank].narrow( + 0, output_offset, padded_chunk_numel + ) + start = rank * chunk_size + length = min(chunk_size, max(grad.size(0) - start, 0)) + if length > 0: + source = grad.narrow(0, start, length).reshape(-1) + destinations.append( + destination.narrow(0, 0, source.numel()) + ) + sources.append(source) + if length < chunk_size: + padding_start = length * trailing_numel + padding_views.append( + destination.narrow( + 0, + padding_start, + padded_chunk_numel - padding_start, + ) + ) + output_offset += padded_chunk_numel + + for destinations, sources in copy_infos.values(): + torch._foreach_copy_(destinations, sources) + if padding_views: + torch._foreach_zero_(padding_views) + # MILES_PATCH_END: fsdp-param-dtype-map + + +def apply_param_dtype_map_patch() -> None: + if getattr(_fsdp_collectives, _PATCH_SENTINEL, False): + return + if torch.__version__ != _EXPECTED_TORCH_VERSION: + raise RuntimeError( + "The Miles FSDP param-dtype patch requires " + f"torch=={_EXPECTED_TORCH_VERSION}, got {torch.__version__}" + ) + + _verify_source("FSDPParamGroup.__init__", _ORIGINAL_PARAM_GROUP_INIT) + _verify_source("FSDPParam.init_dtype_attrs", _ORIGINAL_INIT_DTYPE_ATTRS) + _verify_source( + "_get_param_all_gather_inputs", + _ORIGINAL_GET_PARAM_ALL_GATHER_INPUTS, + ) + _verify_source("foreach_reduce", _ORIGINAL_FOREACH_REDUCE) + _verify_source( + "foreach_reduce_scatter_copy_in", + _ORIGINAL_FOREACH_REDUCE_SCATTER_COPY_IN, + ) + + get_param_all_gather_inputs = _bind_to_collectives( + _patched_get_param_all_gather_inputs, + "_get_param_all_gather_inputs", + no_grad=True, + ) + foreach_reduce_scatter_copy_in = _bind_to_collectives( + _patched_foreach_reduce_scatter_copy_in, + "foreach_reduce_scatter_copy_in", + no_grad=False, + ) + foreach_reduce = _bind_to_collectives( + _patched_foreach_reduce, + "foreach_reduce", + no_grad=True, + ) + + FSDPParamGroup.__init__ = _patched_param_group_init + FSDPParam.init_dtype_attrs = _patched_init_dtype_attrs + _fsdp_collectives._get_param_all_gather_inputs = get_param_all_gather_inputs + _fsdp_collectives.foreach_reduce_scatter_copy_in = ( + foreach_reduce_scatter_copy_in + ) + _fsdp_collectives.foreach_reduce = foreach_reduce + _fsdp_param_group.foreach_reduce = foreach_reduce + setattr(_fsdp_collectives, _PATCH_SENTINEL, True) + + +__all__ = [ + "ParamDtypeMixedPrecisionPolicy", + "apply_param_dtype_map_patch", +] diff --git a/miles/backends/fsdp_utils/mixed_precision.py b/miles/backends/fsdp_utils/mixed_precision.py new file mode 100644 index 00000000..d9f9edf8 --- /dev/null +++ b/miles/backends/fsdp_utils/mixed_precision.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import fnmatch +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import torch +from torch import nn + + +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} + + +@dataclass(frozen=True) +class CompiledParamDtypeMaps: + module_maps: dict[nn.Module, dict[str, torch.dtype]] + root_map: dict[str, torch.dtype] + override_count: int + override_numel: int + + +def compile_param_dtype_maps( + model: nn.Module, + modules: Sequence[nn.Module], + patterns: Mapping[str, str], + default_dtype: torch.dtype, +) -> CompiledParamDtypeMaps: + if not patterns: + return CompiledParamDtypeMaps({}, {}, 0, 0) + + named_params = list(model.named_parameters(remove_duplicate=False)) + param_to_fqn: dict[nn.Parameter, str] = {} + for fqn, param in named_params: + param_to_fqn.setdefault(param, fqn) + overrides: dict[nn.Parameter, torch.dtype] = {} + matched_by: dict[nn.Parameter, str] = {} + for pattern, dtype_name in patterns.items(): + try: + dtype = _DTYPES[dtype_name] + except KeyError as error: + raise ValueError( + f"Unsupported dtype {dtype_name!r} for pattern {pattern!r}" + ) from error + matches: list[nn.Parameter] = [] + matched_params: set[nn.Parameter] = set() + for fqn, param in named_params: + if ( + fnmatch.fnmatchcase(fqn, pattern) + and param not in matched_params + ): + matches.append(param) + matched_params.add(param) + if not matches: + raise ValueError( + f"FSDP parameter dtype pattern {pattern!r} did not match any parameter" + ) + for param in matches: + if (previous := matched_by.get(param)) and previous != pattern: + raise ValueError( + f"Parameter {param_to_fqn[param]!r} matches both {previous!r} and {pattern!r}" + ) + matched_by[param] = pattern + if dtype != default_dtype: + overrides[param] = dtype + + module_names = dict(model.named_modules()) + module_to_name = {module: name for name, module in module_names.items()} + module_maps: dict[nn.Module, dict[str, torch.dtype]] = {} + managed_params: set[nn.Parameter] = set() + for module in modules: + module_name = module_to_name[module] + local_map: dict[str, torch.dtype] = {} + for local_fqn, param in module.named_parameters(): + if param in managed_params: + raise ValueError( + "FSDP wrap modules overlap at parameter " + f"{param_to_fqn.get(param, local_fqn)!r}" + ) + managed_params.add(param) + dtype = overrides.get(param) + if dtype is not None: + local_map[local_fqn] = dtype + if local_map: + module_maps[module] = local_map + + root_map = { + fqn: dtype + for param, dtype in overrides.items() + if param not in managed_params + for fqn in [param_to_fqn[param]] + } + return CompiledParamDtypeMaps( + module_maps, + root_map, + len(overrides), + sum(param.numel() for param in overrides), + ) + + +__all__ = ["CompiledParamDtypeMaps", "compile_param_dtype_maps"] diff --git a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py new file mode 100644 index 00000000..936eae0d --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py @@ -0,0 +1,86 @@ +import copy +import os +from argparse import Namespace + +import torch +import torch.distributed as dist +from torch import nn + +from miles.backends.fsdp_utils.actor import apply_fsdp2 + + +class MixedParamDtypeBlock(nn.Module): + def __init__(self): + super().__init__() + self.low_precision = nn.Linear(8, 15) + self.full_precision = nn.Linear(8, 15) + self.seen_param_dtypes = None + + def forward(self, x): + self.seen_param_dtypes = ( + self.low_precision.weight.dtype, + self.full_precision.weight.dtype, + ) + low_precision_output = self.low_precision(x) + full_precision_output = self.full_precision( + x.to(self.full_precision.weight.dtype) + ) + return low_precision_output + full_precision_output.to( + low_precision_output.dtype + ) + + +class MixedParamDtypeModel(nn.Module): + def __init__(self): + super().__init__() + self.block = MixedParamDtypeBlock() + + def forward(self, x): + return self.block(x) + + +def main(): + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + + torch.manual_seed(42) + model = MixedParamDtypeModel().cuda() + ref_model = copy.deepcopy(model) + ref_model.block.low_precision.to(torch.bfloat16) + model = apply_fsdp2( + model, + args=Namespace( + diffusion_forward_dtype="bf16", + fsdp_reduce_dtype="fp32", + gradient_checkpointing=False, + ), + no_split_modules=["MixedParamDtypeBlock"], + param_dtype_patterns={"block.full_precision.*": "fp32"}, + ) + + torch.manual_seed(43) + inp = torch.randn(4, 8, device="cuda") + output = model(inp) + ref_output = ref_model(inp.to(torch.bfloat16)) + torch.testing.assert_close(output, ref_output) + assert model.block.seen_param_dtypes == (torch.bfloat16, torch.float32) + + output.sum().backward() + ref_output.sum().backward() + for param, ref_param in zip(model.parameters(), ref_model.parameters()): + assert param.grad is not None + assert ref_param.grad is not None + assert param.grad.dtype == torch.float32 + torch.testing.assert_close( + param.grad.full_tensor(), + ref_param.grad.to(torch.float32), + ) + + if dist.get_rank() == 0: + print("OK") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast-gpu/backends/fsdp_utils/test_param_dtype_map.py b/tests/fast-gpu/backends/fsdp_utils/test_param_dtype_map.py new file mode 100644 index 00000000..0b098e49 --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/test_param_dtype_map.py @@ -0,0 +1,39 @@ +"""FSDP2 per-parameter mixed precision.""" + +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=30, + suite="stage-b-3-gpu-h200", + labels=["fsdp"], +) + +import os +import subprocess +import sys +from pathlib import Path + + +_WORKER = Path(__file__).with_name("_param_dtype_map_worker.py") + + +def test_param_dtype_map(): + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc_per_node=2", + str(_WORKER), + ], + env=env, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "OK" in result.stdout diff --git a/tests/fast/backends/fsdp_utils/test_mixed_precision.py b/tests/fast/backends/fsdp_utils/test_mixed_precision.py new file mode 100644 index 00000000..591658e6 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_mixed_precision.py @@ -0,0 +1,94 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=["fsdp"]) + +import pytest +import torch +from torch import nn + +from miles.backends.fsdp_utils.mixed_precision import compile_param_dtype_maps + + +class Block(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.norm = nn.LayerNorm(4) + + +class Model(nn.Module): + def __init__(self): + super().__init__() + self.root_scale = nn.Parameter(torch.ones(1)) + self.blocks = nn.ModuleList([Block(), Block()]) + + +def test_compile_param_dtype_maps_scopes_fqns(): + model = Model() + compiled = compile_param_dtype_maps( + model, + list(model.blocks), + { + "blocks.*.norm.weight": "fp32", + "root_scale": "fp32", + }, + torch.bfloat16, + ) + + assert compiled.module_maps == { + model.blocks[0]: {"norm.weight": torch.float32}, + model.blocks[1]: {"norm.weight": torch.float32}, + } + assert compiled.root_map == {"root_scale": torch.float32} + + +def test_compile_param_dtype_maps_rejects_unmatched_pattern(): + model = Model() + with pytest.raises(ValueError, match="did not match any parameter"): + compile_param_dtype_maps( + model, + list(model.blocks), + {"blocks.*.missing": "fp32"}, + torch.bfloat16, + ) + + +def test_compile_param_dtype_maps_rejects_overlap(): + model = Model() + with pytest.raises(ValueError, match="matches both"): + compile_param_dtype_maps( + model, + list(model.blocks), + { + "blocks.*.norm.weight": "fp32", + "*.0.norm.weight": "fp32", + }, + torch.bfloat16, + ) + + +def test_compile_param_dtype_maps_omits_default_dtype(): + model = Model() + compiled = compile_param_dtype_maps( + model, + list(model.blocks), + {"blocks.*.norm.weight": "bf16"}, + torch.bfloat16, + ) + + assert compiled.module_maps == {} + assert compiled.root_map == {} + + +def test_compile_param_dtype_maps_canonicalizes_shared_parameter_alias(): + model = nn.Module() + model.shared = nn.Linear(4, 4, bias=False) + model.alias = model.shared + compiled = compile_param_dtype_maps( + model, + [], + {"alias.weight": "fp32"}, + torch.bfloat16, + ) + + assert compiled.root_map == {"shared.weight": torch.float32} From 7d0b57e2751202fca4149800d0c811fed7db9105 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 5 Aug 2026 22:07:59 -0700 Subject: [PATCH 2/7] test(fsdp): harden monkeypatch coverage Exercise patch guards, packing paths, policy validation, scoped integration, and frozen FP32 parameters while keeping the feature independent of image rebuilds. --- .github/workflows/docker-build.yml | 4 - .github/workflows/pr-test.yml | 2 +- docker/Dockerfile | 3 - .../fsdp_utils/_param_dtype_map_worker.py | 175 +++++++++++++++++- .../fsdp_utils/test_fsdp_param_dtype_patch.py | 61 ++++++ .../fsdp_utils/test_mixed_precision.py | 26 +++ 6 files changed, 257 insertions(+), 14 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index d3adbfa6..8d051ca0 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -26,10 +26,6 @@ on: paths: - 'docker/Dockerfile' - 'requirements.txt' - - 'miles/backends/fsdp_utils/actor.py' - - 'miles/backends/fsdp_utils/configs/**' - - 'miles/backends/fsdp_utils/fsdp_param_dtype_patch.py' - - 'miles/backends/fsdp_utils/mixed_precision.py' concurrency: group: docker-build diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index be2895ca..6578a042 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -68,7 +68,7 @@ jobs: # Compare the PR merge commit with its current base. CHANGED_FILES=$(git diff --name-only HEAD^1 HEAD) # Same paths as docker-build.yml's post-merge filter. - if printf '%s\n' "$CHANGED_FILES" | grep -qE '^(docker/Dockerfile|requirements\.txt|miles/backends/fsdp_utils/((actor|fsdp_param_dtype_patch|mixed_precision)\.py|configs/.*\.py))$'; then + if printf '%s\n' "$CHANGED_FILES" | grep -qE '^(docker/Dockerfile|requirements\.txt)$'; then echo "changed=true" >> "$GITHUB_OUTPUT" else echo "changed=false" >> "$GITHUB_OUTPUT" diff --git a/docker/Dockerfile b/docker/Dockerfile index bd89a798..6cb7b1fe 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,9 +1,6 @@ ARG SGLANG_IMAGE_TAG=v0.5.12-cu129 FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang -ARG TORCH_VERSION=2.11.0+cu129 -RUN python -c "import torch; assert torch.__version__ == '${TORCH_VERSION}', f'expected torch ${TORCH_VERSION}, got {torch.__version__}'" - # ======================================== Setup ============================================= WORKDIR /root/ diff --git a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py index 936eae0d..eae41332 100644 --- a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py +++ b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py @@ -5,10 +5,18 @@ import torch import torch.distributed as dist from torch import nn +from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard +from miles.backends.fsdp_utils import fsdp_param_dtype_patch from miles.backends.fsdp_utils.actor import apply_fsdp2 +PARAM_DTYPE_PATTERNS = { + "block.full_precision.*": "fp32", + "root_scale": "fp32", +} + + class MixedParamDtypeBlock(nn.Module): def __init__(self): super().__init__() @@ -33,17 +41,123 @@ def forward(self, x): class MixedParamDtypeModel(nn.Module): def __init__(self): super().__init__() + self.root_scale = nn.Parameter(torch.ones(1)) self.block = MixedParamDtypeBlock() + self.seen_root_dtype = None def forward(self, x): - return self.block(x) + self.seen_root_dtype = self.root_scale.dtype + output = self.block(x) + return output * self.root_scale.to(output.dtype) + + +def test_patch_installation(): + fsdp_param_dtype_patch.apply_param_dtype_map_patch() + patched_reduce = fsdp_param_dtype_patch._fsdp_collectives.foreach_reduce + assert patched_reduce is not fsdp_param_dtype_patch._ORIGINAL_FOREACH_REDUCE + assert ( + fsdp_param_dtype_patch._fsdp_param_group.foreach_reduce + is patched_reduce + ) + fsdp_param_dtype_patch.apply_param_dtype_map_patch() + assert fsdp_param_dtype_patch._fsdp_collectives.foreach_reduce is patched_reduce -def main(): - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - dist.init_process_group("nccl") +def test_reduce_scatter_copy_in(): + copy_in = ( + fsdp_param_dtype_patch._fsdp_collectives.foreach_reduce_scatter_copy_in + ) + bf16_grad = torch.arange(6, device="cuda").reshape(3, 2).to(torch.bfloat16) + fp32_grad = torch.arange(10, 16, device="cuda").reshape(2, 3).to(torch.float32) + mixed_output = torch.empty(14, device="cuda", dtype=torch.float32) + copy_in([bf16_grad, fp32_grad], mixed_output, world_size=2) + expected_mixed = torch.tensor( + [ + [0, 1, 2, 3, 10, 11, 12], + [4, 5, 0, 0, 13, 14, 15], + ], + device="cuda", + dtype=torch.float32, + ) + torch.testing.assert_close( + mixed_output.view(2, -1), + expected_mixed, + rtol=0, + atol=0, + ) + + uniform_output = torch.empty(8, device="cuda", dtype=torch.float32) + copy_in([bf16_grad], uniform_output, world_size=2) + expected_uniform = torch.tensor( + [[0, 1, 2, 3], [4, 5, 0, 0]], + device="cuda", + dtype=torch.float32, + ) + torch.testing.assert_close( + uniform_output.view(2, -1), + expected_uniform, + rtol=0, + atol=0, + ) + + +def _assert_policy_error(policy, message): + model = MixedParamDtypeBlock().cuda() + try: + fully_shard(model, mp_policy=policy) + model(torch.randn(2, 8, device="cuda")) + except ValueError as error: + assert message in str(error) + else: + raise AssertionError(f"Expected ValueError containing {message!r}") + + +def test_policy_validation(): + _assert_policy_error( + fsdp_param_dtype_patch.ParamDtypeMixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + param_dtype_map={"missing.weight": torch.float32}, + ), + "param_dtype_map contains FQNs", + ) + _assert_policy_error( + fsdp_param_dtype_patch.ParamDtypeMixedPrecisionPolicy( + param_dtype=torch.bfloat16, + param_dtype_map={ + "full_precision.weight": torch.float32, + "full_precision.bias": torch.float32, + }, + ), + "Mixed parameter dtypes require an explicit reduce_dtype", + ) + + +def test_standard_policy_delegation(): + torch.manual_seed(40) + model = nn.Linear(8, 8, bias=False).cuda() + ref_model = copy.deepcopy(model).to(torch.bfloat16) + fully_shard( + model, + mp_policy=MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + ), + ) + inp = torch.randn(3, 8, device="cuda") + output = model(inp) + ref_output = ref_model(inp.to(torch.bfloat16)) + torch.testing.assert_close(output, ref_output) + output.sum().backward() + ref_output.sum().backward() + torch.testing.assert_close( + model.weight.grad.full_tensor(), + ref_model.weight.grad.to(torch.float32), + ) + + +def test_apply_fsdp2_integration(): torch.manual_seed(42) model = MixedParamDtypeModel().cuda() ref_model = copy.deepcopy(model) @@ -56,7 +170,7 @@ def main(): gradient_checkpointing=False, ), no_split_modules=["MixedParamDtypeBlock"], - param_dtype_patterns={"block.full_precision.*": "fp32"}, + param_dtype_patterns=PARAM_DTYPE_PATTERNS, ) torch.manual_seed(43) @@ -65,6 +179,7 @@ def main(): ref_output = ref_model(inp.to(torch.bfloat16)) torch.testing.assert_close(output, ref_output) assert model.block.seen_param_dtypes == (torch.bfloat16, torch.float32) + assert model.seen_root_dtype == torch.float32 output.sum().backward() ref_output.sum().backward() @@ -77,6 +192,54 @@ def main(): ref_param.grad.to(torch.float32), ) + +def test_frozen_fp32_parameters(): + torch.manual_seed(44) + model = MixedParamDtypeModel().cuda() + model.block.full_precision.requires_grad_(False) + ref_model = copy.deepcopy(model) + ref_model.block.low_precision.to(torch.bfloat16) + model = apply_fsdp2( + model, + args=Namespace( + diffusion_forward_dtype="bf16", + fsdp_reduce_dtype="fp32", + gradient_checkpointing=False, + ), + no_split_modules=["MixedParamDtypeBlock"], + param_dtype_patterns=PARAM_DTYPE_PATTERNS, + ) + + inp = torch.randn(4, 8, device="cuda") + output = model(inp) + ref_output = ref_model(inp.to(torch.bfloat16)) + torch.testing.assert_close(output, ref_output) + output.sum().backward() + ref_output.sum().backward() + assert model.block.full_precision.weight.grad is None + assert ref_model.block.full_precision.weight.grad is None + for param, ref_param in zip( + model.block.low_precision.parameters(), + ref_model.block.low_precision.parameters(), + ): + torch.testing.assert_close( + param.grad.full_tensor(), + ref_param.grad.to(torch.float32), + ) + + +def main(): + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + + test_patch_installation() + test_reduce_scatter_copy_in() + test_policy_validation() + test_standard_policy_delegation() + test_apply_fsdp2_integration() + test_frozen_fp32_parameters() + if dist.get_rank() == 0: print("OK") dist.destroy_process_group() diff --git a/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py b/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py new file mode 100644 index 00000000..753acfe0 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py @@ -0,0 +1,61 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=["fsdp"]) + +import pytest +import torch + +from miles.backends.fsdp_utils import fsdp_param_dtype_patch + + +def test_patch_rejects_unpinned_torch(monkeypatch): + monkeypatch.delattr( + fsdp_param_dtype_patch._fsdp_collectives, + fsdp_param_dtype_patch._PATCH_SENTINEL, + raising=False, + ) + monkeypatch.setattr( + fsdp_param_dtype_patch, + "_EXPECTED_TORCH_VERSION", + "unsupported-version", + ) + + with pytest.raises(RuntimeError, match="requires torch==unsupported-version"): + fsdp_param_dtype_patch.apply_param_dtype_map_patch() + + +def test_patch_rejects_source_drift(monkeypatch): + monkeypatch.delattr( + fsdp_param_dtype_patch._fsdp_collectives, + fsdp_param_dtype_patch._PATCH_SENTINEL, + raising=False, + ) + monkeypatch.setattr( + fsdp_param_dtype_patch, + "_EXPECTED_TORCH_VERSION", + torch.__version__, + ) + monkeypatch.setitem( + fsdp_param_dtype_patch._SOURCE_HASHES, + "FSDPParamGroup.__init__", + "0" * 64, + ) + + with pytest.raises( + RuntimeError, + match="source hash for FSDPParamGroup.__init__ changed", + ): + fsdp_param_dtype_patch.apply_param_dtype_map_patch() + + +def test_param_dtype_policy_keeps_exact_sparse_map(): + param_dtype_map = {"norm.weight": torch.float32} + policy = fsdp_param_dtype_patch.ParamDtypeMixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + param_dtype_map=param_dtype_map, + ) + + assert policy.param_dtype == torch.bfloat16 + assert policy.reduce_dtype == torch.float32 + assert policy.param_dtype_map == param_dtype_map diff --git a/tests/fast/backends/fsdp_utils/test_mixed_precision.py b/tests/fast/backends/fsdp_utils/test_mixed_precision.py index 591658e6..b5441a4f 100644 --- a/tests/fast/backends/fsdp_utils/test_mixed_precision.py +++ b/tests/fast/backends/fsdp_utils/test_mixed_precision.py @@ -40,6 +40,8 @@ def test_compile_param_dtype_maps_scopes_fqns(): model.blocks[1]: {"norm.weight": torch.float32}, } assert compiled.root_map == {"root_scale": torch.float32} + assert compiled.override_count == 3 + assert compiled.override_numel == 9 def test_compile_param_dtype_maps_rejects_unmatched_pattern(): @@ -67,6 +69,28 @@ def test_compile_param_dtype_maps_rejects_overlap(): ) +def test_compile_param_dtype_maps_rejects_unsupported_dtype(): + model = Model() + with pytest.raises(ValueError, match="Unsupported dtype 'float8'"): + compile_param_dtype_maps( + model, + list(model.blocks), + {"root_scale": "float8"}, + torch.bfloat16, + ) + + +def test_compile_param_dtype_maps_rejects_overlapping_wrap_modules(): + model = Model() + with pytest.raises(ValueError, match="FSDP wrap modules overlap"): + compile_param_dtype_maps( + model, + [model.blocks, model.blocks[0]], + {"blocks.0.norm.weight": "fp32"}, + torch.bfloat16, + ) + + def test_compile_param_dtype_maps_omits_default_dtype(): model = Model() compiled = compile_param_dtype_maps( @@ -78,6 +102,8 @@ def test_compile_param_dtype_maps_omits_default_dtype(): assert compiled.module_maps == {} assert compiled.root_map == {} + assert compiled.override_count == 0 + assert compiled.override_numel == 0 def test_compile_param_dtype_maps_canonicalizes_shared_parameter_alias(): From 3b223fb8b3c4166934df01435773916245505b65 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 5 Aug 2026 22:12:23 -0700 Subject: [PATCH 3/7] fix(fsdp): allow PyTorch 2.11 CUDA variants Validate the base PyTorch release and source hashes without tying the monkeypatch to one CUDA wheel suffix. --- .../backends/fsdp_utils/fsdp_param_dtype_patch.py | 5 +++-- .../fsdp_utils/test_fsdp_param_dtype_patch.py | 14 +++----------- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py index 24223f51..9c268fd8 100644 --- a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py +++ b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py @@ -28,7 +28,7 @@ from torch.distributed.tensor import DTensor -_EXPECTED_TORCH_VERSION = "2.11.0+cu129" +_EXPECTED_TORCH_VERSION = "2.11.0" _PATCH_SENTINEL = "_miles_param_dtype_map_patch_applied" _SOURCE_HASHES = { "FSDPParam.init_dtype_attrs": "2cc968770804055cdde959db7cfa47b92a1137badcdaf2e448555741c2fd282c", @@ -540,7 +540,8 @@ def _patched_foreach_reduce_scatter_copy_in( def apply_param_dtype_map_patch() -> None: if getattr(_fsdp_collectives, _PATCH_SENTINEL, False): return - if torch.__version__ != _EXPECTED_TORCH_VERSION: + torch_version = torch.__version__.partition("+")[0] + if torch_version != _EXPECTED_TORCH_VERSION: raise RuntimeError( "The Miles FSDP param-dtype patch requires " f"torch=={_EXPECTED_TORCH_VERSION}, got {torch.__version__}" diff --git a/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py b/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py index 753acfe0..14eca7c7 100644 --- a/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py +++ b/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py @@ -14,13 +14,9 @@ def test_patch_rejects_unpinned_torch(monkeypatch): fsdp_param_dtype_patch._PATCH_SENTINEL, raising=False, ) - monkeypatch.setattr( - fsdp_param_dtype_patch, - "_EXPECTED_TORCH_VERSION", - "unsupported-version", - ) + monkeypatch.setattr(torch, "__version__", "2.12.0+cu130") - with pytest.raises(RuntimeError, match="requires torch==unsupported-version"): + with pytest.raises(RuntimeError, match="requires torch==2.11.0"): fsdp_param_dtype_patch.apply_param_dtype_map_patch() @@ -30,11 +26,7 @@ def test_patch_rejects_source_drift(monkeypatch): fsdp_param_dtype_patch._PATCH_SENTINEL, raising=False, ) - monkeypatch.setattr( - fsdp_param_dtype_patch, - "_EXPECTED_TORCH_VERSION", - torch.__version__, - ) + monkeypatch.setattr(torch, "__version__", "2.11.0+cu130") monkeypatch.setitem( fsdp_param_dtype_patch._SOURCE_HASHES, "FSDPParamGroup.__init__", From c13cc4618ffcee649156acbd6e983d10a72b64ec Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 5 Aug 2026 22:32:06 -0700 Subject: [PATCH 4/7] refactor(fsdp): move precision rules into model plans Make FSDP wrapping and precision policy model-owned so Diffusers and native backends share one explicit parallel-plan interface. --- miles/backends/fsdp_utils/actor.py | 5 +-- .../configs/train_pipeline_config.py | 6 ++-- miles/backends/fsdp_utils/configs/wan2_2.py | 5 --- miles/backends/fsdp_utils/model_backend.py | 32 +++++++++++++++---- miles/backends/fsdp_utils/models/__init__.py | 7 ++-- .../fsdp_utils/models/diffusers/__init__.py | 18 +++++++++++ .../models/diffusers/qwen_image/__init__.py | 1 + .../diffusers/qwen_image/parallel_plan.py | 4 +++ .../models/diffusers/sd3/__init__.py | 1 + .../models/diffusers/sd3/parallel_plan.py | 4 +++ .../models/diffusers/wan2_2/__init__.py | 1 + .../models/diffusers/wan2_2/parallel_plan.py | 10 ++++++ .../fsdp_utils/models/ltx/parallel_plan.py | 5 ++- miles/backends/fsdp_utils/models/package.py | 2 +- .../fsdp_utils/models/parallel_plan.py | 13 ++++++++ .../fsdp_utils/models/ltx/test_ltx.py | 8 +++++ .../backends/fsdp_utils/test_model_backend.py | 26 ++++++++++++++- 17 files changed, 126 insertions(+), 22 deletions(-) create mode 100644 miles/backends/fsdp_utils/models/diffusers/__init__.py create mode 100644 miles/backends/fsdp_utils/models/diffusers/qwen_image/__init__.py create mode 100644 miles/backends/fsdp_utils/models/diffusers/qwen_image/parallel_plan.py create mode 100644 miles/backends/fsdp_utils/models/diffusers/sd3/__init__.py create mode 100644 miles/backends/fsdp_utils/models/diffusers/sd3/parallel_plan.py create mode 100644 miles/backends/fsdp_utils/models/diffusers/wan2_2/__init__.py create mode 100644 miles/backends/fsdp_utils/models/diffusers/wan2_2/parallel_plan.py create mode 100644 miles/backends/fsdp_utils/models/parallel_plan.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 2258b9ab..5ab2beb7 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -138,13 +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), - param_dtype_patterns=self.train_pipeline_config.fsdp_param_dtype_patterns, + 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, diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 8ffbfe69..45852f46 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/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,8 +76,8 @@ 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"] - fsdp_param_dtype_patterns: dict[str, str] = {} needs_timestep_scaling: bool = True optimizer_state_allowed_missing: list[str] = [] # Case-insensitive substrings matched against the checkpoint name (--diffusion-model). diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 7ebb11fd..7e227cd2 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -11,11 +11,6 @@ @register_train_pipeline_config("wan2_2") class Wan2_2TrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("wan2.2", "wan-2.2") - fsdp_param_dtype_patterns = { - "*scale_shift_table": "fp32", - "*time_embedder*": "fp32", - "*.norm2.*": "fp32", - } # High-noise expert ("transformer") handles t >= boundary, low-noise expert # ("transformer_2") the rest. boundary_ratio = 0.875 diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 72a25711..fe1fc8df 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/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,9 +64,15 @@ 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 + def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]: + no_split_modules = self.fsdp_parallel_plan(model).no_split_modules + if no_split_modules is None: + raise ValueError("FSDP parallel plan declares no no-split modules") + return list(no_split_modules) + @abc.abstractmethod def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None: raise NotImplementedError @@ -123,8 +131,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 +161,21 @@ 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) + raise ValueError( + f"{model.__class__.__name__} declares no _no_split_modules for FSDP wrapping" + ) + 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) diff --git a/miles/backends/fsdp_utils/models/__init__.py b/miles/backends/fsdp_utils/models/__init__.py index b6c3fe8d..3cf6731f 100644 --- a/miles/backends/fsdp_utils/models/__init__.py +++ b/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/.py``. +- **Diffusers checkpoint** (has ``model_index.json``): + ``DiffusersModelBackend`` loads it; add + ``models/diffusers//parallel_plan.py`` for its FSDP precision plan. - **Native modeling** (official repo code, non-diffusers checkpoint): add a package ``models//`` 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`` diff --git a/miles/backends/fsdp_utils/models/diffusers/__init__.py b/miles/backends/fsdp_utils/models/diffusers/__init__.py new file mode 100644 index 00000000..ff3cca10 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/__init__.py @@ -0,0 +1,18 @@ +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"] diff --git a/miles/backends/fsdp_utils/models/diffusers/qwen_image/__init__.py b/miles/backends/fsdp_utils/models/diffusers/qwen_image/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/qwen_image/__init__.py @@ -0,0 +1 @@ + diff --git a/miles/backends/fsdp_utils/models/diffusers/qwen_image/parallel_plan.py b/miles/backends/fsdp_utils/models/diffusers/qwen_image/parallel_plan.py new file mode 100644 index 00000000..daec57f0 --- /dev/null +++ b/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() diff --git a/miles/backends/fsdp_utils/models/diffusers/sd3/__init__.py b/miles/backends/fsdp_utils/models/diffusers/sd3/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/sd3/__init__.py @@ -0,0 +1 @@ + diff --git a/miles/backends/fsdp_utils/models/diffusers/sd3/parallel_plan.py b/miles/backends/fsdp_utils/models/diffusers/sd3/parallel_plan.py new file mode 100644 index 00000000..daec57f0 --- /dev/null +++ b/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() diff --git a/miles/backends/fsdp_utils/models/diffusers/wan2_2/__init__.py b/miles/backends/fsdp_utils/models/diffusers/wan2_2/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/wan2_2/__init__.py @@ -0,0 +1 @@ + diff --git a/miles/backends/fsdp_utils/models/diffusers/wan2_2/parallel_plan.py b/miles/backends/fsdp_utils/models/diffusers/wan2_2/parallel_plan.py new file mode 100644 index 00000000..ceb6a866 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/wan2_2/parallel_plan.py @@ -0,0 +1,10 @@ +from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan + + +FSDP_PARALLEL_PLAN = FSDPParallelPlan( + param_dtype_patterns={ + "*scale_shift_table": "fp32", + "*time_embedder*": "fp32", + "*.norm2.*": "fp32", + }, +) diff --git a/miles/backends/fsdp_utils/models/ltx/parallel_plan.py b/miles/backends/fsdp_utils/models/ltx/parallel_plan.py index 63a4517e..d1629f92 100644 --- a/miles/backends/fsdp_utils/models/ltx/parallel_plan.py +++ b/miles/backends/fsdp_utils/models/ltx/parallel_plan.py @@ -8,9 +8,12 @@ import torch +from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan from miles.backends.fsdp_utils.sequence_parallel.plan import SequenceParallelPlan -FSDP_NO_SPLIT_MODULES = ["BasicAVTransformerBlock"] +FSDP_PARALLEL_PLAN = FSDPParallelPlan( + no_split_modules=("BasicAVTransformerBlock",), +) def sequence_parallel_plan(model: torch.nn.Module) -> SequenceParallelPlan: diff --git a/miles/backends/fsdp_utils/models/package.py b/miles/backends/fsdp_utils/models/package.py index 12aeaa78..fa123ecd 100644 --- a/miles/backends/fsdp_utils/models/package.py +++ b/miles/backends/fsdp_utils/models/package.py @@ -6,7 +6,7 @@ materialize_weights=...)``; distributed rank selection stays outside the package - ``modeling`` — ``load_scheduler``, ``enable_gradient_checkpointing``, optional ``flash_attention_entrypoints`` / ``required_flash_kernel_label`` - - ``parallel_plan`` — ``FSDP_NO_SPLIT_MODULES``, ``sequence_parallel_plan`` + - ``parallel_plan`` — ``FSDP_PARALLEL_PLAN``, ``sequence_parallel_plan`` (and optional ``install_sequence_parallel_attention``) - ``attention`` — ``set_attention_backend`` diff --git a/miles/backends/fsdp_utils/models/parallel_plan.py b/miles/backends/fsdp_utils/models/parallel_plan.py new file mode 100644 index 00000000..2aaaf399 --- /dev/null +++ b/miles/backends/fsdp_utils/models/parallel_plan.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class FSDPParallelPlan: + no_split_modules: tuple[str, ...] | None = None + param_dtype_patterns: Mapping[str, str] = field(default_factory=dict) + + +__all__ = ["FSDPParallelPlan"] diff --git a/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py b/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py index ac42e8c5..3a11948b 100644 --- a/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py +++ b/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py @@ -14,6 +14,14 @@ class _LTXConfig: class TestLTXPackage: + def test_fsdp_parallel_plan(self): + plan = MilesModelBackend(_LTXConfig()).fsdp_parallel_plan( + torch.nn.Linear(2, 2) + ) + + assert plan.no_split_modules == ("BasicAVTransformerBlock",) + assert plan.param_dtype_patterns == {} + def test_component_loading_delegates_to_loading_module(self, monkeypatch): backend = MilesModelBackend(_LTXConfig()) sentinel = object() diff --git a/tests/fast/backends/fsdp_utils/test_model_backend.py b/tests/fast/backends/fsdp_utils/test_model_backend.py index 9c84f817..8bb9dc4d 100644 --- a/tests/fast/backends/fsdp_utils/test_model_backend.py +++ b/tests/fast/backends/fsdp_utils/test_model_backend.py @@ -4,7 +4,13 @@ import torch -from miles.backends.fsdp_utils.model_backend import BaseModelBackend, DiffusersModelBackend, MilesModelBackend +from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig +from miles.backends.fsdp_utils.model_backend import ( + BaseModelBackend, + DiffusersModelBackend, + MilesModelBackend, +) +from miles.backends.fsdp_utils.models.diffusers import load_fsdp_parallel_plan class _RecordingModel(torch.nn.Module): @@ -33,7 +39,25 @@ def test_diffusers_implements_model_lifecycle_hooks(self): backend.enable_gradient_checkpointing(model) backend.set_attention_backend(model, "flash") + plan = backend.fsdp_parallel_plan(model) assert model.gradient_checkpointing_enabled assert model.selected == "flash" + assert plan.no_split_modules == ("TransformerBlock",) + assert plan.param_dtype_patterns == {} assert backend.fsdp_no_split_modules(model) == ["TransformerBlock"] + + def test_diffusers_model_family_owns_fsdp_precision_plan(self): + backend = DiffusersModelBackend(Wan2_2TrainPipelineConfig()) + plan = backend.fsdp_parallel_plan(_RecordingModel()) + + assert plan.param_dtype_patterns == { + "*scale_shift_table": "fp32", + "*time_embedder*": "fp32", + "*.norm2.*": "fp32", + } + + def test_all_diffusers_model_plans_load(self): + assert load_fsdp_parallel_plan("sd3").param_dtype_patterns == {} + assert load_fsdp_parallel_plan("qwen_image").param_dtype_patterns == {} + assert load_fsdp_parallel_plan("wan2_2").param_dtype_patterns From 6eb83a404f7732c441fa6d5d04f414195ba5df78 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 5 Aug 2026 22:46:15 -0700 Subject: [PATCH 5/7] fix(fsdp): align and annotate monkeypatch replacements Show each upstream block directly above its replacement and normalize empty dtype maps to the unpatched policy path. --- .../fsdp_utils/fsdp_param_dtype_patch.py | 103 +++++++++++++++--- .../fsdp_utils/_param_dtype_map_worker.py | 38 ++++--- 2 files changed, 111 insertions(+), 30 deletions(-) diff --git a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py index 9c268fd8..66f19bde 100644 --- a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py +++ b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py @@ -94,7 +94,7 @@ def _patched_param_group_init( if isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) else None ) - if param_dtype_map is None: + if not param_dtype_map: _ORIGINAL_PARAM_GROUP_INIT( self, params, @@ -108,7 +108,11 @@ def _patched_param_group_init( ) return - # MILES_PATCH_BEGIN: fsdp-param-dtype-map + # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map + # Upstream forwards one group-wide mixed-precision policy to every + # FSDPParam without resolving per-parameter dtype overrides. + # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map managed_params = set(params) fqn_to_param: dict[str, nn.Parameter] = {} for module in modules: @@ -137,7 +141,7 @@ def _patched_param_group_init( } if len(effective_dtypes) > 1 and mp_policy.reduce_dtype is None: raise ValueError("Mixed parameter dtypes require an explicit reduce_dtype") - # MILES_PATCH_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map _ORIGINAL_PARAM_GROUP_INIT( self, @@ -151,7 +155,10 @@ def _patched_param_group_init( offload_policy, ) - # MILES_PATCH_BEGIN: fsdp-param-dtype-map + # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map + # self.fsdp_params keeps the same group-wide mp_policy after construction. + # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map for fsdp_param, param in zip(self.fsdp_params, params, strict=True): override = param_overrides.get(param) fsdp_param._param_dtype_override = override @@ -163,18 +170,31 @@ def _patched_param_group_init( param_dtype=effective_param_dtype, param_dtype_map=None, ) - # MILES_PATCH_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map def _patched_init_dtype_attrs( self: FSDPParam, mp_policy: TorchMixedPrecisionPolicy, ) -> None: - if not isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy): + if ( + not isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) + or not mp_policy.param_dtype_map + ): _ORIGINAL_INIT_DTYPE_ATTRS(self, mp_policy) return - # MILES_PATCH_BEGIN: fsdp-param-dtype-map + # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map + # param_dtype, reduce_dtype = (mp_policy.param_dtype, mp_policy.reduce_dtype) + # self.orig_dtype = self.sharded_param.dtype + # if reduce_dtype == param_dtype: + # reduce_dtype = None + # if param_dtype == self.orig_dtype: + # param_dtype = None + # self.param_dtype = param_dtype + # self.reduce_dtype = reduce_dtype + # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map param_dtype = ( self._param_dtype_override if self._param_dtype_override is not None @@ -188,7 +208,7 @@ def _patched_init_dtype_attrs( param_dtype = None self.param_dtype = param_dtype self.reduce_dtype = reduce_dtype - # MILES_PATCH_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map # Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. @@ -209,7 +229,42 @@ def use_foreach_copy(fsdp_param: FSDPParam) -> bool: ) param_all_gather_inputs: list[list[torch.Tensor]] = [[] for _ in fsdp_params] - # MILES_PATCH_BEGIN: fsdp-param-dtype-map + # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map + # foreach_copy_indices: list[int] = [] + # foreach_copy_inputs: list[torch.Tensor] = [] + # foreach_copy_input_numels: list[int] = [] + # + # for i, fsdp_param in enumerate(fsdp_params): + # if use_foreach_copy(fsdp_param): + # foreach_copy_indices.append(i) + # all_gather_input = ( + # fsdp_param._sharded_param_data + # if fsdp_param.sharded_state == ShardedState.SHARDED + # else cast( + # torch.Tensor, fsdp_param._sharded_post_forward_param_data + # ) + # ) + # foreach_copy_inputs.append(all_gather_input) + # foreach_copy_input_numels.append(all_gather_input.numel()) + # else: + # param_all_gather_inputs[i] = fsdp_param.all_gather_inputs + # + # if foreach_copy_inputs: + # fsdp_param_0 = fsdp_params[foreach_copy_indices[0]] + # param_dtype, device = fsdp_param_0.param_dtype, fsdp_param_0.device + # flat_foreach_copy_input = torch.empty( + # (sum(foreach_copy_input_numels),), + # device=device, + # dtype=param_dtype, + # ) + # splits = torch.split( + # flat_foreach_copy_input, foreach_copy_input_numels + # ) + # torch._foreach_copy_(splits, foreach_copy_inputs) + # for i, split in zip(foreach_copy_indices, splits): + # param_all_gather_inputs[i] = [split] + # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map foreach_copy_infos: dict[ torch.dtype, tuple[list[int], list[torch.Tensor], list[int]] ] = {} @@ -243,7 +298,7 @@ def use_foreach_copy(fsdp_param: FSDPParam) -> bool: torch._foreach_copy_(splits, inputs) for i, split in zip(indices, splits): param_all_gather_inputs[i] = [split] - # MILES_PATCH_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map return param_all_gather_inputs @@ -278,7 +333,19 @@ def _patched_foreach_reduce( autograd, so clearing the list frees the gradients. """ - # MILES_PATCH_BEGIN: fsdp-param-dtype-map + # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map + # grad_dtypes = {grad.dtype for grad in unsharded_grads} + # if len(grad_dtypes) != 1: + # # Check this at runtime since it could be a real runtime error if e.g. + # # fp8 weights do not produce the correct higher precision gradients + # _raise_assert_with_print( + # "FSDP reduce-scatter expects uniform gradient dtype but got " + # f"{grad_dtypes}" + # ) + # grad_dtype = unsharded_grads[0].dtype + # reduce_dtype = reduce_dtype or grad_dtype + # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map grad_dtypes = {grad.dtype for grad in unsharded_grads} if reduce_dtype is None and len(grad_dtypes) != 1: _raise_assert_with_print( @@ -286,7 +353,7 @@ def _patched_foreach_reduce( f"gradient dtypes but got {grad_dtypes}" ) reduce_dtype = reduce_dtype or unsharded_grads[0].dtype - # MILES_PATCH_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map (predivide_factor, postdivide_factor, reduce_scatter_op, all_reduce_op) = ( _get_gradient_divide_factors( reduce_scatter_group, @@ -487,7 +554,15 @@ def _patched_foreach_reduce_scatter_copy_in( world_size: int, ) -> None: reduce_scatter_input = reduce_scatter_input.view(world_size, -1) - # MILES_PATCH_BEGIN: fsdp-param-dtype-map + # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map + # torch.ops.fsdp.chunk_cat( + # unsharded_grads, + # dim=0, + # num_chunks=world_size, + # out=reduce_scatter_input, + # ) + # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map if len({grad.dtype for grad in unsharded_grads}) == 1: torch.ops.fsdp.chunk_cat( unsharded_grads, @@ -534,7 +609,7 @@ def _patched_foreach_reduce_scatter_copy_in( torch._foreach_copy_(destinations, sources) if padding_views: torch._foreach_zero_(padding_views) - # MILES_PATCH_END: fsdp-param-dtype-map + # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map def apply_param_dtype_map_patch() -> None: diff --git a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py index eae41332..3378804e 100644 --- a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py +++ b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py @@ -135,26 +135,32 @@ def test_policy_validation(): def test_standard_policy_delegation(): - torch.manual_seed(40) - model = nn.Linear(8, 8, bias=False).cuda() - ref_model = copy.deepcopy(model).to(torch.bfloat16) - fully_shard( - model, - mp_policy=MixedPrecisionPolicy( + policies = ( + MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32, ), + fsdp_param_dtype_patch.ParamDtypeMixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + param_dtype_map={}, + ), ) - inp = torch.randn(3, 8, device="cuda") - output = model(inp) - ref_output = ref_model(inp.to(torch.bfloat16)) - torch.testing.assert_close(output, ref_output) - output.sum().backward() - ref_output.sum().backward() - torch.testing.assert_close( - model.weight.grad.full_tensor(), - ref_model.weight.grad.to(torch.float32), - ) + for policy in policies: + torch.manual_seed(40) + model = nn.Linear(8, 8, bias=False).cuda() + ref_model = copy.deepcopy(model).to(torch.bfloat16) + fully_shard(model, mp_policy=policy) + inp = torch.randn(3, 8, device="cuda") + output = model(inp) + ref_output = ref_model(inp.to(torch.bfloat16)) + torch.testing.assert_close(output, ref_output) + output.sum().backward() + ref_output.sum().backward() + torch.testing.assert_close( + model.weight.grad.full_tensor(), + ref_model.weight.grad.to(torch.float32), + ) def test_apply_fsdp2_integration(): From 541584442549bfb0bf3496336affeb3075c7e648 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 5 Aug 2026 22:52:31 -0700 Subject: [PATCH 6/7] docs(fsdp): show exact upstream constructor block Replace the prose summary with the original FSDPParam construction code so the monkeypatch remains directly reviewable. --- .../fsdp_utils/fsdp_param_dtype_patch.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py index 66f19bde..0c25f224 100644 --- a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py +++ b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py @@ -109,8 +109,22 @@ def _patched_param_group_init( return # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map - # Upstream forwards one group-wide mixed-precision policy to every - # FSDPParam without resolving per-parameter dtype overrides. + # self.modules = modules # permit ref cycle because 1:1 lifetime + # param_module_infos = _get_param_module_infos(params, modules) + # + # self.fsdp_params = [ + # FSDPParam( + # param, + # module_info, + # mesh_info, + # post_forward_mesh_info, + # device, + # shard_placement_fn, + # mp_policy, + # offload_policy, + # ) + # for param, module_info in zip(params, param_module_infos) + # ] # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map managed_params = set(params) @@ -141,7 +155,6 @@ def _patched_param_group_init( } if len(effective_dtypes) > 1 and mp_policy.reduce_dtype is None: raise ValueError("Mixed parameter dtypes require an explicit reduce_dtype") - # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map _ORIGINAL_PARAM_GROUP_INIT( self, @@ -155,10 +168,6 @@ def _patched_param_group_init( offload_policy, ) - # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map - # self.fsdp_params keeps the same group-wide mp_policy after construction. - # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map - # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map for fsdp_param, param in zip(self.fsdp_params, params, strict=True): override = param_overrides.get(param) fsdp_param._param_dtype_override = override From 51e82b3f23c7d894b475e2b477a43610e88cd1d7 Mon Sep 17 00:00:00 2001 From: rockdu Date: Thu, 6 Aug 2026 17:25:17 -0700 Subject: [PATCH 7/7] refactor(fsdp): synchronize the vendored dtype patch Keep the monkeypatch structurally aligned with the proposed PyTorch change while preserving the zero-temporary mixed-gradient packing path. This PR was authored with Claude. Co-authored-by: Cursor --- .../fsdp_utils/fsdp_param_dtype_patch.py | 296 ++++++++++++------ 1 file changed, 194 insertions(+), 102 deletions(-) diff --git a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py index 0c25f224..002fc17e 100644 --- a/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py +++ b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py @@ -13,8 +13,10 @@ from torch import nn from torch.distributed.fsdp import MixedPrecisionPolicy as TorchMixedPrecisionPolicy from torch.distributed.fsdp._fully_shard import _fsdp_collectives, _fsdp_param_group -from torch.distributed.fsdp._fully_shard._fsdp_api import ReduceScatter +from torch.distributed.fsdp._fully_shard._fsdp_api import OffloadPolicy, ReduceScatter from torch.distributed.fsdp._fully_shard._fsdp_collectives import ( + AllGather, + AllGatherResult, _div_if_needed, _get_device_handle, _get_dim0_padded_size, @@ -23,14 +25,28 @@ _to_dtype_if_needed, compiled_autograd_enabled, ) +from torch.distributed.fsdp._fully_shard._fsdp_common import ( + DataParallelMeshInfo, + FSDPMeshInfo, +) from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState -from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup -from torch.distributed.tensor import DTensor +from torch.distributed.fsdp._fully_shard._fsdp_param_group import ( + _get_param_module_infos, + _ModuleToHandleDict, + AllReduceState, + DefaultAllGather, + DefaultReduceScatter, + FSDPCommContext, + FSDPParamGroup, + TrainingState, +) +from torch.distributed.tensor import DTensor, Shard _EXPECTED_TORCH_VERSION = "2.11.0" _PATCH_SENTINEL = "_miles_param_dtype_map_patch_applied" _SOURCE_HASHES = { + "FSDPParam.__init__": "5973449ece76930fa71e8d8d2aa43b481660bcd79b1b40fa9556b686664a9466", "FSDPParam.init_dtype_attrs": "2cc968770804055cdde959db7cfa47b92a1137badcdaf2e448555741c2fd282c", "FSDPParamGroup.__init__": "323868f31033bb5696eaf30f278838901c0d6b7e0da031c517177707b716409b", "_get_param_all_gather_inputs": "a70fbae57b8aa3dc669d01a409d30c74428546bc4d69635a90cffe9380a5785f", @@ -38,6 +54,7 @@ "foreach_reduce_scatter_copy_in": "559a065467abbaa578f348bd6a8c6478cfc8ee516e602227b795bc3f6727eb22", } +_ORIGINAL_PARAM_INIT = FSDPParam.__init__ _ORIGINAL_PARAM_GROUP_INIT = FSDPParamGroup.__init__ _ORIGINAL_INIT_DTYPE_ATTRS = FSDPParam.init_dtype_attrs _ORIGINAL_GET_PARAM_ALL_GATHER_INPUTS = _fsdp_collectives._get_param_all_gather_inputs @@ -78,40 +95,22 @@ def _bind_to_collectives(fn, name: str, *, no_grad: bool): return torch.no_grad()(bound) if no_grad else bound +# Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. def _patched_param_group_init( self, params: list[nn.Parameter], modules: tuple[nn.Module, ...], - mesh_info, - post_forward_mesh_info, + mesh_info: DataParallelMeshInfo, + post_forward_mesh_info: FSDPMeshInfo | None, device: torch.device, - shard_placement_fn, + shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None, mp_policy: TorchMixedPrecisionPolicy, - offload_policy, + offload_policy: OffloadPolicy, ) -> None: - param_dtype_map = ( - mp_policy.param_dtype_map - if isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) - else None - ) - if not param_dtype_map: - _ORIGINAL_PARAM_GROUP_INIT( - self, - params, - modules, - mesh_info, - post_forward_mesh_info, - device, - shard_placement_fn, - mp_policy, - offload_policy, - ) - return + self.modules = modules # permit ref cycle because 1:1 lifetime + param_module_infos = _get_param_module_infos(params, modules) # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map - # self.modules = modules # permit ref cycle because 1:1 lifetime - # param_module_infos = _get_param_module_infos(params, modules) - # # self.fsdp_params = [ # FSDPParam( # param, @@ -127,83 +126,173 @@ def _patched_param_group_init( # ] # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map - managed_params = set(params) - fqn_to_param: dict[str, nn.Parameter] = {} - for module in modules: - for fqn, param in module.named_parameters(): - if param not in managed_params: - continue - previous = fqn_to_param.get(fqn) - if previous is not None and previous is not param: - raise ValueError( - f"param_dtype_map FQN {fqn!r} is ambiguous across the fully_shard modules" - ) - fqn_to_param[fqn] = param - unknown_fqns = sorted(set(param_dtype_map).difference(fqn_to_param)) - if unknown_fqns: - raise ValueError( - "param_dtype_map contains FQNs that do not name a parameter managed " - f"by this fully_shard call: {unknown_fqns}" - ) - param_overrides = { - fqn_to_param[fqn]: dtype for fqn, dtype in param_dtype_map.items() - } - effective_dtypes = { - param_overrides.get(param, mp_policy.param_dtype) or param.dtype - for param in params - if param.requires_grad - } - if len(effective_dtypes) > 1 and mp_policy.reduce_dtype is None: - raise ValueError("Mixed parameter dtypes require an explicit reduce_dtype") - - _ORIGINAL_PARAM_GROUP_INIT( - self, - params, - modules, - mesh_info, - post_forward_mesh_info, - device, - shard_placement_fn, - mp_policy, - offload_policy, + param_dtype_map = ( + mp_policy.param_dtype_map + if isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) + else None ) + param_overrides: dict[nn.Parameter, torch.dtype] = {} + if param_dtype_map: + managed_params = set(params) + fqn_to_param: dict[str, nn.Parameter] = {} + for module in modules: + for fqn, param in module.named_parameters(): + if param not in managed_params: + continue + previous = fqn_to_param.get(fqn) + # A managed FQN must not identify different parameter tensors. + if previous is not None and previous is not param: + raise ValueError( + f"param_dtype_map FQN {fqn!r} is ambiguous across the fully_shard modules" + ) + fqn_to_param[fqn] = param + unknown_fqns = sorted(set(param_dtype_map).difference(fqn_to_param)) + if unknown_fqns: + raise ValueError( + "param_dtype_map contains FQNs that do not name a parameter managed " + f"by this fully_shard call: {unknown_fqns}" + ) + param_overrides = { + fqn_to_param[fqn]: dtype for fqn, dtype in param_dtype_map.items() + } + effective_dtypes = { + param_overrides.get(param, mp_policy.param_dtype) or param.dtype + for param in params + if param.requires_grad + } + if len(effective_dtypes) > 1 and mp_policy.reduce_dtype is None: + raise ValueError( + "Mixed parameter dtypes require an explicit reduce_dtype" + ) - for fsdp_param, param in zip(self.fsdp_params, params, strict=True): + self.fsdp_params = [] + for param, module_info in zip(params, param_module_infos): override = param_overrides.get(param) - fsdp_param._param_dtype_override = override - effective_param_dtype = ( - override if override is not None else mp_policy.param_dtype + param_mp_policy = ( + replace( + mp_policy, + param_dtype=( + override if override is not None else mp_policy.param_dtype + ), + param_dtype_map=None, + ) + if param_dtype_map + else mp_policy ) - fsdp_param.mp_policy = replace( - mp_policy, - param_dtype=effective_param_dtype, - param_dtype_map=None, + fsdp_param = FSDPParam( + param, + module_info, + mesh_info, + post_forward_mesh_info, + device, + shard_placement_fn, + param_mp_policy, + offload_policy, ) + fsdp_param._param_dtype_override = override + self.fsdp_params.append(fsdp_param) # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map + self.mesh_info = mesh_info + self.post_forward_mesh_info = post_forward_mesh_info + # pyrefly: ignore [read-only] + self.device = device + self.device_handle = _get_device_handle(device.type) + self.mp_policy = mp_policy + self.offload_policy = offload_policy + self._training_state = TrainingState.IDLE + # Group's sharded state always matches its parameters' sharded states + self._sharded_state = ShardedState.SHARDED + self._module_fqn: str | None = None # prefixed from root module + # Only consider resetting sharded parameters once in lazy init since it + # can incur nontrivial overhead to reset them + self._reset_sharded_params: bool = False + + # - Hook state + self._module_to_pre_save_state_dict_hook_handle: _ModuleToHandleDict = {} + self._module_to_pre_load_state_dict_hook_handle: _ModuleToHandleDict = {} + self._all_reduce_hook: Callable[[torch.Tensor], None] | None = None + self._all_gather_comm: AllGather = DefaultAllGather() + self._all_gather_output = torch.empty(0, device=self.device) + self._reduce_scatter_comm: ReduceScatter = DefaultReduceScatter() + # Optional stream to run the user-defined all-reduce hook in + # Saved here and not in the comm. context because we allow the user to + # specify it, possibly at construction time before lazy init + self._all_reduce_hook_stream: torch.cuda.Stream | None = None + + # - Communication and communication/computation overlap + self.comm_ctx = FSDPCommContext() + # Group's indices in the shared post-forward order + self._post_forward_indices: list[int] = [] + # Whether to reduce gradients at all (whether for FSDP or HSDP) + self.reduce_grads: bool = True + # Whether to all-reduce gradients for HSDP; only used if + # `self.reduce_grads` is true, in which case setting this to false + # means reduce-scatter but no all-reduce + self.all_reduce_grads: bool = True + # Whether to reshard parameters after backward (only useful for + # gradient accumulation) + self.reshard_after_backward: bool = True + # Optional custom factor for the gradient reduction op (e.g. to divide + # by a factor other than the world size) + self.gradient_divide_factor: float | None = None + # Whether reduce-scatter and all-reduce should be issued using only + # summations, potentially with separate pre-/post-scaling. + self.force_sum_reduction_for_comms: bool = False + # `async_op` arg used for pre-forward/pre-backward unshard; can be + # overridden to only do explicit prefetching and avoid inter-stream + # fragmentation from using separate unshard streams + self.unshard_async_op: bool = False + # Whether to unshard in backward: can be overridden by the user if the + # parameters in this group are not needed for backward (e.g. embedding) + self.unshard_in_backward: bool = True + + # - CUDA events for stream synchronization + # Holds the all-gather output buffer, sync objects, and metadata + self._all_gather_result: AllGatherResult | None = None + # Holds the reduce-scatter/all-reduce view-out CUDA event that marks the end of + # the group's post-backward (e.g. reduce-scatter, all-reduce and div), which + # should be waited on at the end of backward + self._post_reduce_event: torch.Event | None = None + # Holds the reshard-after-forward CUDA event when resharding to a + # different world size, which should be waited on in the next unshard + self._reshard_after_forward_event: torch.Event | None = None + + # Only for HSDP, if accumulating gradients without all-reduce, save the + # partial reduce output (only reduce-scattered but not all-reduced) + self._partial_reduce_output: torch.Tensor | None = None + # Holds the all-reduce input and all-reduce event to keep it alive + # until the end of backward (critical when doing bf16 reduction with + # fp32 parameters since the all-reduce input is allocated in the RS + # stream and will have no refs to it after being upcast to fp32) + self._all_reduce_state: AllReduceState | None = None + +# Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. def _patched_init_dtype_attrs( self: FSDPParam, mp_policy: TorchMixedPrecisionPolicy, ) -> None: - if ( - not isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) - or not mp_policy.param_dtype_map - ): - _ORIGINAL_INIT_DTYPE_ATTRS(self, mp_policy) - return - # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map # param_dtype, reduce_dtype = (mp_policy.param_dtype, mp_policy.reduce_dtype) # self.orig_dtype = self.sharded_param.dtype + # # Clamp `reduce_dtype` to `None` if no casting is required: since + # # gradients are computed in `param_dtype`, if `reduce_dtype` matches, + # # then we do not need extra casting # if reduce_dtype == param_dtype: # reduce_dtype = None + # # Clamp `param_dtype` to `None` if no casting is required # if param_dtype == self.orig_dtype: # param_dtype = None # self.param_dtype = param_dtype # self.reduce_dtype = reduce_dtype + # # None indicates that the mixed precision is not enabled # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map + has_param_dtype_map = ( + isinstance(mp_policy, ParamDtypeMixedPrecisionPolicy) + and bool(mp_policy.param_dtype_map) + ) param_dtype = ( self._param_dtype_override if self._param_dtype_override is not None @@ -211,12 +300,18 @@ def _patched_init_dtype_attrs( ) reduce_dtype = mp_policy.reduce_dtype self.orig_dtype = self.sharded_param.dtype - if mp_policy.param_dtype_map is None and reduce_dtype == param_dtype: + # Clamp `reduce_dtype` to `None` if no casting is required: since + # gradients are computed in `param_dtype`, if `reduce_dtype` matches, + # then we do not need extra casting + # Per-parameter mixed dtypes require one explicit group reduce dtype. + if not has_param_dtype_map and reduce_dtype == param_dtype: reduce_dtype = None + # Clamp `param_dtype` to `None` if no casting is required if param_dtype == self.orig_dtype: param_dtype = None self.param_dtype = param_dtype self.reduce_dtype = reduce_dtype + # None indicates that the mixed precision is not enabled # MILES_PATCH_REPLACEMENT_END: fsdp-param-dtype-map @@ -243,32 +338,29 @@ def use_foreach_copy(fsdp_param: FSDPParam) -> bool: # foreach_copy_inputs: list[torch.Tensor] = [] # foreach_copy_input_numels: list[int] = [] # + # # 1st pass: for foreach-copy parameters, get inputs and metadata for the + # # foreach copy, and for the others, actually get their all-gather inputs # for i, fsdp_param in enumerate(fsdp_params): # if use_foreach_copy(fsdp_param): # foreach_copy_indices.append(i) # all_gather_input = ( # fsdp_param._sharded_param_data # if fsdp_param.sharded_state == ShardedState.SHARDED - # else cast( - # torch.Tensor, fsdp_param._sharded_post_forward_param_data - # ) + # else cast(torch.Tensor, fsdp_param._sharded_post_forward_param_data) # ) # foreach_copy_inputs.append(all_gather_input) # foreach_copy_input_numels.append(all_gather_input.numel()) # else: # param_all_gather_inputs[i] = fsdp_param.all_gather_inputs # + # # 2nd pass: use foreach copy to compute the remaining all-gather inputs # if foreach_copy_inputs: # fsdp_param_0 = fsdp_params[foreach_copy_indices[0]] # param_dtype, device = fsdp_param_0.param_dtype, fsdp_param_0.device # flat_foreach_copy_input = torch.empty( - # (sum(foreach_copy_input_numels),), - # device=device, - # dtype=param_dtype, - # ) - # splits = torch.split( - # flat_foreach_copy_input, foreach_copy_input_numels + # (sum(foreach_copy_input_numels),), device=device, dtype=param_dtype # ) + # splits = torch.split(flat_foreach_copy_input, foreach_copy_input_numels) # torch._foreach_copy_(splits, foreach_copy_inputs) # for i, split in zip(foreach_copy_indices, splits): # param_all_gather_inputs[i] = [split] @@ -323,10 +415,10 @@ def _patched_foreach_reduce( reduce_dtype: torch.dtype | None, device: torch.device, gradient_divide_factor: float | None, - all_reduce_group: dist.ProcessGroup | None, + all_reduce_group: dist.ProcessGroup | None, # not `None` iff HSDP all_reduce_stream: torch.Stream, all_reduce_grads: bool, - partial_reduce_output: torch.Tensor | None, + partial_reduce_output: torch.Tensor | None, # only used for HSDP all_reduce_hook: Callable[[torch.Tensor], None] | None, force_sum_reduction_for_comms: bool = False, ) -> tuple[ @@ -348,8 +440,7 @@ def _patched_foreach_reduce( # # Check this at runtime since it could be a real runtime error if e.g. # # fp8 weights do not produce the correct higher precision gradients # _raise_assert_with_print( - # "FSDP reduce-scatter expects uniform gradient dtype but got " - # f"{grad_dtypes}" + # f"FSDP reduce-scatter expects uniform gradient dtype but got {grad_dtypes}" # ) # grad_dtype = unsharded_grads[0].dtype # reduce_dtype = reduce_dtype or grad_dtype @@ -565,10 +656,7 @@ def _patched_foreach_reduce_scatter_copy_in( reduce_scatter_input = reduce_scatter_input.view(world_size, -1) # MILES_PATCH_UPSTREAM_BEGIN: fsdp-param-dtype-map # torch.ops.fsdp.chunk_cat( - # unsharded_grads, - # dim=0, - # num_chunks=world_size, - # out=reduce_scatter_input, + # unsharded_grads, dim=0, num_chunks=world_size, out=reduce_scatter_input # ) # MILES_PATCH_UPSTREAM_END: fsdp-param-dtype-map # MILES_PATCH_REPLACEMENT_BEGIN: fsdp-param-dtype-map @@ -581,6 +669,9 @@ def _patched_foreach_reduce_scatter_copy_in( ) return + # Pack each parameter's padded rank chunks directly into the rank-major + # reduce-scatter input, grouping by source dtype to batch cast-and-copy + # without an intermediate buffer. copy_infos: dict[ torch.dtype, tuple[list[torch.Tensor], list[torch.Tensor]] ] = {} @@ -588,7 +679,7 @@ def _patched_foreach_reduce_scatter_copy_in( output_offset = 0 for grad in unsharded_grads: chunk_size = math.ceil(grad.size(0) / world_size) - trailing_numel = grad.numel() // grad.size(0) + trailing_numel = math.prod(grad.shape[1:]) padded_chunk_numel = chunk_size * trailing_numel destinations, sources = copy_infos.setdefault(grad.dtype, ([], [])) for rank in range(world_size): @@ -631,8 +722,9 @@ def apply_param_dtype_map_patch() -> None: f"torch=={_EXPECTED_TORCH_VERSION}, got {torch.__version__}" ) - _verify_source("FSDPParamGroup.__init__", _ORIGINAL_PARAM_GROUP_INIT) + _verify_source("FSDPParam.__init__", _ORIGINAL_PARAM_INIT) _verify_source("FSDPParam.init_dtype_attrs", _ORIGINAL_INIT_DTYPE_ATTRS) + _verify_source("FSDPParamGroup.__init__", _ORIGINAL_PARAM_GROUP_INIT) _verify_source( "_get_param_all_gather_inputs", _ORIGINAL_GET_PARAM_ALL_GATHER_INPUTS,