Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / fully_shard / param.py: 82%
589 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15# Adapted from https://github.com/pytorch/pytorch/blob/release/2.6/torch/distributed/fsdp/_fully_shard/_fsdp_param.py
16# enhanced with fully_shard parameter management
17# ============================================================================
18"""HSDP parameter"""
19# pylint: disable=W0212
20from dataclasses import dataclass
21from typing import Callable, List, Optional, cast
23import torch
24import torch.distributed as dist
25from torch import nn
26from torch._prims_common import make_contiguous_strides_for
28from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
29from hyper_parallel.core.dtensor.dtensor import DTensor
30from hyper_parallel.core.dtensor.layout import Layout
31from hyper_parallel.core.dtensor.placement_types import Replicate, Shard, StridedShard
32from hyper_parallel.core.fully_shard.hsdp_param import HSDPParamV2
33from hyper_parallel.core.fully_shard.hsdp_utils import (
34 ParamModuleInfo,
35 ShardedState,
36 apply_gradient_scaling_factor,
37)
38from hyper_parallel.core.fully_shard.utils import (
39 CPUOffloadPolicy,
40 DataParallelMeshInfo,
41 DDPMeshInfo,
42 FSDPMeshInfo,
43 HSDPMeshInfo,
44 MixedPrecisionPolicy,
45 OffloadPolicy,
46 SourceShardMetaInfo,
47)
48from hyper_parallel.core.utils import compute_local_shape_and_global_offset_by_ceil_chunk
51def _copy_without_bumping_version(dst: torch.Tensor, src: torch.Tensor) -> None:
52 """Copy into ``dst`` while preserving its autograd version counter."""
53 # pylint: disable=W0212
54 with torch.autograd._unsafe_preserve_version_counter(dst):
55 dst.copy_(src)
58@dataclass
59class ReduceScatterCommCtx:
60 """Per-parameter reduce-scatter output and asynchronous work."""
62 reduce_scatter_output: Optional[torch.Tensor] = None
63 reduce_scatter_handle: Optional[dist.Work] = None
66@dataclass
67class AllReduceCommCtx:
68 """Per-parameter all-reduce output and asynchronous work."""
70 all_reduce_output: Optional[torch.Tensor] = None
71 all_reduce_handle: Optional[dist.Work] = None
74@dataclass
75class AllGatherCommCtx:
76 """Per-parameter all-gather output and asynchronous work."""
78 allgather_output: Optional[torch.Tensor] = None
79 allgather_handle: Optional[dist.Work] = None
82class ParameterHookMigrator:
83 """Preserve parameter backward hooks across HSDP parameter replacement."""
85 def __init__(self) -> None:
86 self._orig_param_hooks: List[Callable] = []
87 self._saved_hook_ids: set[int] = set()
89 def _save_backward_hooks(self, param: nn.Parameter) -> None:
90 """Save backward hooks from a parameter, deduplicated by hook identity."""
91 if not hasattr(param, "_backward_hooks") or param._backward_hooks is None:
92 return
94 for _, hook_func in param._backward_hooks.items():
95 hook_func_id = id(hook_func)
96 if hook_func_id not in self._saved_hook_ids:
97 self._orig_param_hooks.append(hook_func)
98 self._saved_hook_ids.add(hook_func_id)
100 def _migrate_backward_hooks(self, new_param: nn.Parameter) -> None:
101 """Register saved backward hooks on a replacement parameter once."""
102 if not self._orig_param_hooks or hasattr(new_param, "migrate_backward_hooks_run_once"):
103 return
105 for hook_func in self._orig_param_hooks:
106 try:
107 if new_param.requires_grad:
108 new_param.register_hook(hook_func)
109 except RuntimeError:
110 # Skip hook registration if the parameter does not require gradients.
111 pass
112 new_param.migrate_backward_hooks_run_once = True
115class TorchHSDPParamV2(HSDPParamV2):
116 """
117 Torch HSDP parameter.
118 """
120 def __init__(
121 self,
122 param: nn.Parameter,
123 module_info: ParamModuleInfo,
124 mesh_info: DataParallelMeshInfo,
125 shard_placement_fn: Optional[Callable[[nn.Parameter], Optional[Shard]]] = None,
126 mp_policy: Optional[MixedPrecisionPolicy] = None,
127 offload_policy: Optional[OffloadPolicy] = None,
128 device: Optional[torch.device] = None,
129 source_shard_info: Optional[SourceShardMetaInfo] = None,
130 ):
131 """
132 Initialize TorchHSDPParamV2 and shard the parameter.
134 Args:
135 param (nn.Parameter): The original full parameter to shard.
136 module_info (ParamModuleInfo): Ownership and shared-weight metadata.
137 mesh_info (DataParallelMeshInfo): Mesh topology for shard/replicate dimensions.
138 shard_placement_fn (Callable, optional): Returns a Shard placement for the parameter,
139 or None to use default (Shard(0)).
140 mp_policy (MixedPrecisionPolicy, optional): Mixed precision dtype policy.
141 offload_policy (OffloadPolicy, optional): CPU offload policy.
142 device (torch.device, optional): Target device for the sharded parameter.
143 source_shard_info (SourceShardMetaInfo, optional): Source TP/EP layout metadata, built by
144 the owning ``HSDPState``. Must be supplied (with ``origin_is_dtensor=True``)
145 whenever ``param`` is a native DTensor, and omitted otherwise.
147 Raises:
148 ValueError: If ``source_shard_info.origin_is_dtensor`` disagrees with whether
149 ``param`` is a native DTensor.
150 """
151 self._module_info: ParamModuleInfo = module_info
152 self.mesh_info = mesh_info
153 self.mp_policy = mp_policy
154 self.device = device
155 self.orig_dtype = None
156 self.param_dtype = None
157 self.reduce_dtype = None
158 self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy)
159 self.pin_memory = (
160 self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory
161 )
162 self._parameter_hook_migrator = ParameterHookMigrator()
163 # ``source_shard_info`` is built and validated by the owning state
164 # (``_build_param_source_shard_info``): for a native DTensor parameter it always
165 # describes that parameter's own mesh/placements. Only the agreement
166 # between the two is re-checked here, because it is the invariant the
167 # sharding math below depends on.
168 if isinstance(param, DTensor) != (
169 source_shard_info is not None and source_shard_info.origin_is_dtensor
170 ):
171 raise ValueError(
172 "source_shard_info.origin_is_dtensor must be True exactly for native DTensor parameters, "
173 f"got parameter type {type(param).__name__} and source_shard_info={source_shard_info}"
174 )
175 self.source_shard_info = source_shard_info
176 self._orig_param_is_dtensor = (
177 source_shard_info is not None and source_shard_info.origin_is_dtensor
178 )
179 self._orig_dtensor_mesh = source_shard_info.mesh if self._orig_param_is_dtensor else None
180 self._orig_dtensor_placements = (
181 tuple(source_shard_info.placements) if self._orig_param_is_dtensor else None
182 )
183 self._storage_source_layout = self._build_storage_source_layout()
184 self._spmd_shard_mesh_dim = self.mesh_info.shard_mesh_dim
185 self._spmd_replicate_mesh_dim = self.mesh_info.replicate_mesh_dim
186 self._init_sharded_param(param, shard_placement_fn)
187 self.unsharded_accumulated_grad = None
188 self._param_fqn: Optional[str] = None
189 self.unsharded_param_buffers: List[torch.Tensor] = []
190 self.allgather_comm_ctx = AllGatherCommCtx()
191 self._post_load_hook_handle = (
192 module_info.module.register_load_state_dict_post_hook(
193 lambda *args, **kwargs: self.reset_sharded_param()
194 )
195 )
196 self.reduce_scatter_comm_ctx = ReduceScatterCommCtx()
197 self.all_reduce_comm_ctx = AllReduceCommCtx()
198 self._parameter_hook_migrator._save_backward_hooks(param)
199 self._grad = None
200 # Keep reduce-scatter accumulation in reduce_dtype until the final
201 # micro-step performs the optional replicate all-reduce.
202 self._reduce_partial_output = None
203 self.gradient_scaling_factor = None
205 def _get_base_spmd_placements(self) -> tuple:
206 if self.source_shard_info is not None:
207 # Preserve the source distributed layout and prefix the explicit
208 # DP/FSDP mesh dimensions on the unified mesh.
209 source_mesh, source_placements = self._storage_source_layout
210 self._spmd_mesh = DeviceMesh.concatenate([self.mesh_info.mesh, source_mesh])
211 dp_prefix_placements = tuple(Replicate() for _ in range(self.mesh_info.mesh.ndim))
212 return dp_prefix_placements + source_placements
214 self._spmd_mesh = self.mesh_info.mesh
215 return tuple(Replicate() for _ in range(self._spmd_mesh.ndim))
217 def _build_storage_source_layout(self) -> tuple[Optional[DeviceMesh], tuple]:
218 """Build the immutable source layout used by FSDP storage and reduction.
220 Native DTensor parameters may already describe data-parallel axes that
221 are also covered by ``mesh_info.mesh``. Exclude those axes from the
222 concatenated storage mesh without changing ``source_shard_info``, whose
223 mesh and placements always retain the original compute-layout meaning.
224 Explicit metadata for plain production-mode parameters already contains
225 only its TP/EP source layout and must remain unchanged.
227 Ownership is derived from rank topology rather than dimension names:
228 a source axis is redundant exactly when varying it at this rank's own
229 coordinate never leaves the FSDP domain's rank set. Name matching
230 cannot express this because trainer meshes rename flattened data
231 axes (e.g. ``dp``/``cp`` become ``fsdp_replicate``/``fsdp_shard``).
233 Returns:
234 Source mesh and placements to append to the FSDP mesh, or an empty
235 layout when this parameter has no source-layout metadata.
236 """
237 if self.source_shard_info is None:
238 return None, ()
239 source_mesh = self.source_shard_info.mesh
240 source_placements = tuple(self.source_shard_info.placements)
241 if not self.source_shard_info.origin_is_dtensor:
242 return source_mesh, source_placements
244 source_mesh_dim_names = source_mesh.mesh_dim_names or ()
245 fsdp_ranks = self._fsdp_domain_ranks()
246 coordinate = source_mesh.get_coordinate() if fsdp_ranks is not None else None
247 if coordinate is None:
248 # Rank topology is unavailable (e.g. a topology-only mesh), so no
249 # axis can be proven redundant. Keep the full layout, matching the
250 # pre-dedup behavior.
251 return source_mesh, source_placements
253 storage_dim_names = tuple(
254 dim_name
255 for dim_index, dim_name in enumerate(source_mesh_dim_names)
256 if not self._source_dim_covered_by_fsdp(
257 source_mesh, dim_index, coordinate, fsdp_ranks
258 )
259 )
260 if not storage_dim_names or storage_dim_names == source_mesh_dim_names:
261 return source_mesh, source_placements
263 placement_by_dim = dict(zip(source_mesh_dim_names, source_placements))
264 return (
265 source_mesh[storage_dim_names],
266 tuple(placement_by_dim[dim_name] for dim_name in storage_dim_names),
267 )
269 def _fsdp_domain_ranks(self) -> Optional[frozenset]:
270 """Ranks owned by this FSDP unit's data-parallel domain.
272 Returns ``None`` when the mesh carries no usable rank topology, in
273 which case no source axis can be proven redundant.
274 """
275 mesh = getattr(self.mesh_info, "mesh", None)
276 mesh_tensor = getattr(mesh, "mesh", None)
277 if not isinstance(mesh_tensor, torch.Tensor):
278 return None
279 return frozenset(int(rank) for rank in mesh_tensor.flatten().tolist())
281 @staticmethod
282 def _source_dim_covered_by_fsdp(source_mesh, dim_index, coordinate, fsdp_ranks) -> bool:
283 """Whether varying one source dim stays inside the FSDP rank domain.
285 The fiber is taken at this rank's own coordinate: peer ranks reached
286 by varying a data-parallel axis share this rank's other coordinates
287 (e.g. its TP slot), which is exactly the population FSDP reduces over.
288 """
289 index = list(coordinate)
290 mesh_tensor = source_mesh.mesh
291 for extent in range(source_mesh.size(dim_index)):
292 index[dim_index] = extent
293 if int(mesh_tensor[tuple(index)]) not in fsdp_ranks:
294 return False
295 return True
297 def _apply_data_parallel_placements(self, placements: list, shard_placement: Shard) -> tuple:
298 if len(placements) != self._spmd_mesh.ndim:
299 raise AssertionError(
300 f"Expected {self._spmd_mesh.ndim} unified placements, got {len(placements)}: {placements}"
301 )
302 if (
303 isinstance(self.mesh_info, DDPMeshInfo)
304 and self._spmd_replicate_mesh_dim is not None
305 and not self._orig_param_is_dtensor
306 ):
307 placements[self._spmd_replicate_mesh_dim] = Replicate()
308 if (
309 isinstance(self.mesh_info, FSDPMeshInfo)
310 and self._spmd_shard_mesh_dim is not None
311 ):
312 # If TP/EP already shards the same tensor dimension, fully_shard must
313 # use StridedShard so the unified placement preserves the intended
314 # shard order on the concatenated mesh.
315 split_factor = 1
316 for mesh_idx, placement in enumerate(placements):
317 if mesh_idx == self._spmd_shard_mesh_dim:
318 continue
319 if placement.is_shard(shard_placement.dim):
320 split_factor *= self._spmd_mesh.mesh_shape[mesh_idx]
321 placements[self._spmd_shard_mesh_dim] = (
322 StridedShard(shard_placement.dim, split_factor=split_factor)
323 if split_factor > 1
324 else shard_placement
325 )
326 return tuple(placements)
328 def _build_sharding_spec(
329 self,
330 source_param: nn.Parameter,
331 source_local_tensor: torch.Tensor,
332 ) -> Layout:
333 """Build the final layout after data and model parallel sharding.
335 A dual-mode parameter is a plain tensor containing only its TP/EP-local
336 shard, so its local shape cannot represent the final logical shape.
337 Source distribution metadata restores that shape before the FSDP
338 placements are applied. This method only constructs the layout and has
339 no parameter lifecycle side effects.
341 Args:
342 source_param: Parameter received before FSDP partitioning.
343 source_local_tensor: Local tensor partitioned by FSDP.
345 Returns:
346 Layout containing all FSDP, TP, and EP placements.
347 """
348 logical_global_stride = None
349 if isinstance(source_param, DTensor):
350 logical_global_size = source_param.size()
351 logical_global_stride = source_param.layout.tensor_stride
352 elif self.source_shard_info is not None:
353 source_sharding_spec = Layout.from_device_mesh(self.source_shard_info.mesh)
354 source_sharding_spec.set_placements(self.source_shard_info.placements)
355 source_sharding_spec.placement_to_tensor_map(source_local_tensor.ndim)
356 logical_global_size = source_sharding_spec.get_global_shape(
357 source_local_tensor.size()
358 )
359 else:
360 logical_global_size = source_local_tensor.size()
362 if logical_global_stride is None:
363 logical_global_stride = make_contiguous_strides_for(logical_global_size)
365 sharding_spec = Layout.from_device_mesh(self._spmd_mesh)
366 sharding_spec.set_placements(self._spmd_placements)
367 sharding_spec.placement_to_tensor_map(source_local_tensor.ndim)
368 sharding_spec.set_tensor_meta(
369 logical_global_size,
370 logical_global_stride,
371 source_local_tensor.dtype,
372 )
373 return sharding_spec
375 @property
376 def reduce_partial_output(self) -> Optional[torch.Tensor]:
377 """Return reduce-scatter results accumulated before the final micro-step."""
378 return self._reduce_partial_output
380 @reduce_partial_output.setter
381 def reduce_partial_output(self, value: Optional[torch.Tensor]) -> None:
382 self._reduce_partial_output = value
384 def reduce_comm_dtype(self, grad: Optional[torch.Tensor] = None) -> torch.dtype:
385 """Resolve the communication dtype owned by this parameter.
387 Args:
388 grad: Optional gradient used when no mixed-precision reduction dtype
389 is configured.
391 Returns:
392 The dtype used by reduce-scatter and all-reduce buffers.
393 """
394 if self.reduce_dtype is not None:
395 return self.reduce_dtype
396 if grad is not None:
397 return grad.dtype
398 if self.unsharded_accumulated_grad is not None:
399 return self.unsharded_accumulated_grad_data.dtype
400 if self.unsharded_param.grad is not None:
401 return self.unsharded_grad_data.dtype
402 return self.orig_dtype
404 def reduce_scatter_output(self) -> Optional[torch.Tensor]:
405 """
406 Get the reduce-scatter output tensor and wait for asynchronous operation to complete.
408 Returns:
409 torch.Tensor: The sharded gradient tensor after reduce-scatter operation.
410 """
411 if self.reduce_scatter_comm_ctx.reduce_scatter_handle is not None:
412 self.reduce_scatter_comm_ctx.reduce_scatter_handle.wait()
413 self._grad.untyped_storage().resize_(0)
414 self._grad = None
415 self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
416 return self.reduce_scatter_comm_ctx.reduce_scatter_output
418 def clear_reduce_scatter_output(self) -> None:
419 """Clear the reduce-scatter output tensor to free memory."""
420 self.reduce_scatter_comm_ctx.reduce_scatter_output = None
421 self._grad = None
423 def all_reduce_output(self) -> Optional[torch.Tensor]:
424 """
425 Get the all-reduce output tensor and wait for asynchronous operation to complete.
427 Returns:
428 torch.Tensor: The reduced gradient tensor after all-reduce operation.
429 """
430 if self.all_reduce_comm_ctx.all_reduce_handle is not None:
431 self.all_reduce_comm_ctx.all_reduce_handle.wait()
432 self.all_reduce_comm_ctx.all_reduce_handle = None
433 return self.all_reduce_comm_ctx.all_reduce_output
435 def clear_all_reduce_output(self) -> None:
436 """Clear the all-reduce output tensor to free memory."""
437 self.all_reduce_comm_ctx.all_reduce_output = None
439 def clear_unsharded_source_grad(self):
440 if self.unsharded_accumulated_grad_data is not None:
441 self.unsharded_accumulated_grad = None
442 elif self.unsharded_param.grad is not None:
443 self.unsharded_param.grad = None
445 def apply_reduced_grad(self, reduced_grad: torch.Tensor) -> bool:
446 """
447 Apply reduced gradient to the sharded parameter.
449 Reshapes ``reduced_grad`` to match the local shard, optionally
450 offloads to CPU, then accumulates or assigns onto
451 ``hsdp_param.sharded_param.grad``.
453 Args:
454 reduced_grad (torch.Tensor): Gradient after reduce-scatter
455 and/or all-reduce.
457 Returns:
458 Whether the current stream must synchronize for CPU offload.
459 """
460 sharded_grad = None
461 if not self.mp_policy.apply_grad_on_fp32_main_grad:
462 sharded_grad = self.sharded_param.grad
463 else:
464 if not hasattr(self.sharded_param, "main_grad"):
465 self.sharded_param.main_grad = None
466 sharded_grad = self.sharded_param.main_grad
467 reduced_grad = (
468 reduced_grad.reshape(-1)
469 .narrow(0, 0, self.sharded_size.numel())
470 .view(self.sharded_size)
471 )
472 if (
473 not self.mp_policy.apply_grad_on_fp32_main_grad
474 and reduced_grad.dtype != self.orig_dtype
475 ):
476 reduced_grad = reduced_grad.to(self.orig_dtype)
477 to_accumulate_grad = sharded_grad is not None
478 need_synchronize = False
479 if self.offload_to_cpu:
480 non_blocking = self.pin_memory and not to_accumulate_grad
481 reduced_grad = reduced_grad.to(
482 torch.device("cpu"), non_blocking=non_blocking
483 )
484 need_synchronize = True
485 if sharded_grad is None:
486 if not self.mp_policy.apply_grad_on_fp32_main_grad:
487 self.sharded_param.grad = self.to_sharded_dtensor(reduced_grad)
488 else:
489 self.sharded_param.main_grad = self.to_sharded_dtensor(reduced_grad)
490 self.sharded_param.grad = None
491 else:
492 if not self.mp_policy.apply_grad_on_fp32_main_grad:
493 self.sharded_param.grad._local_tensor += reduced_grad
494 else:
495 self.sharded_param.main_grad._local_tensor += reduced_grad
496 self.sharded_param.grad = None
497 if self.unsharded_accumulated_grad_data is not None:
498 self.unsharded_accumulated_grad = None
499 elif self.unsharded_param.grad is not None:
500 self.unsharded_param.grad = None
501 return need_synchronize
503 def _resolve_hsdp_placement(
504 self,
505 param: nn.Parameter,
506 shard_placement_fn: Optional[Callable],
507 ) -> Shard:
508 """Validate and normalize the fully_shard placement for one parameter."""
509 if param.device != self.device and param.device.type != "meta":
510 raise AssertionError(
511 f"Expects the parameter to already be moved to device {self.device} but got {param.device}"
512 )
514 hsdp_placement = shard_placement_fn(param) if shard_placement_fn else None
515 if hsdp_placement is None:
516 hsdp_placement = Shard(0)
517 elif hsdp_placement.dim < 0:
518 # if dim is negative, add the number of dimensions of the parameter
519 hsdp_placement = Shard(hsdp_placement.dim + param.ndim)
521 if not isinstance(hsdp_placement, Shard):
522 raise AssertionError(
523 f"Expected Shard, got {type(hsdp_placement)}: {hsdp_placement}"
524 )
525 return hsdp_placement
527 def _init_shard_metadata(
528 self,
529 param: nn.Parameter,
530 hsdp_placement: Shard,
531 ) -> tuple[list, torch.Tensor, int, int]:
532 """Initialize parameter shape and mesh metadata used by sharding."""
533 self.hsdp_placement = hsdp_placement
534 base_placements = list(self._get_base_spmd_placements())
535 param_data = param.to_local() if self._orig_param_is_dtensor else param
536 shard_dim = hsdp_placement.dim
537 if param_data.ndim == 0:
538 raise ValueError("fully_shard does not support scalar parameters")
539 if shard_dim < 0 or shard_dim >= param_data.ndim:
540 raise ValueError(
541 f"Invalid fully_shard dim {shard_dim} for parameter "
542 f"{self._module_info.param_name} with shape {tuple(param_data.shape)}"
543 )
544 self._orig_size = param_data.size()
545 self._contiguous_orig_stride = make_contiguous_strides_for(self._orig_size)
546 if isinstance(self.mesh_info, FSDPMeshInfo):
547 self.shard_rank = self.mesh_info.shard_mesh_rank
548 self.shard_world_size = self.mesh_info.shard_mesh_size
549 else:
550 self.shard_rank = 0
551 self.shard_world_size = 1
553 if isinstance(self.mesh_info, DDPMeshInfo):
554 self.replicate_world_size = self.mesh_info.replicate_mesh_size
555 else:
556 self.replicate_world_size = 1
557 self.is_replicate_param = (
558 isinstance(self.mesh_info, DDPMeshInfo)
559 and not isinstance(self.mesh_info, HSDPMeshInfo)
560 )
561 dim_shard_size = (param_data.size(shard_dim) + self.shard_world_size - 1) // self.shard_world_size
562 return base_placements, param_data, shard_dim, dim_shard_size
564 def _init_shard_placements(
565 self,
566 param_data: torch.Tensor,
567 shard_dim: int,
568 base_placements: list,
569 ) -> None:
570 """Build data-parallel placements for the local shard."""
571 if shard_dim != 0 and param_data.size(shard_dim) % self.shard_world_size != 0:
572 raise NotImplementedError(
573 f"fully_shard only supports uneven sharding on dim=0, but parameter "
574 f"{self._module_info.param_name} has shape {tuple(param_data.shape)}, "
575 f"shard dim {shard_dim}, and world size {self.shard_world_size}"
576 )
577 spmd_placements = list(
578 self._apply_data_parallel_placements(
579 base_placements,
580 self.hsdp_placement,
581 )
582 )
583 if param_data.size(shard_dim) % self.shard_world_size != 0:
584 if self._spmd_shard_mesh_dim is None:
585 raise AssertionError("Uneven FSDP sharding requires a shard mesh dimension")
586 fsdp_placement = spmd_placements[self._spmd_shard_mesh_dim]
587 if isinstance(fsdp_placement, StridedShard):
588 fsdp_placement = StridedShard(
589 fsdp_placement.dim,
590 fsdp_placement.split_factor,
591 uneven_shard=True,
592 )
593 else:
594 fsdp_placement = Shard(fsdp_placement.dim, uneven_shard=True)
595 spmd_placements[self._spmd_shard_mesh_dim] = fsdp_placement
596 self._spmd_placements = tuple(spmd_placements)
598 def _build_sharded_param_data(
599 self,
600 param_data: torch.Tensor,
601 shard_dim: int,
602 dim_shard_size: int,
603 ) -> torch.Tensor:
604 """Create the actual local shard and its fixed-size communication storage."""
605 local_shape, global_offset = compute_local_shape_and_global_offset_by_ceil_chunk(
606 param_data.size(),
607 shard_dim,
608 self.shard_world_size,
609 self.shard_rank,
610 )
611 actual_shard_offset = global_offset[shard_dim]
612 actual_shard_length = local_shape[shard_dim]
613 sharded_param = param_data.narrow(
614 shard_dim,
615 actual_shard_offset,
616 actual_shard_length,
617 ).clone().contiguous()
618 self.sharded_size = sharded_param.size()
619 self.contiguous_sharded_stride = make_contiguous_strides_for(self.sharded_size)
620 padded_sharded_size = list(param_data.size())
621 padded_sharded_size[shard_dim] = dim_shard_size
622 self.padded_sharded_param_size = torch.Size(padded_sharded_size)
623 if self.offload_to_cpu and not sharded_param.is_meta:
624 sharded_param = sharded_param.cpu()
625 if self.pin_memory:
626 sharded_param = sharded_param.pin_memory()
628 if self.sharded_size == self.padded_sharded_param_size:
629 self._sharded_param_data = sharded_param.view(-1)
630 else:
631 padded_sharded_param = sharded_param.new_zeros(self.padded_sharded_param_size)
632 if self.pin_memory and not padded_sharded_param.is_meta:
633 padded_sharded_param = padded_sharded_param.pin_memory()
634 if sharded_param.numel() > 0:
635 padded_sharded_param.narrow(
636 shard_dim,
637 0,
638 actual_shard_length,
639 ).copy_(sharded_param)
640 self._sharded_param_data = padded_sharded_param.view(-1)
641 sharded_param = padded_sharded_param.narrow(
642 shard_dim,
643 0,
644 actual_shard_length,
645 )
646 return sharded_param
648 @torch.no_grad()
649 def _init_sharded_param(
650 self,
651 param: nn.Parameter,
652 shard_placement_fn: Optional[Callable],
653 ) -> None:
654 """Initialize the persistent sharded parameter and communication storage."""
655 hsdp_placement = self._resolve_hsdp_placement(param, shard_placement_fn)
656 base_placements, param_data, shard_dim, dim_shard_size = self._init_shard_metadata(
657 param,
658 hsdp_placement,
659 )
660 self._init_shard_placements(
661 param_data,
662 shard_dim,
663 base_placements,
664 )
665 sharded_param = self._build_sharded_param_data(
666 param_data,
667 shard_dim,
668 dim_shard_size,
669 )
670 self._sharding_spec = self._build_sharding_spec(param, param_data)
672 self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))
673 self.sharded_param._layout = self._sharding_spec
674 self.sharded_param._placements = tuple(self._sharding_spec.placements)
675 self.sharded_param.requires_grad_(param.requires_grad)
676 self._setattr_on_modules(self.sharded_param)
677 # after init, self.sharded_param replaces original param, gradients must accumulate to this Parameter's grad
678 self.sharded_param._hsdp_param_initialized = True
679 self.sharded_state = ShardedState.SHARDED
680 self.param_dtype = None
681 self.reduce_dtype = None
683 def init_dtype_attrs(self, mp_policy: MixedPrecisionPolicy) -> None:
684 """Initialize param_dtype and reduce_dtype from the mixed precision policy."""
685 param_dtype, reduce_dtype = (mp_policy.param_dtype, mp_policy.reduce_dtype)
686 self.orig_dtype = self.sharded_param.dtype
687 if reduce_dtype == param_dtype:
688 reduce_dtype = None
689 if param_dtype == self.orig_dtype:
690 param_dtype = None
691 self.param_dtype = param_dtype
692 self.reduce_dtype = reduce_dtype
694 def init_unsharded_param_buffers(
695 self,
696 all_gather_input_numels: list[int],
697 all_gather_input_dtypes: list[torch.dtype],
698 world_size: int,
699 device: torch.device,
700 force_recreate: bool = False,
701 ):
702 """
703 Allocate buffers that hold unsharded parameter data.
705 Args:
706 all_gather_input_numels: Number of elements per input shard.
707 all_gather_input_dtypes: Dtype of each input shard.
708 world_size: Number of ranks in the shard process group.
709 device: Device on which to allocate the output buffers.
710 force_recreate: If True, always recreate buffers even if already initialized.
711 """
712 if not force_recreate and len(self.unsharded_param_buffers) > 0:
713 return # already initialized
714 if force_recreate and hasattr(self, "_unsharded_param"):
715 raise RuntimeError(
716 "Cannot recreate unsharded_param_buffers after initializing the stable "
717 "unsharded parameter."
718 )
719 self.unsharded_param_buffers = [
720 torch.empty(torch.Size([numel * world_size]), dtype=dtype, device=device)
721 for numel, dtype in zip(all_gather_input_numels, all_gather_input_dtypes)
722 ]
724 def init_unsharded_param(self) -> None:
725 """Initialize the stable unsharded parameter from its final output storage."""
726 if len(self.unsharded_param_buffers) != 1:
727 raise AssertionError(
728 f"Expected 1 unsharded_param_buffer, got {len(self.unsharded_param_buffers)}"
729 )
731 if self.allgather_comm_ctx.allgather_output is not None:
732 packed_shape = list(self.sharded_size)
733 packed_shape[0] *= self.shard_world_size
734 chunks = torch.chunk(
735 self.allgather_comm_ctx.allgather_output.view(packed_shape),
736 self.shard_world_size,
737 dim=0,
738 )
739 # pylint: disable=W0212
740 with torch.autograd._unsafe_preserve_version_counter(
741 self.unsharded_param_buffers[0]
742 ):
743 torch.cat(
744 chunks,
745 dim=self.hsdp_placement.dim,
746 out=self.unsharded_param_buffers[0].view(self._orig_size),
747 )
748 self.allgather_comm_ctx.allgather_output.untyped_storage().resize_(0)
749 self.allgather_comm_ctx.allgather_output = None
751 if hasattr(self, "_unsharded_param"):
752 # Keep one stable ``_unsharded_param`` object across unshard cycles:
753 # autograd-facing module state captured during forward must still be
754 # the same object in backward. Only its storage is refreshed above,
755 # which also avoids reusing stale weights after ``optimizer.step()``
756 # mutates the sharded local shard alone (the non-dim-0 unpack path
757 # materializes a contiguous copy, so a stale ``.data`` would not see
758 # the update).
759 return
761 unsharded_param = torch.as_strided(
762 self.unsharded_param_buffers[0],
763 size=self._orig_size,
764 stride=self._contiguous_orig_stride,
765 storage_offset=0,
766 )
767 if self.source_shard_info is not None and self.source_shard_info.origin_is_dtensor:
768 unsharded_param = DTensor.from_local(
769 unsharded_param,
770 self.source_shard_info.mesh,
771 self.source_shard_info.placements,
772 )
773 self._unsharded_param = nn.Parameter(
774 unsharded_param,
775 requires_grad=self.sharded_param.requires_grad,
776 )
778 def to_sharded(self) -> None:
779 self._setattr_on_modules(self.sharded_param)
780 if self.unsharded_param_buffers[0] is not self._sharded_param_data:
781 self.free_unsharded_param()
782 self.sharded_state = ShardedState.SHARDED
784 def to_unsharded(self) -> None:
785 set_requires_grad_if_needed(self.sharded_param, self._unsharded_param)
786 self._setattr_on_modules(self._unsharded_param)
787 self.sharded_state = ShardedState.UNSHARDED
789 def _setattr_on_modules(self, param: nn.Parameter) -> None:
790 """Set parameter on module and shared modules, preserving pointer consistency."""
791 if getattr(self._module_info.module.__setattr__, "__func__", None) is nn.Module.__setattr__:
792 # fast path
793 self._module_info.module._parameters[self._module_info.param_name] = param
794 else:
795 # slow path
796 setattr(self._module_info.module, self._module_info.param_name, param)
797 self._parameter_hook_migrator._save_backward_hooks(self.sharded_param)
798 self._parameter_hook_migrator._migrate_backward_hooks(param)
799 # Iterate through all modules that share this parameter to prevent pointer desync.
800 for shared_module, shared_param_name in zip(
801 self._module_info.shared_modules, self._module_info.shared_param_names
802 ):
803 if getattr(shared_module.__setattr__, "__func__", None) is nn.Module.__setattr__:
804 shared_module._parameters[shared_param_name] = param
805 else:
806 setattr(shared_module, shared_param_name, param)
808 def to_sharded_dtensor(self, tensor: torch.Tensor) -> DTensor:
809 """
810 Converts a local tensor representing either the sharded parameter or
811 sharded gradient to DTensor.
812 """
813 sharded_dtensor = DTensor.from_local(
814 tensor,
815 self._sharding_spec.mesh,
816 self._sharding_spec.placements,
817 shape=self._sharding_spec.tensor_shape,
818 stride=self._sharding_spec.tensor_stride,
819 )
820 sharded_dtensor._layout = self._sharding_spec
821 sharded_dtensor._placements = tuple(self._sharding_spec.placements)
822 return sharded_dtensor
824 def to_accumulated_grad_if_needed(self) -> None:
825 if self._unsharded_param.grad is None:
826 return
827 # Keep local gradients alive across no-sync / delayed-sync steps even
828 # after the parameter transitions back to the sharded view.
829 unsharded_grad = self._unsharded_param.grad
830 self._unsharded_param.grad = None
831 if self.reduce_dtype is not None and unsharded_grad.dtype != self.reduce_dtype:
832 unsharded_grad = unsharded_grad.to(self.reduce_dtype)
833 if self.unsharded_accumulated_grad is None:
834 self.unsharded_accumulated_grad = unsharded_grad
835 else:
836 self.unsharded_accumulated_grad += unsharded_grad
838 def accumulate_unsharded_grad_if_needed(self) -> None:
839 if (
840 self.unsharded_accumulated_grad is not None
841 and self.unsharded_param.grad is not None
842 ):
843 grad = self.unsharded_param.grad
844 if self.reduce_dtype is not None and grad.dtype != self.reduce_dtype:
845 grad = grad.to(self.reduce_dtype)
846 self.unsharded_param.grad = None
847 self.unsharded_accumulated_grad += grad
849 def alloc_unsharded_param_buffers(self) -> None:
850 """
851 Restore unsharded parameter buffers to their full capacity.
852 unsharded_param_buffer is the final storage which should be reffereced by self._unsharded_param
853 """
854 for tensor in self.unsharded_param_buffers:
855 expected_size = tensor.numel() * tensor.itemsize
856 storage = tensor.untyped_storage()
857 if storage.size() != expected_size:
858 storage.resize_(expected_size)
860 def free_unsharded_param(self) -> None:
861 """Release storage of the unsharded parameter buffers."""
862 for tensor in self.unsharded_param_buffers:
863 storage = tensor.untyped_storage()
864 if storage.size() != 0:
865 storage.resize_(0)
867 @property
868 def all_gather_inputs(self) -> list[torch.Tensor]:
869 """Return the local sharded tensor to use as input for all-gather, applying dtype cast if needed."""
870 self._assert_in_states(ShardedState.SHARDED)
871 sharded_param_data = self._sharded_param_data
872 if self.offload_to_cpu:
873 sharded_param_data = sharded_param_data.to(
874 self.device, non_blocking=True
875 )
876 if self.param_dtype is not None and self.param_dtype != sharded_param_data.dtype:
877 return [sharded_param_data.to(self.param_dtype)]
878 return [sharded_param_data]
880 @property
881 def _sharded_local_tensor(self) -> torch.Tensor:
882 """Return the local tensor backing the persistent sharded DTensor."""
883 return self.sharded_param._local_tensor
885 @property
886 def unsharded_param(self) -> nn.Parameter:
887 """Return the full unsharded parameter after all-gather."""
888 return self._unsharded_param
890 @property
891 def unsharded_grad_data(self) -> torch.Tensor:
892 """
893 Get the unsharded gradient data as a local tensor.
894 """
895 return self._unsharded_param.grad
897 @property
898 def unsharded_accumulated_grad_data(self) -> torch.Tensor:
899 """
900 Get the unsharded accumulated gradient data as a local tensor.
901 """
902 return self.unsharded_accumulated_grad
904 def _assert_in_states(self, *states: ShardedState) -> None:
905 """Assert current state is one of expected states."""
906 if self.sharded_state not in states:
907 raise AssertionError(
908 f"Expected sharded_state in {states}, got {self.sharded_state}"
909 )
911 def _resolve_reset_param(self):
912 """Resolve the (possibly swapped) module param for ``reset_sharded_param``.
914 Refreshes ``self.sharded_param`` for the DTensor case and returns the
915 current module parameter for the caller to re-shard.
916 """
917 module_info = self._module_info
918 new_param = getattr(module_info.module, module_info.param_name)
919 if new_param is self.sharded_param:
920 return new_param
921 # Ensure object identity is preserved after parameter conversion.
922 if torch.__future__.get_swap_module_params_on_conversion():
923 raise AssertionError(
924 f"Expects swap_tensors to preserve object but got {new_param} "
925 f"instead of {self.sharded_param}"
926 )
927 if isinstance(new_param, DTensor):
928 self.sharded_param = new_param
929 if not getattr(self.sharded_param, "_hsdp_param_initialized", None):
930 # reset _hsdp_param_initialized flag.
931 self.sharded_param._hsdp_param_initialized = True
932 # If new_param is a plain Tensor, keep the existing 'self.sharded_param' ref;
933 # only its _local_tensor / _sharded_param_data are refreshed below.
934 return new_param
936 def _is_same_sharded_local_tensor(self, local_tensor: torch.Tensor) -> bool:
937 """Return whether communication storage already aliases the local tensor."""
938 if not isinstance(self._sharded_param_data, torch.Tensor):
939 return False
940 sharded_data_ptr = self._sharded_param_data.untyped_storage().data_ptr()
941 return (
942 # Empty shards may have a zero data pointer and must still be rebuilt.
943 sharded_data_ptr > 0
944 and sharded_data_ptr == local_tensor.untyped_storage().data_ptr()
945 )
947 def _validate_reset_local_tensor(
948 self,
949 local_tensor: torch.Tensor,
950 ) -> torch.Tensor:
951 """Validate the actual local shape after parameter conversion."""
952 if local_tensor.size() != self.sharded_size:
953 raise AssertionError(
954 f"Expected sharded_size to be {self.sharded_size}, got {local_tensor.size()}"
955 )
956 return local_tensor
958 def _pin_reset_local_tensor_if_needed(
959 self,
960 local_tensor: torch.Tensor,
961 ) -> tuple[torch.Tensor, bool]:
962 """Move reset storage to pinned CPU memory when the policy requires it."""
963 if self.pin_memory and not local_tensor.is_pinned():
964 return local_tensor.cpu().pin_memory(), True
965 return local_tensor, False
967 def _refresh_sharded_local_tensor(
968 self,
969 local_tensor: torch.Tensor,
970 shard_dim: int,
971 ) -> None:
972 """Rebuild padded communication storage and refresh the DTensor local view."""
973 actual_shard_length = self.sharded_size[shard_dim]
974 if self.sharded_size == self.padded_sharded_param_size:
975 local_tensor = local_tensor.contiguous()
976 self._sharded_param_data = local_tensor.view(-1)
977 local_view = local_tensor.detach()
978 else:
979 padded_local_tensor = local_tensor.new_zeros(self.padded_sharded_param_size)
980 if self.pin_memory:
981 padded_local_tensor = padded_local_tensor.pin_memory()
982 if local_tensor.numel() > 0:
983 padded_local_tensor.narrow(
984 shard_dim,
985 0,
986 actual_shard_length,
987 ).copy_(local_tensor)
988 self._sharded_param_data = padded_local_tensor.view(-1)
989 local_view = padded_local_tensor.narrow(
990 shard_dim,
991 0,
992 actual_shard_length,
993 ).detach()
994 set_requires_grad_if_needed(self.sharded_param, local_view)
995 self.sharded_param._local_tensor = local_view
996 self._update_shardedparam_storage_forcely()
997 if not self.sharded_param._local_tensor.is_contiguous():
998 raise AssertionError(
999 "Expected sharded_param._local_tensor to be contiguous"
1000 )
1002 def reset_sharded_param(self) -> None:
1003 """Reset sharded param after load_state_dict."""
1004 new_param = self._resolve_reset_param()
1005 local_tensor = new_param._local_tensor if isinstance(new_param, DTensor) else new_param
1006 if local_tensor.is_meta:
1007 return
1008 # local_tensor can be padded twice
1009 # 1st time in fully_shard(model)
1010 # 2nd time in model(input) lazy_init
1011 # 2nd time should be no-op if parameters remain unchanged
1012 # 2nd time shouldn't be no-op if people call model.load_state_dict(...) before lazy_init
1013 # this makes it possible for trainer to call `sd = model.state_dict()` before the training loop
1014 # and use `sd` without calling .state_dict() per iteration
1015 same_local_tensor = self._is_same_sharded_local_tensor(local_tensor)
1016 shard_dim = self.hsdp_placement.dim
1017 if not same_local_tensor:
1018 local_tensor = self._validate_reset_local_tensor(local_tensor)
1019 local_tensor, pinned_local_tensor = self._pin_reset_local_tensor_if_needed(local_tensor)
1020 if not isinstance(self.sharded_param, DTensor):
1021 raise AssertionError(f"Expected DTensor, got {type(self.sharded_param)}")
1022 if not same_local_tensor or pinned_local_tensor:
1023 self._refresh_sharded_local_tensor(local_tensor, shard_dim)
1024 self._sharding_spec.set_tensor_meta(
1025 self._sharding_spec.tensor_shape,
1026 self._sharding_spec.tensor_stride,
1027 local_tensor.dtype,
1028 )
1029 self.sharded_param._layout = self._sharding_spec
1030 self.sharded_param._placements = tuple(self._sharding_spec.placements)
1031 # After ``to_empty`` replaces the module parameter with a plain tensor,
1032 # re-install the DTensor ``nn.Parameter`` so the optimizer and forward
1033 # hooks see the correct object. Idempotent when the module already
1034 # holds ``self.sharded_param`` (same data_ptr → no-op in practice).
1035 self._setattr_on_modules(self.sharded_param)
1037 def _update_shardedparam_storage_forcely(self,):
1038 sharded_param_data = self.sharded_param.data
1039 local_tensor_data = self.sharded_param._local_tensor.data
1040 if (
1041 sharded_param_data.device != local_tensor_data.device
1042 or sharded_param_data.data_ptr() != local_tensor_data.data_ptr()
1043 ):
1044 local_tensor_data.requires_grad_(self.sharded_param.requires_grad)
1045 # TensorImpl keeps storage_ptr, storage_offset, sizes, strides, dtype and metadatas
1046 # so swap TensorImpl can make self.sharded_param ref the self.sharded_param._local_tensor's TensorImpl
1047 # local_tensor_data TensorImpl will deconstruct after this function finish.
1048 torch._C._swap_tensor_impl(self.sharded_param, local_tensor_data) # pylint: disable=W0212
1050 @torch.no_grad()
1051 def _get_unsharded_param_data(self, async_op: bool = False) -> None:
1052 """
1053 Perform all-gather to get unsharded parameter data.
1055 Args:
1056 async_op: Whether to execute asynchronously.
1058 The output buffer and optional asynchronous handle are stored in the
1059 parameter communication context.
1060 """
1061 all_gather_input = self.all_gather_inputs[0]
1063 if self.shard_world_size <= 1 or self.mesh_info.shard_process_group is None:
1064 if len(self.unsharded_param_buffers) == 0:
1065 self.unsharded_param_buffers = [all_gather_input]
1066 elif self.unsharded_param_buffers[0] is not all_gather_input:
1067 # if param_dtype cast or tensor.to caused by cpu_offload
1068 self.alloc_unsharded_param_buffers()
1069 _copy_without_bumping_version(self.unsharded_param_buffers[0], all_gather_input)
1070 self.allgather_comm_ctx.allgather_output = None
1071 self.allgather_comm_ctx.allgather_handle = None
1072 return
1074 self.init_unsharded_param_buffers(
1075 all_gather_input_numels=[all_gather_input.numel()],
1076 all_gather_input_dtypes=[all_gather_input.dtype],
1077 world_size=self.shard_world_size,
1078 device=self.device,
1079 )
1080 self.alloc_unsharded_param_buffers()
1082 self.allgather_comm_ctx.allgather_output = self.unsharded_param_buffers[0]
1083 if self.hsdp_placement.dim != 0:
1084 # Non-dim-0 sharding uses an extra all-gather buffer before
1085 # restoring the original dimension with chunk + cat.
1086 self.allgather_comm_ctx.allgather_output = torch.empty_like(
1087 self.unsharded_param_buffers[0]
1088 )
1089 # pylint: disable=W0212
1090 with torch.autograd._unsafe_preserve_version_counter(
1091 self.allgather_comm_ctx.allgather_output
1092 ):
1093 self.allgather_comm_ctx.allgather_handle = dist.all_gather_into_tensor(
1094 self.allgather_comm_ctx.allgather_output,
1095 all_gather_input,
1096 group=self.mesh_info.shard_process_group,
1097 async_op=async_op,
1098 )
1100 if self.allgather_comm_ctx.allgather_output is self.unsharded_param_buffers[0]:
1101 self.allgather_comm_ctx.allgather_output = None
1104 def unshard(self, async_op: bool = False) -> None:
1105 if self.allgather_comm_ctx.allgather_handle is not None:
1106 # Already triggered by HSDPState.prefetch(), so return directly.
1107 return # no-op
1108 self._get_unsharded_param_data(async_op=async_op)
1111 def wait_for_unshard(self) -> None:
1112 self._assert_in_states(ShardedState.SHARDED)
1114 if self.allgather_comm_ctx.allgather_handle is not None:
1115 self.allgather_comm_ctx.allgather_handle.wait()
1116 self.allgather_comm_ctx.allgather_handle = None
1118 self.init_unsharded_param()
1119 self.to_unsharded()
1121 def shard(self) -> None:
1122 """
1123 Transition parameter from unsharded back to sharded state.
1124 """
1125 self._assert_in_states(ShardedState.UNSHARDED)
1126 self.to_sharded()
1128 def reduce_scatter_grad(
1129 self,
1130 async_op: bool = True,
1131 reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
1132 output_buffer: Optional[torch.Tensor] = None,
1133 ) -> None:
1134 """
1135 Perform reduce-scatter on gradient to reduce and shard the full gradient.
1137 Args:
1138 async_op: Whether to execute asynchronously.
1139 reduce_op: do reduce-scatter avg or sum.
1140 output_buffer: Optional pre-allocated output buffer for fused all-reduce.
1141 When provided, reduce_scatter writes directly into this buffer,
1142 enabling zero-copy fusion with subsequent all_reduce operations.
1143 The buffer must have ``padded_sharded_param_size.numel()`` elements
1144 and the reduction dtype.
1146 The output and optional asynchronous work are stored in
1147 ``reduce_scatter_comm_ctx``.
1148 """
1149 if self.unsharded_accumulated_grad is not None:
1150 grad = self.unsharded_accumulated_grad_data
1151 else:
1152 grad = self.unsharded_grad_data
1153 self._grad = grad.to(self.reduce_comm_dtype(grad))
1154 shard_dim = self.hsdp_placement.dim
1155 if self.shard_world_size <= 1:
1156 self._grad = self._grad.view(-1)
1157 elif shard_dim != 0:
1158 grad_chunks = torch.chunk(self._grad, self.shard_world_size, dim=shard_dim)
1159 self._grad = torch.cat(grad_chunks, dim=0).contiguous().view(-1)
1160 else:
1161 padded_unsharded_dim0 = self.padded_sharded_param_size[0] * self.shard_world_size
1162 if self._grad.size(0) != padded_unsharded_dim0:
1163 padded_unsharded_size = torch.Size((padded_unsharded_dim0, *self._grad.shape[1:]))
1164 padded_grad = self._grad.new_zeros(padded_unsharded_size)
1165 padded_grad.narrow(0, 0, self._grad.size(0)).copy_(self._grad)
1166 self._grad = padded_grad
1167 self._grad = self._grad.view(-1)
1168 apply_gradient_scaling_factor(self._grad, self.gradient_scaling_factor)
1170 shard_process_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None
1171 if shard_process_group is None or self.shard_world_size <= 1:
1172 if output_buffer is not None:
1173 output_buffer.copy_(self._grad)
1174 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
1175 else:
1176 self.reduce_scatter_comm_ctx.reduce_scatter_output = self._grad
1177 self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
1178 return
1180 # Calculate output size
1181 output_numel = self._grad.numel() // self.shard_world_size
1182 # Use provided output buffer or allocate a new one
1183 if output_buffer is not None:
1184 if output_buffer.numel() != output_numel:
1185 raise ValueError(
1186 f"output_buffer size mismatch: expected {output_numel}, got {output_buffer.numel()}"
1187 )
1188 if output_buffer.dtype != self._grad.dtype:
1189 raise ValueError(
1190 f"output_buffer dtype mismatch: expected {self._grad.dtype}, got {output_buffer.dtype}"
1191 )
1192 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
1193 else:
1194 self.reduce_scatter_comm_ctx.reduce_scatter_output = torch.empty(
1195 output_numel,
1196 dtype=self._grad.dtype,
1197 device=self._grad.device,
1198 )
1199 # Execute reduce_scatter_tensor
1200 self.reduce_scatter_comm_ctx.reduce_scatter_handle = dist.reduce_scatter_tensor(
1201 self.reduce_scatter_comm_ctx.reduce_scatter_output,
1202 self._grad,
1203 op=reduce_op,
1204 group=shard_process_group,
1205 async_op=async_op,
1206 )
1208 def all_reduce_grad(
1209 self,
1210 async_op: bool = True,
1211 reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
1212 ) -> None:
1213 """
1214 All-reduce the current reduce-scatter output across the replicate mesh.
1216 Args:
1217 async_op: Whether to execute asynchronously.
1218 reduce_op: Reduction operation for the replicate dimension.
1220 The output and optional asynchronous work are stored in
1221 ``all_reduce_comm_ctx``.
1222 """
1223 grad = self.reduce_scatter_comm_ctx.reduce_scatter_output
1224 if grad is None:
1225 raise RuntimeError("all_reduce_grad requires a completed reduce-scatter output.")
1226 reduce_dtype = self.reduce_comm_dtype(grad)
1227 if grad.dtype != reduce_dtype:
1228 grad = grad.to(reduce_dtype)
1230 replicate_process_group = (
1231 self.mesh_info.replicate_process_group if isinstance(self.mesh_info, DDPMeshInfo) else None
1232 )
1233 if replicate_process_group is None or self.replicate_world_size <= 1:
1234 self.all_reduce_comm_ctx.all_reduce_output = grad
1235 self.all_reduce_comm_ctx.all_reduce_handle = None
1236 return
1238 self.all_reduce_comm_ctx.all_reduce_handle = dist.all_reduce(
1239 grad,
1240 op=reduce_op,
1241 group=replicate_process_group,
1242 async_op=async_op,
1243 )
1244 self.all_reduce_comm_ctx.all_reduce_output = grad
1246 def all_reduce_source_replicate_grad_inplace(
1247 self,
1248 reduced_grad: torch.Tensor,
1249 reduce_op: dist.ReduceOp,
1250 ) -> None:
1251 """All-reduce a final gradient over replicated source-layout axes, in place.
1253 ``reduced_grad`` is modified in place; the caller keeps using its own
1254 reference and must not expect a returned tensor. No-op when the source
1255 layout has no replicated mesh axis, so callers may invoke it
1256 unconditionally.
1258 Args:
1259 reduced_grad: Final local gradient after FSDP/HSDP reduction.
1260 reduce_op: Reduction operation shared with the DP communication path.
1261 """
1262 if self.source_shard_info is None or not self.source_shard_info.placements:
1263 return
1264 # Use the same deduplicated source layout as sharded storage. Native
1265 # DTensor compute metadata can include DP/CP axes already reduced by
1266 # FSDP; reducing those axes again would multiply validate-mode grads.
1267 source_mesh, source_placements = self._storage_source_layout
1268 replicate_mesh_dims = tuple(
1269 mesh_dim
1270 for mesh_dim, placement in enumerate(source_placements)
1271 if placement.is_replicate()
1272 )
1273 if not replicate_mesh_dims:
1274 return
1275 mesh_dim_names = source_mesh.mesh_dim_names
1276 if mesh_dim_names is None:
1277 raise ValueError(
1278 "TP shard mesh must define mesh_dim_names to all-reduce replicated gradients."
1279 )
1280 replicate_mesh_dim_names = tuple(mesh_dim_names[mesh_dim] for mesh_dim in replicate_mesh_dims)
1281 replicate_mesh = source_mesh[replicate_mesh_dim_names].flatten()
1282 replicate_world_size = replicate_mesh.size()
1283 if replicate_world_size <= 1:
1284 return
1285 dist.all_reduce(
1286 reduced_grad,
1287 op=reduce_op,
1288 group=replicate_mesh.get_group(),
1289 async_op=False,
1290 )
1293def set_requires_grad_if_needed(
1294 src_tensor: torch.Tensor, dst_tensor: torch.Tensor
1295) -> None:
1296 """set dst_tensor requires_grads from src_tensor if needed."""
1297 if src_tensor.requires_grad != dst_tensor.requires_grad:
1298 dst_tensor.requires_grad_(src_tensor.requires_grad)