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/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d9430378..5ab2beb7 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 @@ -137,12 +138,14 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty raise RuntimeError(f"{component} did not honor meta initialization") checkpoint.sync_model_dtypes(model) full_state = model.state_dict() if rank == 0 else {} + fsdp_parallel_plan = self.model_backend.fsdp_parallel_plan(model) model = apply_fsdp2( model, mesh=self.parallel_state.get_mesh("fsdp"), cpu_offload=self.args.fsdp_cpu_offload, args=self.args, - no_split_modules=self.model_backend.fsdp_no_split_modules(model), + no_split_modules=fsdp_parallel_plan.no_split_modules, + param_dtype_patterns=fsdp_parallel_plan.param_dtype_patterns, ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -609,7 +612,14 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None): +def apply_fsdp2( + model, + mesh=None, + cpu_offload=False, + args=None, + no_split_modules=None, + param_dtype_patterns=None, +): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard offload_policy = CPUOffloadPolicy() if cpu_offload else None @@ -621,19 +631,49 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules param_dtype = _resolve_dtype(args.diffusion_forward_dtype) reduce_dtype = _resolve_dtype(args.fsdp_reduce_dtype) + param_dtype_maps = compile_param_dtype_maps( + model, + modules, + param_dtype_patterns or {}, + param_dtype, + ) + has_param_dtype_overrides = bool( + param_dtype_maps.module_maps or param_dtype_maps.root_map + ) + param_dtype_policy_cls = None + if has_param_dtype_overrides: + from .fsdp_param_dtype_patch import ( + ParamDtypeMixedPrecisionPolicy, + apply_param_dtype_map_patch, + ) + + apply_param_dtype_map_patch() + param_dtype_policy_cls = ParamDtypeMixedPrecisionPolicy logger.info( - f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, reduce_dtype={reduce_dtype}" + f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, " + f"param_dtype={param_dtype}, reduce_dtype={reduce_dtype}, " + f"param_dtype_overrides={param_dtype_maps.override_count} " + f"({param_dtype_maps.override_numel:,} parameters)" ) fsdp_kwargs = { - "mp_policy": MixedPrecisionPolicy( - param_dtype=param_dtype, - reduce_dtype=reduce_dtype, - ), "offload_policy": offload_policy, "mesh": mesh, } + def make_mp_policy(param_dtype_map): + if param_dtype_map: + assert param_dtype_policy_cls is not None + return param_dtype_policy_cls( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + param_dtype_map=param_dtype_map, + ) + return MixedPrecisionPolicy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + ) + if args.gradient_checkpointing: # MixedPrecisionPolicy does not cast buffers; a buffer above param_dtype # makes the ckpt recompute dtype-diverge from the forward and abort. @@ -644,8 +684,16 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules module.register_buffer(name, buf.to(param_dtype), persistent=persistent) for module in modules: - fully_shard(module, **fsdp_kwargs) + fully_shard( + module, + mp_policy=make_mp_policy(param_dtype_maps.module_maps.get(module)), + **fsdp_kwargs, + ) - fully_shard(model, **fsdp_kwargs) + fully_shard( + model, + mp_policy=make_mp_policy(param_dtype_maps.root_map), + **fsdp_kwargs, + ) return model diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index c392ea57..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,6 +76,7 @@ def get_train_pipeline_config_cls(family: str) -> type[TrainPipelineConfig]: class TrainPipelineConfig(abc.ABC): """Base class. Subclass per model family.""" + model_family: str | None = None lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"] needs_timestep_scaling: bool = True optimizer_state_allowed_missing: list[str] = [] 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..002fc17e --- /dev/null +++ b/miles/backends/fsdp_utils/fsdp_param_dtype_patch.py @@ -0,0 +1,768 @@ +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 OffloadPolicy, ReduceScatter +from torch.distributed.fsdp._fully_shard._fsdp_collectives import ( + AllGather, + AllGatherResult, + _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_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 ( + _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", + "foreach_reduce": "46bcdaa0df40823e13922359db131cf329f8adf12eefa72ff16ceb41f9650d90", + "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 +_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 + + +# Copied from PyTorch v2.11.0 at 70d99e998b4955e0049d13a98d77ae1b14db1f45. +def _patched_param_group_init( + self, + params: list[nn.Parameter], + modules: tuple[nn.Module, ...], + mesh_info: DataParallelMeshInfo, + post_forward_mesh_info: FSDPMeshInfo | None, + device: torch.device, + shard_placement_fn: Callable[[nn.Parameter], Shard | None] | None, + mp_policy: TorchMixedPrecisionPolicy, + offload_policy: OffloadPolicy, +) -> None: + 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.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 + 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" + ) + + self.fsdp_params = [] + for param, module_info in zip(params, param_module_infos): + override = param_overrides.get(param) + 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 = 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: + # 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 + else mp_policy.param_dtype + ) + reduce_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 + # 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 + + +# 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_UPSTREAM_BEGIN: fsdp-param-dtype-map + # foreach_copy_indices: list[int] = [] + # 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) + # ) + # 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) + # 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]] + ] = {} + + # 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_REPLACEMENT_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, # not `None` iff HSDP + all_reduce_stream: torch.Stream, + all_reduce_grads: bool, + 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[ + 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_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( + # 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 + # 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( + "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_REPLACEMENT_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_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, + dim=0, + num_chunks=world_size, + out=reduce_scatter_input, + ) + 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]] + ] = {} + padding_views: list[torch.Tensor] = [] + output_offset = 0 + for grad in unsharded_grads: + chunk_size = math.ceil(grad.size(0) / world_size) + 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): + 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_REPLACEMENT_END: fsdp-param-dtype-map + + +def apply_param_dtype_map_patch() -> None: + if getattr(_fsdp_collectives, _PATCH_SENTINEL, False): + return + 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__}" + ) + + _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, + ) + _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/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-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..3378804e --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_worker.py @@ -0,0 +1,255 @@ +import copy +import os +from argparse import Namespace + +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__() + 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.root_scale = nn.Parameter(torch.ones(1)) + self.block = MixedParamDtypeBlock() + self.seen_root_dtype = None + + def forward(self, 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 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(): + 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={}, + ), + ) + 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(): + 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=PARAM_DTYPE_PATTERNS, + ) + + 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) + assert model.seen_root_dtype == 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), + ) + + +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() + + +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/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_fsdp_param_dtype_patch.py b/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py new file mode 100644 index 00000000..14eca7c7 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_fsdp_param_dtype_patch.py @@ -0,0 +1,53 @@ +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(torch, "__version__", "2.12.0+cu130") + + with pytest.raises(RuntimeError, match="requires torch==2.11.0"): + 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(torch, "__version__", "2.11.0+cu130") + 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 new file mode 100644 index 00000000..b5441a4f --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_mixed_precision.py @@ -0,0 +1,120 @@ +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} + assert compiled.override_count == 3 + assert compiled.override_numel == 9 + + +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_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( + model, + list(model.blocks), + {"blocks.*.norm.weight": "bf16"}, + torch.bfloat16, + ) + + 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(): + 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} 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