Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / hsdp_param.py: 89%
599 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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: MixedPrecisionPolicy = MixedPrecisionPolicy(),
140 offload_policy: Optional[OffloadPolicy] = None,
141 device: Optional[torch.device] = None,
142 source_shard_info: Optional[SourceShardMetaInfo] = None,
143 ) -> None:
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._src_parameter_id = id(param)
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.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.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.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.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 dtype attributes, applying source-parameter overrides before policy defaults."""
668 custom_dtypes = (mp_policy.custom_params or {}).get(self._src_parameter_id, {})
669 param_dtype = custom_dtypes.get("param_dtype", mp_policy.param_dtype)
670 reduce_dtype = custom_dtypes.get("reduce_dtype", mp_policy.reduce_dtype)
671 self.orig_dtype = self.sharded_param.dtype
672 if reduce_dtype == param_dtype:
673 reduce_dtype = None
674 if param_dtype == self.orig_dtype:
675 param_dtype = None
676 self.param_dtype = param_dtype
677 self.reduce_dtype = reduce_dtype
678 self.apply_grad_on_fp32_main_grad = mp_policy.apply_grad_on_fp32_main_grad
680 def init_unsharded_param_buffers(
681 self,
682 all_gather_input_numels: list[int],
683 all_gather_input_dtypes: list[torch.dtype],
684 world_size: int,
685 device: torch.device,
686 force_recreate: bool = False,
687 ):
688 """
689 Allocate buffers that hold unsharded parameter data.
691 Args:
692 all_gather_input_numels: Number of elements per input shard.
693 all_gather_input_dtypes: Dtype of each input shard.
694 world_size: Number of ranks in the shard process group.
695 device: Device on which to allocate the output buffers.
696 force_recreate: If True, always recreate buffers even if already initialized.
697 """
698 if not force_recreate and len(self.unsharded_param_buffers) > 0:
699 return # already initialized
700 if force_recreate and hasattr(self, "_unsharded_param"):
701 raise RuntimeError(
702 "Cannot recreate unsharded_param_buffers after initializing the stable "
703 "unsharded parameter."
704 )
705 self.unsharded_param_buffers = [
706 torch.empty(torch.Size([numel * world_size]), dtype=dtype, device=device)
707 for numel, dtype in zip(all_gather_input_numels, all_gather_input_dtypes)
708 ]
710 def init_unsharded_param(self) -> None:
711 """Initialize the stable unsharded parameter from its final output storage."""
712 if len(self.unsharded_param_buffers) != 1:
713 raise AssertionError(
714 f"Expected 1 unsharded_param_buffer, got {len(self.unsharded_param_buffers)}"
715 )
717 if self.allgather_comm_ctx.allgather_output is not None:
718 packed_shape = list(self.sharded_size)
719 packed_shape[0] *= self.shard_world_size
720 chunks = torch.chunk(
721 self.allgather_comm_ctx.allgather_output.view(packed_shape),
722 self.shard_world_size,
723 dim=0,
724 )
725 # pylint: disable=W0212
726 with torch.autograd._unsafe_preserve_version_counter(
727 self.unsharded_param_buffers[0]
728 ):
729 torch.cat(
730 chunks,
731 dim=self.hsdp_placement.dim,
732 out=self.unsharded_param_buffers[0].view(self._orig_size),
733 )
734 self.allgather_comm_ctx.allgather_output.untyped_storage().resize_(0)
735 self.allgather_comm_ctx.allgather_output = None
737 if hasattr(self, "_unsharded_param"):
738 # Keep one stable ``_unsharded_param`` object across unshard cycles:
739 # autograd-facing module state captured during forward must still be
740 # the same object in backward. Only its storage is refreshed above,
741 # which also avoids reusing stale weights after ``optimizer.step()``
742 # mutates the sharded local shard alone (the non-dim-0 unpack path
743 # materializes a contiguous copy, so a stale ``.data`` would not see
744 # the update).
745 return
747 unsharded_param = torch.as_strided(
748 self.unsharded_param_buffers[0],
749 size=self._orig_size,
750 stride=self._contiguous_orig_stride,
751 storage_offset=0,
752 )
753 if self.source_shard_info is not None and self.source_shard_info.origin_is_dtensor:
754 unsharded_param = DTensor.from_local(
755 unsharded_param,
756 self.source_shard_info.mesh,
757 self.source_shard_info.placements,
758 )
759 self._unsharded_param = nn.Parameter(
760 unsharded_param,
761 requires_grad=self.sharded_param.requires_grad,
762 )
764 def to_sharded(self) -> None:
765 """Install the sharded parameter on managed modules and update state."""
766 self._setattr_on_modules(self.sharded_param)
767 if self.unsharded_param_buffers[0] is not self._sharded_param_data:
768 self.free_unsharded_param()
769 self.sharded_state = ShardedState.SHARDED
771 def to_unsharded(self) -> None:
772 """Install the unsharded parameter on managed modules and update state."""
773 set_requires_grad_if_needed(self.sharded_param, self._unsharded_param)
774 self._setattr_on_modules(self._unsharded_param)
775 self.sharded_state = ShardedState.UNSHARDED
777 def _setattr_on_modules(self, param: nn.Parameter) -> None:
778 """Set parameter on module and shared modules, preserving pointer consistency."""
779 if getattr(self._module_info.module.__setattr__, "__func__", None) is nn.Module.__setattr__:
780 # fast path
781 self._module_info.module._parameters[self._module_info.param_name] = param
782 else:
783 # slow path
784 setattr(self._module_info.module, self._module_info.param_name, param)
785 self._parameter_hook_migrator._save_backward_hooks(self.sharded_param)
786 self._parameter_hook_migrator._migrate_backward_hooks(param)
787 # Iterate through all modules that share this parameter to prevent pointer desync.
788 for shared_module, shared_param_name in zip(
789 self._module_info.shared_modules, self._module_info.shared_param_names
790 ):
791 if getattr(shared_module.__setattr__, "__func__", None) is nn.Module.__setattr__:
792 shared_module._parameters[shared_param_name] = param
793 else:
794 setattr(shared_module, shared_param_name, param)
796 def to_sharded_dtensor(self, tensor: torch.Tensor) -> DTensor:
797 """
798 Converts a local tensor representing either the sharded parameter or
799 sharded gradient to DTensor.
801 The parameter-owned sharding specification already contains the
802 logical shape, stride, dtype, and placements. Reusing it avoids the
803 metadata deepcopy performed by ``DTensor.from_local`` when explicit
804 shape and stride are supplied. The returned DTensor intentionally
805 shares this layout with the sharded parameter.
806 """
807 return DTensor.from_local_with_layout(
808 tensor,
809 self._sharding_spec,
810 shape=self._sharding_spec.tensor_shape,
811 )
813 def to_accumulated_grad_if_needed(self) -> None:
814 """Move the current unsharded gradient into the accumulation buffer."""
815 if not hasattr(self, "_unsharded_param") or self._unsharded_param.grad is None:
816 return
817 # Keep local gradients alive across no-sync / delayed-sync steps even
818 # after the parameter transitions back to the sharded view.
819 unsharded_grad = self._unsharded_param.grad
820 self._unsharded_param.grad = None
821 if self.reduce_dtype is not None and unsharded_grad.dtype != self.reduce_dtype:
822 unsharded_grad = unsharded_grad.to(self.reduce_dtype)
823 if self.unsharded_accumulated_grad is None:
824 self.unsharded_accumulated_grad = unsharded_grad
825 else:
826 self.unsharded_accumulated_grad += unsharded_grad
828 def accumulate_unsharded_grad_if_needed(self) -> None:
829 """Add the current unsharded gradient to an existing accumulation buffer."""
830 if not hasattr(self, "_unsharded_param"):
831 return
832 if (
833 self.unsharded_accumulated_grad is not None
834 and self.unsharded_param.grad is not None
835 ):
836 grad = self.unsharded_param.grad
837 if self.reduce_dtype is not None and grad.dtype != self.reduce_dtype:
838 grad = grad.to(self.reduce_dtype)
839 self.unsharded_param.grad = None
840 self.unsharded_accumulated_grad += grad
842 def alloc_unsharded_param_buffers(self) -> None:
843 """
844 Restore unsharded parameter buffers to their full capacity.
845 unsharded_param_buffer is the final storage which should be reffereced by self._unsharded_param
846 """
847 for tensor in self.unsharded_param_buffers:
848 expected_size = tensor.numel() * tensor.itemsize
849 storage = tensor.untyped_storage()
850 if storage.size() != expected_size:
851 storage.resize_(expected_size)
853 def free_unsharded_param(self) -> None:
854 """Release storage of the unsharded parameter buffers."""
855 for tensor in self.unsharded_param_buffers:
856 storage = tensor.untyped_storage()
857 if storage.size() != 0:
858 storage.resize_(0)
860 @property
861 def all_gather_inputs(self) -> list[torch.Tensor]:
862 """Return the local sharded tensor to use as input for all-gather, applying dtype cast if needed."""
863 self._assert_in_states(ShardedState.SHARDED)
864 sharded_param_data = self._sharded_param_data
865 if self.offload_to_cpu:
866 sharded_param_data = sharded_param_data.to(
867 self.device, non_blocking=True
868 )
869 if self.param_dtype is not None and self.param_dtype != sharded_param_data.dtype:
870 return [sharded_param_data.to(self.param_dtype)]
871 return [sharded_param_data]
873 @property
874 def unsharded_param(self) -> nn.Parameter:
875 """Return the full unsharded parameter after all-gather."""
876 return self._unsharded_param
878 @property
879 def unsharded_grad_data(self) -> torch.Tensor:
880 """
881 Get the unsharded gradient data as a local tensor.
882 """
883 return self._unsharded_param.grad
885 @property
886 def unsharded_accumulated_grad_data(self) -> torch.Tensor:
887 """
888 Get the unsharded accumulated gradient data as a local tensor.
889 """
890 return self.unsharded_accumulated_grad
892 def _assert_in_states(self, *states: ShardedState) -> None:
893 """Assert current state is one of expected states."""
894 if self.sharded_state not in states:
895 raise AssertionError(
896 f"Expected sharded_state in {states}, got {self.sharded_state}"
897 )
899 def _resolve_reset_param(self):
900 """Resolve the (possibly swapped) module param for ``reset_sharded_param``.
902 Refreshes ``self.sharded_param`` for the DTensor case and returns the
903 current module parameter for the caller to re-shard.
904 """
905 module_info = self._module_info
906 new_param = getattr(module_info.module, module_info.param_name)
907 if new_param is self.sharded_param:
908 return new_param
909 # Ensure object identity is preserved after parameter conversion.
910 if torch.__future__.get_swap_module_params_on_conversion():
911 raise AssertionError(
912 f"Expects swap_tensors to preserve object but got {new_param} "
913 f"instead of {self.sharded_param}"
914 )
915 if isinstance(new_param, DTensor):
916 self.sharded_param = new_param
917 if not getattr(self.sharded_param, "_hsdp_param_initialized", None):
918 # reset _hsdp_param_initialized flag.
919 self.sharded_param._hsdp_param_initialized = True
920 # If new_param is a plain Tensor, keep the existing 'self.sharded_param' ref;
921 # only its _local_tensor / _sharded_param_data are refreshed below.
922 return new_param
924 def _is_same_sharded_local_tensor(self, local_tensor: torch.Tensor) -> bool:
925 """Return whether communication storage already aliases the local tensor."""
926 if not isinstance(self._sharded_param_data, torch.Tensor):
927 return False
928 sharded_data_ptr = self._sharded_param_data.untyped_storage().data_ptr()
929 return (
930 # Empty shards may have a zero data pointer and must still be rebuilt.
931 sharded_data_ptr > 0
932 and sharded_data_ptr == local_tensor.untyped_storage().data_ptr()
933 )
935 def _validate_reset_local_tensor(
936 self,
937 local_tensor: torch.Tensor,
938 ) -> torch.Tensor:
939 """Validate the actual local shape after parameter conversion."""
940 if local_tensor.size() != self.sharded_size:
941 raise AssertionError(
942 f"Expected sharded_size to be {self.sharded_size}, got {local_tensor.size()}"
943 )
944 return local_tensor
946 def _pin_reset_local_tensor_if_needed(
947 self,
948 local_tensor: torch.Tensor,
949 ) -> tuple[torch.Tensor, bool]:
950 """Move reset storage to pinned CPU memory when the policy requires it."""
951 if self.pin_memory and not local_tensor.is_pinned():
952 return local_tensor.cpu().pin_memory(), True
953 return local_tensor, False
955 def _refresh_sharded_local_tensor(
956 self,
957 local_tensor: torch.Tensor,
958 shard_dim: int,
959 ) -> None:
960 """Rebuild padded communication storage and refresh the DTensor local view."""
961 # A converted parameter may still reference the forward graph through
962 # DTensor's local TensorImpl. Storage owned by HSDP must be an
963 # off-graph leaf so this reference cannot retain the previous device
964 # allocation after an offload or parameter conversion.
965 local_tensor = local_tensor.detach()
966 actual_shard_length = self.sharded_size[shard_dim]
967 if self.sharded_size == self.padded_sharded_param_size:
968 local_tensor = local_tensor.contiguous()
969 self._sharded_param_data = local_tensor.view(-1)
970 # Keep the communication view's base non-trainable while the
971 # local shard has its own requires_grad flag and shared storage.
972 local_view = local_tensor.detach()
973 else:
974 padded_local_tensor = local_tensor.new_zeros(self.padded_sharded_param_size)
975 if self.pin_memory:
976 padded_local_tensor = padded_local_tensor.pin_memory()
977 if local_tensor.numel() > 0:
978 padded_local_tensor.narrow(
979 shard_dim,
980 0,
981 actual_shard_length,
982 ).copy_(local_tensor)
983 self._sharded_param_data = padded_local_tensor.view(-1)
984 local_view = padded_local_tensor.narrow(
985 shard_dim,
986 0,
987 actual_shard_length,
988 ).detach()
989 set_requires_grad_if_needed(self.sharded_param, local_view)
990 self.sharded_param._local_tensor = local_view
991 self._update_shardedparam_storage_forcely()
992 if not self.sharded_param._local_tensor.is_contiguous():
993 raise AssertionError(
994 "Expected sharded_param._local_tensor to be contiguous"
995 )
997 def reset_sharded_param(self) -> None:
998 """Reset sharded param after load_state_dict.
1000 This method may replace the DTensor local storage. It intentionally
1001 detaches the rebuilt persistent storage from autograd because it is
1002 not part of a forward computation.
1003 """
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 """Start gathering the full parameter when no gather is already pending.
1107 Args:
1108 async_op: Whether to launch the all-gather asynchronously.
1109 """
1110 if self.allgather_comm_ctx.allgather_handle is not None:
1111 # Already triggered by HSDPState.prefetch(), so return directly.
1112 return # no-op
1113 self._get_unsharded_param_data(async_op=async_op)
1116 def wait_for_unshard(self) -> None:
1117 """Wait for the parameter all-gather and expose the unsharded parameter."""
1118 self._assert_in_states(ShardedState.SHARDED)
1120 if self.allgather_comm_ctx.allgather_handle is not None:
1121 self.allgather_comm_ctx.allgather_handle.wait()
1122 self.allgather_comm_ctx.allgather_handle = None
1124 self.init_unsharded_param()
1125 self.to_unsharded()
1127 def shard(self) -> None:
1128 """
1129 Transition parameter from unsharded back to sharded state.
1130 """
1131 self._assert_in_states(ShardedState.UNSHARDED)
1132 self.to_sharded()
1134 def reduce_scatter_grad(
1135 self,
1136 async_op: bool = True,
1137 reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
1138 output_buffer: Optional[torch.Tensor] = None,
1139 ) -> None:
1140 """
1141 Perform reduce-scatter on gradient to reduce and shard the full gradient.
1143 Args:
1144 async_op: Whether to execute asynchronously.
1145 reduce_op: do reduce-scatter avg or sum.
1146 output_buffer: Optional pre-allocated output buffer for fused all-reduce.
1147 When provided, reduce_scatter writes directly into this buffer,
1148 enabling zero-copy fusion with subsequent all_reduce operations.
1149 The buffer must have ``padded_sharded_param_size.numel()`` elements
1150 and the reduction dtype.
1151 """
1152 if self.unsharded_accumulated_grad is not None:
1153 self._grad = self.unsharded_accumulated_grad_data
1154 self.unsharded_accumulated_grad = None
1155 else:
1156 self._grad = self.unsharded_grad_data
1157 self.unsharded_param.grad = None
1158 self._grad = self._grad.to(self.reduce_comm_dtype(self._grad))
1159 shard_dim = self.hsdp_placement.dim
1160 if self.shard_world_size <= 1:
1161 self._grad = self._grad.view(-1)
1162 elif shard_dim != 0:
1163 grad_chunks = torch.chunk(self._grad, self.shard_world_size, dim=shard_dim)
1164 self._grad = torch.cat(grad_chunks, dim=0).contiguous().view(-1)
1165 grad_chunks = None
1166 else:
1167 padded_unsharded_dim0 = self.padded_sharded_param_size[0] * self.shard_world_size
1168 if self._grad.size(0) != padded_unsharded_dim0:
1169 padded_unsharded_size = torch.Size((padded_unsharded_dim0, *self._grad.shape[1:]))
1170 padded_grad = self._grad.new_zeros(padded_unsharded_size)
1171 padded_grad.narrow(0, 0, self._grad.size(0)).copy_(self._grad)
1172 self._grad = padded_grad
1173 self._grad = self._grad.view(-1)
1174 apply_gradient_scaling_factor(self._grad, self.gradient_scaling_factor)
1176 shard_process_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None
1177 if shard_process_group is None or self.shard_world_size <= 1:
1178 if output_buffer is not None:
1179 output_buffer.copy_(self._grad)
1180 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
1181 else:
1182 self.reduce_scatter_comm_ctx.reduce_scatter_output = self._grad
1183 self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
1184 self._grad = None
1185 return
1187 # Calculate output size
1188 output_numel = self._grad.numel() // self.shard_world_size
1189 # Use provided output buffer or allocate a new one
1190 if output_buffer is not None:
1191 if output_buffer.numel() != output_numel:
1192 raise ValueError(
1193 f"output_buffer size mismatch: expected {output_numel}, got {output_buffer.numel()}"
1194 )
1195 if output_buffer.dtype != self._grad.dtype:
1196 raise ValueError(
1197 f"output_buffer dtype mismatch: expected {self._grad.dtype}, got {output_buffer.dtype}"
1198 )
1199 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
1200 else:
1201 self.reduce_scatter_comm_ctx.reduce_scatter_output = torch.empty(
1202 output_numel,
1203 dtype=self._grad.dtype,
1204 device=self._grad.device,
1205 )
1206 # Execute reduce_scatter_tensor
1207 self.reduce_scatter_comm_ctx.reduce_scatter_handle = dist.reduce_scatter_tensor(
1208 self.reduce_scatter_comm_ctx.reduce_scatter_output,
1209 self._grad,
1210 op=reduce_op,
1211 group=shard_process_group,
1212 async_op=async_op,
1213 )
1214 if not async_op:
1215 self._grad = None
1217 def all_reduce_grad(
1218 self,
1219 async_op: bool = True,
1220 reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
1221 ) -> None:
1222 """
1223 All-reduce the current reduce-scatter output across the replicate mesh.
1225 Args:
1226 async_op: Whether to execute asynchronously.
1227 reduce_op: Reduction operation for the replicate dimension.
1229 The output and optional asynchronous work are stored in
1230 ``all_reduce_comm_ctx``.
1231 """
1232 grad = self.reduce_scatter_comm_ctx.reduce_scatter_output
1233 if grad is None:
1234 raise RuntimeError("all_reduce_grad requires a completed reduce-scatter output.")
1235 reduce_dtype = self.reduce_comm_dtype(grad)
1236 if grad.dtype != reduce_dtype:
1237 grad = grad.to(reduce_dtype)
1239 replicate_process_group = (
1240 self.mesh_info.replicate_process_group if isinstance(self.mesh_info, DDPMeshInfo) else None
1241 )
1242 if replicate_process_group is None or self.replicate_world_size <= 1:
1243 self.all_reduce_comm_ctx.all_reduce_output = grad
1244 self.all_reduce_comm_ctx.all_reduce_handle = None
1245 return
1247 self.all_reduce_comm_ctx.all_reduce_handle = dist.all_reduce(
1248 grad,
1249 op=reduce_op,
1250 group=replicate_process_group,
1251 async_op=async_op,
1252 )
1253 self.all_reduce_comm_ctx.all_reduce_output = grad
1255 def all_reduce_source_replicate_grad_inplace(
1256 self,
1257 reduced_grad: torch.Tensor,
1258 reduce_op: dist.ReduceOp,
1259 ) -> None:
1260 """All-reduce a final gradient over replicated source-layout axes, in place.
1262 ``reduced_grad`` is modified in place; the caller keeps using its own
1263 reference and must not expect a returned tensor. No-op when the source
1264 layout has no replicated mesh axis, so callers may invoke it
1265 unconditionally.
1267 Args:
1268 reduced_grad: Final local gradient after FSDP/HSDP reduction.
1269 reduce_op: Reduction operation shared with the DP communication path.
1270 """
1271 if self.source_shard_info is None or not self.source_shard_info.placements:
1272 return
1273 # Use the same deduplicated source layout as sharded storage. Native
1274 # DTensor compute metadata can include DP/CP axes already reduced by
1275 # FSDP; reducing those axes again would multiply validate-mode grads.
1276 source_mesh, source_placements = self._storage_source_layout
1277 replicate_mesh_dims = tuple(
1278 mesh_dim
1279 for mesh_dim, placement in enumerate(source_placements)
1280 if placement.is_replicate()
1281 )
1282 if not replicate_mesh_dims:
1283 return
1284 mesh_dim_names = source_mesh.mesh_dim_names
1285 if mesh_dim_names is None:
1286 raise ValueError(
1287 "TP shard mesh must define mesh_dim_names to all-reduce replicated gradients."
1288 )
1289 replicate_mesh_dim_names = tuple(mesh_dim_names[mesh_dim] for mesh_dim in replicate_mesh_dims)
1290 replicate_mesh = source_mesh[replicate_mesh_dim_names].flatten()
1291 replicate_world_size = replicate_mesh.size()
1292 if replicate_world_size <= 1:
1293 return
1294 dist.all_reduce(
1295 reduced_grad,
1296 op=reduce_op,
1297 group=replicate_mesh.get_group(),
1298 async_op=False,
1299 )
1302def set_requires_grad_if_needed(
1303 src_tensor: torch.Tensor, dst_tensor: torch.Tensor
1304) -> None:
1305 """set dst_tensor requires_grads from src_tensor if needed."""
1306 if src_tensor.requires_grad != dst_tensor.requires_grad:
1307 dst_tensor.requires_grad_(src_tensor.requires_grad)