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