Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / fully_shard / param_group.py: 82%
541 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15# Adapted from https://github.com/pytorch/pytorch/blob/release/2.6/torch/distributed/fsdp/_fully_shard/_fsdp_param.py
16# enhanced with fully_shard parameter management
17# ============================================================================
18"""HSDP parameter group.
20This module implements fused communication for HSDP (Hybrid Shard Data Parallel) parameters.
21Instead of issuing one all-gather / reduce-scatter per parameter, ``HSDPParamGroup`` packs all
22parameters within a module into a single contiguous buffer and performs one collective operation,
23which reduces kernel launch overhead and improves bandwidth utilization.
25Key components:
26- ``HSDPParamGroup``: Groups all HSDP parameters in a module for fused all-gather (forward)
27 and fused reduce-scatter + all-reduce (backward).
28- ``AllGatherMetadata`` / ``AllGatherMetadataCache``: Caches per-group metadata (dtypes, numels,
29 split sizes) to avoid recomputation across iterations.
30- ``CommContext``: Global context that tracks the in-flight async communication handle and the
31 param group that owns it, enabling pipelined overlap between communication and computation.
32"""
33from typing import List, Optional, NamedTuple, Any
34from dataclasses import dataclass, field
35from contextlib import ExitStack
36import torch
37import torch.distributed as dist
38from torch.distributed import Work
39from hyper_parallel.core.fully_shard.hsdp_utils import apply_gradient_scaling_factor
40from hyper_parallel.core.fully_shard.utils import (
41 MixedPrecisionPolicy,
42 FSDPMeshInfo,
43 DDPMeshInfo,
44 HSDPMeshInfo,
45)
46from hyper_parallel.platform.torch.fully_shard.pack_utils import (
47 build_rs_plan,
48 pack_for_reduce_scatter,
49)
50from hyper_parallel.platform.torch.fully_shard.param import TorchHSDPParamV2
53def get_all_gather_metadata(hsdp_params):
54 """Collect metadata required for fused all-gather from all HSDP parameters.
56 Iterates over each parameter's local shard inputs and records their dtypes and
57 element counts. All parameters must share the same dtype (heterogeneous dtypes
58 are not yet supported).
60 Args:
61 hsdp_params: List of ``TorchHSDPParamV2`` whose ``all_gather_inputs`` will
62 be inspected.
64 Returns:
65 AllGatherMetadata: Aggregated metadata used by ``foreach_all_gather`` to
66 allocate the fused output buffer and perform copy-in/copy-out.
68 Raises:
69 ValueError: If parameters have different dtypes.
70 """
71 param_input_dtypes = []
72 param_input_numels = []
73 inp_split_sizes = []
74 total_input_numel = 0
75 first_dtype = None
77 for hsdp_param in hsdp_params:
78 inputs = hsdp_param.all_gather_inputs
79 if first_dtype is None:
80 first_dtype = inputs[0].dtype
81 elif first_dtype != inputs[0].dtype:
82 raise ValueError("All parameters in the group must have a uniform dtype.")
83 param_dtypes = [t.dtype for t in inputs]
84 param_numels = [t.numel() for t in inputs]
85 param_input_dtypes.append(param_dtypes)
86 param_input_numels.append(param_numels)
87 inp_split_sizes.extend(param_numels)
88 total_input_numel += sum(param_numels)
90 return AllGatherMetadata(
91 param_input_dtypes,
92 param_input_numels,
93 first_dtype,
94 inp_split_sizes,
95 total_input_numel
96 )
99@dataclass
100class AllGatherMetadata:
101 """Metadata describing the fused all-gather buffer layout.
103 Attributes:
104 param_input_dtypes: Per-parameter list of input tensor dtypes.
105 param_input_numels: Per-parameter list of input tensor element counts.
106 dtype: Uniform dtype of all inputs (used to allocate the fused buffer).
107 inp_split_sizes: Flat list of element counts for each input tensor across
108 all parameters, used by ``torch.split`` / ``split_with_sizes_copy`` to
109 slice the fused buffer back into per-parameter chunks.
110 total_input_numel: Total number of elements from all local shards (one rank's
111 contribution); the full all-gather output has ``total_input_numel * world_size``
112 elements.
113 hash_key: Computed in ``__post_init__`` for use as a cache key.
114 """
115 param_input_dtypes: list[list[torch.dtype]]
116 param_input_numels: list[list[int]]
117 dtype: torch.dtype
118 inp_split_sizes: list[int]
119 total_input_numel: int
120 hash_key: int = field(init=False)
122 def __post_init__(self):
123 self.hash_key = hash((
124 tuple(tuple(d) for d in self.param_input_dtypes),
125 tuple(tuple(n) for n in self.param_input_numels),
126 self.dtype,
127 tuple(self.inp_split_sizes),
128 self.total_input_numel
129 ))
132class AllGatherResult(NamedTuple):
133 """Result of a fused all-gather operation.
135 Attributes:
136 all_gather_output: The contiguous output buffer holding gathered data from all ranks.
137 metadata: The ``AllGatherMetadata`` used to interpret the buffer layout.
138 handle: Async work handle from ``dist.all_gather_into_tensor``; ``None`` when
139 the operation was synchronous or when ``shard_world_size == 1``.
140 """
141 all_gather_output: torch.Tensor
142 metadata: AllGatherMetadata
143 handle: Optional[Work]
146@dataclass
147class CommContext:
148 """Global communication context for pipelining fused gradient reduction.
150 For FSDP (shard-only), the reduce-scatter handle is stored in ``comm_handle``
151 and the next module's backward hook waits on it before issuing its own reduction.
153 For HSDP (shard + replicate), a two-phase pipeline is used:
154 Phase 1 (``wait_reduce_scatter_and_issue_all_reduce``): wait for
155 reduce-scatter, then issue one or more async all-reduces stored on
156 the owning ``HSDPParamGroup``.
157 Phase 2 (``wait_all_reduce_and_apply_grad``): wait for all-reduce and
158 write reduced gradients back.
160 This allows three-way overlap:
161 Layer N reduce_scatter ↔ Layer N-1 backward compute
162 Layer N all_reduce ↔ Layer N-1 reduce_scatter
163 """
164 comm_handle: Optional[Work] = None
165 all_reduce_handle: Optional[Work] = None
166 pre_param_group = None
167 # Param group whose all_reduce has been issued but grad not yet applied
168 all_reduce_param_group = None
171comm_ctx = CommContext()
174def get_comm_ctx():
175 """Return the global ``CommContext`` singleton."""
176 return comm_ctx
179@dataclass
180class ReplicateBucket:
181 """One fused all-reduce bucket sharing the same replicate process group."""
183 key: int
184 group: Any
185 group_size: int
186 param_indices: list[int]
187 flat_numel: int
188 buffer: Optional[torch.Tensor] = None
191@dataclass
192class PendingBucketAllReduce:
193 """One in-flight async all-reduce launched for a replicate bucket."""
195 bucket_key: int
196 handle: Any
199class AllGatherMetadataCache:
200 """Cache for ``AllGatherMetadata`` to avoid recomputation across iterations.
202 The cache key is derived from ``(id(param), param.version)`` tuples so that
203 it invalidates automatically when parameters are re-sharded or replaced.
204 """
205 _cache: dict[int, AllGatherMetadata] = {}
207 @classmethod
208 def get_metadata(cls, hsdp_params, fn):
209 """Return cached metadata or compute via *fn* and cache the result."""
210 param_key = tuple((id(p), getattr(p, 'version', 0)) for p in hsdp_params)
211 key = hash(param_key)
213 if key in cls._cache:
214 return cls._cache[key]
215 metadata = fn(hsdp_params)
216 cls._cache[key] = metadata
217 return metadata
220def all_gather_copy_in(all_gather_inputs, all_gather_output, inp_split_sizes, all_gather_input_numel, rank):
221 """Copy per-parameter local shards into the fused all-gather input buffer.
223 The fused output buffer has shape ``(total_input_numel * world_size,)``. Each rank
224 writes its local shards into the slice ``[input_numel * rank : input_numel * (rank+1)]``
225 using ``torch._foreach_copy_`` for efficient batched copy.
227 Args:
228 all_gather_inputs: Flat list of local shard tensors from all parameters.
229 all_gather_output: The pre-allocated fused output buffer.
230 inp_split_sizes: Element counts for splitting the rank-local slice.
231 all_gather_input_numel: Total elements for one rank's local shards.
232 rank: This rank's index within the shard process group.
234 Returns:
235 Tuple of (rank-local input slice, full output buffer).
236 """
237 all_gather_input = all_gather_output.narrow(0, all_gather_input_numel * rank, all_gather_input_numel)
238 foreach_copy_dsts = torch.split(all_gather_input, inp_split_sizes)
239 with torch.no_grad():
240 # pylint: disable=W0212
241 torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs)
242 return all_gather_input, all_gather_output
245def reduce_scatter_copy_in(
246 hsdp_params: List[TorchHSDPParamV2],
247 unsharded_grads: List[torch.Tensor],
248 reduce_scatter_input: torch.Tensor,
249 world_size: int,
250) -> None:
251 """Pack unsharded gradients into the fused reduce-scatter input buffer.
253 Uses ``torch._chunk_cat`` to interleave chunks from each gradient tensor so that
254 the buffer layout matches what ``dist.reduce_scatter_tensor`` expects: the buffer
255 is viewed as ``(world_size, total_numel // world_size)`` where row *i* contains
256 the slice destined for rank *i* after reduction.
258 Args:
259 hsdp_params: Parameters whose layout determines the pack plan per gradient.
260 unsharded_grads: Full (unsharded) gradients from all parameters.
261 reduce_scatter_input: Pre-allocated flat buffer of size ``sum(g.numel() for g in unsharded_grads)``.
262 world_size: Number of ranks in the shard process group.
263 """
264 if len(hsdp_params) != len(unsharded_grads):
265 raise AssertionError(
266 "reduce_scatter_copy_in expects one hsdp_param per unsharded_grad, but got "
267 f"{len(hsdp_params)} params and {len(unsharded_grads)} grads"
268 )
269 packed_rows = reduce_scatter_input.view(world_size, -1)
270 col_offset = 0
271 with torch.no_grad():
272 for hsdp_param, grad in zip(hsdp_params, unsharded_grads):
273 grad = grad.contiguous()
274 plan = build_rs_plan(hsdp_param, grad, world_size)
275 packed_grad = pack_for_reduce_scatter(grad, plan)
276 next_col_offset = col_offset + packed_grad.size(1)
277 packed_rows[:, col_offset:next_col_offset].copy_(packed_grad)
278 col_offset = next_col_offset
279 if col_offset != packed_rows.size(1):
280 raise AssertionError(
281 "reduce_scatter_copy_in packed an unexpected number of elements: "
282 f"{col_offset} != {packed_rows.size(1)}"
283 )
286class HSDPParamGroup:
287 """Groups all HSDP parameters within a module for fused collective communication.
289 Instead of issuing per-parameter all-gather (forward) and reduce-scatter (backward),
290 this class packs all parameter shards into a single contiguous buffer and performs one
291 fused collective, reducing NCCL/HCCL kernel launch overhead.
293 Lifecycle within one training iteration:
294 1. **Forward** — ``unshard()`` → ``foreach_all_gather()`` packs local shards into
295 ``ag_output`` and issues a single ``all_gather_into_tensor``.
296 2. **Forward (wait)** — ``wait_for_unshard()`` → ``foreach_all_gather_copy_out()``
297 waits on the handle and scatters gathered data back to per-parameter buffers.
298 3. **Backward** — ``foreach_reduce()`` packs unsharded gradients, issues fused
299 ``reduce_scatter_tensor`` (+ optional ``all_reduce`` for HSDP replicate dim),
300 and stores the handle in ``CommContext`` for pipelined overlap.
301 4. **Backward (apply)** — ``apply_fusion_reduced_grad()`` waits on the handle and
302 writes reduced gradient slices back to each parameter's ``.grad`` or ``.main_grad``.
304 Args:
305 hsdp_params: List of ``TorchHSDPParamV2`` belonging to this module.
306 mesh_info: Mesh info providing shard/replicate process groups.
307 device: Target device for buffer allocation.
308 mp_policy: Mixed-precision policy controlling reduce dtype and grad dtype.
309 """
311 def __init__(
312 self,
313 hsdp_params,
314 mesh_info: FSDPMeshInfo,
315 device: Optional[torch.device] = None,
316 mp_policy: Optional[MixedPrecisionPolicy] = None,
317 enable_zero_copy: bool = True,
318 ):
319 self.mesh_info = mesh_info
320 self.device = device
321 self.hsdp_params = hsdp_params
322 if isinstance(self.mesh_info, (FSDPMeshInfo, HSDPMeshInfo)):
323 self.shard_rank = self.mesh_info.shard_mesh_rank
324 self.shard_world_size = self.mesh_info.shard_mesh_size
325 else:
326 self.shard_rank = 0
327 self.shard_world_size = 1
328 self.shard_group = self.mesh_info.shard_process_group
329 self.replicate_group = None
330 if isinstance(self.mesh_info, (HSDPMeshInfo, DDPMeshInfo)):
331 self.replicate_group = self.mesh_info.replicate_process_group
332 elif isinstance(self.mesh_info, FSDPMeshInfo):
333 self.replicate_group = self._infer_layout_replicate_group()
334 self.device = device
335 self._all_gather_output = torch.empty(0, device=self.device)
336 self.ag_output = None # Fused all-gather output buffer, lazily allocated
337 self.metadata_cache = None
338 self.mp_policy = mp_policy
339 self.enable_zero_copy = enable_zero_copy
340 self._result = None # Pending AllGatherResult from async all-gather
341 self._reduce_output = None # Fused reduce-scatter output, consumed by apply_fusion_reduced_grad
342 self._reduce_op = None # Reduce op saved from foreach_reduce for use in apply_fusion_reduced_grad
343 self._needs_avg_div = False # Whether AVG was split into SUM + deferred div
344 self._reduce_hsdp_params = None
345 self._active_replicate_buckets: dict[int, ReplicateBucket] = {}
346 self._active_param_flat_offsets: list[int] = []
347 self._pending_all_reduce_handles: list[PendingBucketAllReduce] = []
348 self._init_mp_dtypes()
349 self._flat_param_buffer = None # Contiguous buffer holding all params' sharded data
350 self._flat_cast_buffer = None # Cast buffer for mixed precision (param_dtype)
351 if self.enable_zero_copy:
352 self._init_flat_param_buffer()
353 self.gradient_scaling_factor = None
355 def _infer_layout_replicate_group(self):
356 """Infer a compatibility all-reduce group from params' final DTensor layout when mesh_info has none.
358 DTENSOR_UNIFIED parameters may still carry replicate axes from the original
359 DTensor layout, for example a ``(tp, ep)`` mesh where ``ep`` is replicate-only.
360 The non-fused path derives this group from each param's layout-driven
361 ``unsharded_group_info``. ``comm_fusion`` now buckets by those groups, so
362 this helper only preserves the historical ``self.replicate_group`` field
363 for compatibility with simpler single-group paths.
364 """
365 replicate_groups = []
366 for hsdp_param in self.hsdp_params:
367 group_info = getattr(hsdp_param, "unsharded_group_info", None)
368 group = getattr(group_info, "group", None)
369 if group is None or getattr(hsdp_param, "replicate_world_size", 1) <= 1:
370 continue
371 replicate_groups.append((group, getattr(hsdp_param, "_param_fqn", "<unknown>")))
373 if not replicate_groups:
374 return None
376 ref_group, _ = replicate_groups[0]
377 return ref_group
379 def _build_active_replicate_buckets(self, hsdp_params):
380 """Group active params by their layout-driven replicate all-reduce group."""
381 buckets: dict[int, ReplicateBucket] = {}
382 for idx, hsdp_param in enumerate(hsdp_params):
383 group_info = getattr(hsdp_param, "unsharded_group_info", None)
384 group = getattr(group_info, "group", None)
385 group_size = getattr(
386 group_info,
387 "rank_size",
388 getattr(hsdp_param, "replicate_world_size", 1),
389 )
390 if not isinstance(group_size, int):
391 fallback_group_size = getattr(hsdp_param, "replicate_world_size", 1)
392 group_size = fallback_group_size if isinstance(fallback_group_size, int) else 1
393 if group is None or group_size <= 1:
394 continue
396 key = id(group)
397 if key not in buckets:
398 buckets[key] = ReplicateBucket(
399 key=key,
400 group=group,
401 group_size=group_size,
402 param_indices=[],
403 flat_numel=0,
404 )
405 buckets[key].param_indices.append(idx)
406 buckets[key].flat_numel += hsdp_param.sharded_size.numel()
407 return buckets
409 def _allocate_bucket_buffers_if_needed(self, device, dtype):
410 """Allocate or resize per-bucket temporary all-reduce buffers."""
411 for bucket in self._active_replicate_buckets.values():
412 if bucket.flat_numel == 0:
413 continue
414 needs_new_buffer = (
415 bucket.buffer is None
416 or bucket.buffer.numel() != bucket.flat_numel
417 or bucket.buffer.device != device
418 or bucket.buffer.dtype != dtype
419 )
420 if needs_new_buffer:
421 bucket.buffer = torch.empty(bucket.flat_numel, device=device, dtype=dtype)
423 def _pack_bucket_from_reduce_output(self, bucket: ReplicateBucket) -> torch.Tensor:
424 """Pack one replicate bucket's scattered shards into a contiguous all-reduce buffer."""
425 if bucket.buffer is None:
426 raise AssertionError("Bucket buffer must be allocated before packing from reduce output")
427 if self._reduce_output is None or self._reduce_hsdp_params is None:
428 raise AssertionError("Bucket packing requires an active fused reduce output")
429 dst_offset = 0
430 for idx in bucket.param_indices:
431 hsdp_param = self._reduce_hsdp_params[idx]
432 src_offset = self._active_param_flat_offsets[idx]
433 numel = hsdp_param.sharded_size.numel()
434 bucket.buffer.narrow(0, dst_offset, numel).copy_(
435 self._reduce_output.narrow(0, src_offset, numel)
436 )
437 dst_offset += numel
438 return bucket.buffer
440 def _unpack_bucket_to_reduce_output(self, bucket: ReplicateBucket) -> None:
441 """Write one bucket's post-all-reduce data back into the fused reduce output."""
442 if bucket.buffer is None:
443 raise AssertionError("Bucket buffer must exist before unpacking to reduce output")
444 if self._reduce_output is None or self._reduce_hsdp_params is None:
445 raise AssertionError("Bucket unpack requires an active fused reduce output")
446 src_offset = 0
447 for idx in bucket.param_indices:
448 hsdp_param = self._reduce_hsdp_params[idx]
449 dst_offset = self._active_param_flat_offsets[idx]
450 numel = hsdp_param.sharded_size.numel()
451 self._reduce_output.narrow(0, dst_offset, numel).copy_(
452 bucket.buffer.narrow(0, src_offset, numel)
453 )
454 src_offset += numel
456 def _init_flat_param_buffer(self):
457 """Initialize a contiguous flat buffer and rebase all params' sharded data into it.
459 This enables zero-copy all-gather by making all local shards contiguous in memory,
460 so they can be passed directly to ``all_gather_into_tensor`` without ``foreach_copy_``.
461 When mixed-precision casting is needed, a separate cast buffer is also allocated.
462 """
463 if self.shard_world_size <= 1:
464 return
465 if len(self.hsdp_params) == 0:
466 return
467 if any(p.offload_to_cpu or p.sharded_param.device.type == "meta" for p in self.hsdp_params):
468 return
470 total_numel = sum(p._sharded_param_data.numel() for p in self.hsdp_params)
471 orig_dtype = self.hsdp_params[0]._sharded_param_data.dtype
472 flat_buffer = torch.empty(total_numel, dtype=orig_dtype, device=self.device)
474 offset = 0
475 for hsdp_param in self.hsdp_params:
476 numel = hsdp_param._sharded_param_data.numel()
477 flat_slice = flat_buffer.narrow(0, offset, numel)
478 flat_slice.copy_(hsdp_param._sharded_param_data)
479 # Rebase _sharded_param_data to be a view into the flat buffer
480 hsdp_param._sharded_param_data = flat_slice
481 # Rebase DTensor's local tensor so optimizer in-place updates write to flat buffer
482 new_local = flat_slice.view(hsdp_param.sharded_size)
483 req_grad = hsdp_param.sharded_param.requires_grad
484 hsdp_param.sharded_param._local_tensor = new_local
485 hsdp_param.sharded_param.data = new_local
486 if req_grad:
487 new_local.requires_grad_(True)
488 hsdp_param.sharded_param.requires_grad_(True)
489 offset += numel
491 self._flat_param_buffer = flat_buffer
493 # Allocate cast buffer for mixed precision if needed
494 has_param_dtype = any(p.param_dtype is not None for p in self.hsdp_params)
495 if has_param_dtype:
496 cast_dtype = next(p.param_dtype for p in self.hsdp_params if p.param_dtype is not None)
497 self._flat_cast_buffer = torch.empty(total_numel, dtype=cast_dtype, device=self.device)
499 def _is_flat_buffer_valid(self):
500 """Check if the flat buffer is still backing the params' sharded data.
502 The flat buffer becomes invalid after ``load_state_dict`` triggers
503 ``reset_sharded_param``, which re-assigns ``_sharded_param_data``.
504 """
505 if self._flat_param_buffer is None or len(self.hsdp_params) == 0:
506 return False
507 return self.hsdp_params[0]._sharded_param_data.data_ptr() == self._flat_param_buffer.data_ptr()
509 def unshard(self, async_op: bool = False):
510 """Trigger fused all-gather to reconstruct full parameters from shards.
512 If a prefetch has already been issued (``_result is not None``), this is a no-op.
513 For ``shard_world_size == 1`` (no sharding), skips the collective entirely.
515 Args:
516 async_op: If True, the all-gather runs asynchronously and must be
517 completed later via ``wait_for_unshard()``.
518 """
519 # Already prefetched — skip
520 if self._result is not None:
521 return
522 if self.shard_world_size == 1:
523 self._result = AllGatherResult(self._all_gather_output, None, None)
524 return
525 self.foreach_all_gather(async_op=async_op)
527 def _init_mp_dtypes(self):
528 """Initialize and validate mixed-precision dtypes across all trainable parameters.
530 All trainable parameters in the group must have a uniform ``orig_dtype`` and
531 ``reduce_dtype``; heterogeneous dtypes would cause incorrect buffer slicing.
532 """
533 for hsdp_param in self.hsdp_params:
534 hsdp_param.init_dtype_attrs(self.mp_policy)
535 trainable_params: list[TorchHSDPParamV2] = [
536 p for p in self.hsdp_params if p.sharded_param.requires_grad
537 ]
538 orig_dtypes = {p.orig_dtype for p in trainable_params}
539 reduce_dtypes = {p.reduce_dtype for p in trainable_params}
540 if len(trainable_params) > 0 and len(orig_dtypes) != 1:
541 raise AssertionError(
542 f"hsdp expects uniform original parameter dtype but got {orig_dtypes}"
543 )
544 self._orig_dtype = next(iter(orig_dtypes)) if trainable_params else None
545 if len(trainable_params) > 0 and len(reduce_dtypes) != 1:
546 raise AssertionError(
547 f"hsdp expects uniform reduce dtype but got {reduce_dtypes}"
548 )
549 self._reduce_dtype = next(iter(reduce_dtypes)) if trainable_params else None
551 def wait_for_unshard(self):
552 """Wait for the async all-gather to complete and scatter data to per-parameter buffers.
554 For ``shard_world_size == 1``, simply copies the local shard as the full parameter.
555 Otherwise, calls ``foreach_all_gather_copy_out`` to split the fused buffer and
556 write each parameter's all-gather output. Finally, initializes unsharded parameters.
557 """
558 if self._result is None:
559 return
560 if self.shard_world_size == 1:
561 for hsdp_param in self.hsdp_params:
562 all_gather_input = hsdp_param.all_gather_inputs[0]
563 hsdp_param.init_all_gather_outputs(
564 [all_gather_input.numel()],
565 [all_gather_input.dtype],
566 self.shard_world_size,
567 self.device
568 )
569 hsdp_param.alloc_all_gather_outputs()
570 # pylint: disable=W0212
571 with torch.autograd._unsafe_preserve_version_counter(hsdp_param.all_gather_outputs[0]):
572 # pylint: disable=W0212
573 hsdp_param.all_gather_outputs[0].copy_(all_gather_input)
574 else:
575 self.foreach_all_gather_copy_out()
576 for hsdp_param in self.hsdp_params:
577 hsdp_param.init_unsharded_param()
578 hsdp_param.to_unsharded()
580 def alloc_all_gather_output(self, total_output_numel):
581 """Resize the fused all-gather buffer storage to fit ``total_output_numel`` elements.
583 Uses ``untyped_storage().resize_()`` to avoid reallocating the tensor object,
584 enabling storage reuse across iterations.
585 """
586 storage = self.ag_output.untyped_storage()
587 expected_size = total_output_numel * self.ag_output.itemsize
588 if storage.size() != expected_size:
589 storage.resize_(expected_size)
591 def free_all_gather_output(self):
592 """Release device memory of the fused all-gather buffer by resizing storage to 0."""
593 storage = self.ag_output.untyped_storage()
594 if storage.size() != 0:
595 storage.resize_(0)
597 @torch.no_grad()
598 def foreach_all_gather(self, async_op=False):
599 """Perform a fused all-gather for all parameters in the group.
601 When a flat parameter buffer is available (see ``_init_flat_param_buffer``),
602 the local shards are already contiguous and can be passed directly to
603 ``all_gather_into_tensor`` without any copy-in. Otherwise falls back to
604 the ``all_gather_copy_in`` path.
606 Args:
607 async_op: If True, the collective runs asynchronously.
608 """
609 if self.metadata_cache is None:
610 self.metadata_cache = AllGatherMetadataCache()
611 # pylint: disable=W0108
612 metadata = self.metadata_cache.get_metadata(self.hsdp_params, lambda p: get_all_gather_metadata(p))
613 if metadata.total_input_numel == 0:
614 return
615 world_size, rank = self.shard_group.size(), self.shard_group.rank()
616 total_output_numel = metadata.total_input_numel * world_size
617 if self.ag_output is None:
618 self.ag_output = torch.empty(size=(total_output_numel,),
619 dtype=metadata.dtype, device=self.device)
620 else:
621 self.alloc_all_gather_output(total_output_numel)
623 if self.enable_zero_copy and not self._is_flat_buffer_valid():
624 self._init_flat_param_buffer()
625 use_flat_buffer = self.enable_zero_copy and self._flat_param_buffer is not None
626 if use_flat_buffer:
627 # Zero-copy path: flat buffer already holds contiguous shard data
628 if self._flat_cast_buffer is not None:
629 # Mixed precision: single contiguous cast instead of N small copies
630 self._flat_cast_buffer.copy_(self._flat_param_buffer)
631 all_gather_input = self._flat_cast_buffer
632 else:
633 all_gather_input = self._flat_param_buffer
634 else:
635 # Fallback: collect inputs and copy into the rank-local slice of ag_output
636 all_gather_inputs = []
637 for hsdp_param in self.hsdp_params:
638 all_gather_inputs.extend(hsdp_param.all_gather_inputs)
639 if len(all_gather_inputs) == 0:
640 return
641 all_gather_input, _ = all_gather_copy_in(
642 all_gather_inputs,
643 self.ag_output,
644 metadata.inp_split_sizes,
645 metadata.total_input_numel,
646 rank
647 )
648 del all_gather_inputs # Free references to individual shard tensors
650 handle = dist.all_gather_into_tensor(self.ag_output, all_gather_input, self.shard_group, async_op)
651 self._result = AllGatherResult(self.ag_output, metadata, handle)
653 @torch.no_grad()
654 def foreach_all_gather_copy_out(self):
655 """Wait for the fused all-gather and scatter results back to per-parameter buffers.
657 After the collective completes, the fused output is viewed as ``(world_size, -1)``
658 and split along dim=1 according to ``inp_split_sizes``. Each slice is copied into
659 the corresponding parameter's ``all_gather_outputs`` buffer using
660 ``split_with_sizes_copy`` for zero-extra-allocation copy-out.
662 Version counters are preserved via ``_unsafe_preserve_version_counter`` to avoid
663 triggering autograd version checks on parameter tensors that alias these buffers.
664 """
665 (ag_output, metadata, _) = self._result
666 if self._result.handle is not None:
667 self._result.handle.wait()
668 device = ag_output.device
669 world_size = self.shard_group.size()
670 split_with_sizes_out = []
671 for input_numels, input_dtypes, hsdp_param in zip(
672 metadata.param_input_numels, metadata.param_input_dtypes, self.hsdp_params
673 ):
674 hsdp_param.init_all_gather_outputs(input_numels, input_dtypes, world_size, device)
675 hsdp_param.alloc_all_gather_outputs()
676 split_with_sizes_out.extend(hsdp_param.all_gather_outputs)
677 ag_output = ag_output.view(world_size, -1)
678 out = [t.view(world_size, -1) for t in split_with_sizes_out]
679 non_inference_outs = [o for o in out if not o.is_inference()]
680 if len(non_inference_outs) > 0:
681 # Older torch variants only accept one tensor per context manager.
682 # Preserve all version counters explicitly for cross-version compatibility.
683 # pylint: disable=W0212
684 with ExitStack() as stack:
685 for tensor in non_inference_outs:
686 stack.enter_context(torch.autograd._unsafe_preserve_version_counter(tensor))
687 torch.split_with_sizes_copy(ag_output, metadata.inp_split_sizes, dim=1, out=out)
688 else:
689 torch.split_with_sizes_copy(ag_output, metadata.inp_split_sizes, dim=1, out=out)
690 self._result = None
691 self.free_all_gather_output() # Immediately release fused buffer memory
693 @torch.no_grad()
694 def foreach_reduce(
695 self,
696 reduce_scatter_reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
697 async_op: bool = True,
698 ) -> Optional[torch.Tensor]:
699 """Perform fused gradient reduction (reduce-scatter + optional all-reduce).
701 Collects unsharded gradients from all parameters, packs them into a single
702 contiguous buffer, and issues one ``reduce_scatter_tensor``. For HSDP (2D mesh),
703 a follow-up ``all_reduce`` across the replicate dimension is also performed.
705 When ``async_op=True``, the communication handle is stored in the global
706 ``CommContext`` so that the next module's backward hook can overlap computation
707 with this reduction. The actual gradient write-back is deferred to
708 ``apply_fusion_reduced_grad()``.
710 Args:
711 reduce_scatter_reduce_op: Reduction operator (default: AVG).
712 async_op: If True, run collectives asynchronously for compute-comm overlap.
713 """
714 # Collect unsharded gradients (from accumulated grad or .grad)
715 hsdp_params: List[TorchHSDPParamV2] = []
716 unsharded_grads: List[torch.Tensor] = []
717 for hsdp_param in self.hsdp_params:
718 if not hasattr(hsdp_param, '_unsharded_param'):
719 continue
720 if hsdp_param.unsharded_accumulated_grad is not None:
721 hsdp_params.append(hsdp_param)
722 unsharded_grads.append(hsdp_param.unsharded_accumulated_grad_data)
723 elif hsdp_param._unsharded_param.grad is not None: # pylint: disable=W0212
724 hsdp_params.append(hsdp_param)
725 unsharded_grads.append(hsdp_param.unsharded_grad_data)
726 if not hsdp_params:
727 return
728 grad_dtypes = {g.dtype for g in unsharded_grads}
729 if len(grad_dtypes) != 1:
730 raise ValueError(
731 f"FSDP reduce-scatter expects uniform grad dtype but got {grad_dtypes}"
732 )
733 grad_dtype = unsharded_grads[0].dtype
734 reduce_dtype = self._reduce_dtype or grad_dtype
735 world_size = self.shard_group.size()
736 reduce_scatter_input_numel = sum(s.numel() for s in unsharded_grads)
737 reduce_scatter_output_numel = reduce_scatter_input_numel // world_size
738 device = unsharded_grads[0].device
739 # Pack all gradients into a contiguous buffer for fused reduce-scatter
740 reduce_scatter_input = torch.empty((reduce_scatter_input_numel,), dtype=reduce_dtype, device=device)
741 reduce_scatter_copy_in(hsdp_params, unsharded_grads, reduce_scatter_input, world_size)
742 unsharded_grads.clear() # Release references to full gradients
743 # Captured here, consumed once in _apply_reduced_grad after all collectives
744 # complete. Async paths cross method boundaries, so the field is unavoidable.
745 reduce_output = reduce_scatter_input.new_empty((reduce_scatter_output_numel,))
746 self._needs_avg_div = reduce_scatter_reduce_op == dist.ReduceOp.AVG
747 comm_op = dist.ReduceOp.SUM if self._needs_avg_div else reduce_scatter_reduce_op
748 self._reduce_op = comm_op
749 self._reduce_hsdp_params = hsdp_params
750 self._active_param_flat_offsets = []
751 flat_offset = 0
752 for hsdp_param in hsdp_params:
753 self._active_param_flat_offsets.append(flat_offset)
754 flat_offset += hsdp_param.sharded_size.numel()
755 self._active_replicate_buckets = self._build_active_replicate_buckets(hsdp_params)
756 self._allocate_bucket_buffers_if_needed(reduce_output.device, reduce_output.dtype)
757 self._pending_all_reduce_handles = []
758 apply_gradient_scaling_factor(reduce_scatter_input, self.gradient_scaling_factor)
759 rs_handle = dist.reduce_scatter_tensor(
760 output=reduce_output,
761 input=reduce_scatter_input,
762 group=self.shard_group,
763 op=comm_op,
764 async_op=async_op
765 )
766 comm_ctx.comm_handle = rs_handle
767 # Step 2 (HSDP only): All-reduce is deferred to apply_fusion_reduced_grad()
768 self._reduce_output = reduce_output
769 if async_op:
770 # Register this group for deferred grad application by the next backward hook
771 comm_ctx.pre_param_group = self
772 else:
773 self.apply_fusion_reduced_grad()
775 def wait_reduce_scatter_and_issue_all_reduce(self):
776 """Phase 1 of pipelined HSDP gradient reduction.
778 Waits for the async reduce-scatter to complete, then issues an async
779 all-reduce for each active replicate bucket. The bucket handles are
780 stored on this ``HSDPParamGroup`` so they can overlap with the next
781 layer's reduce-scatter (Phase 2 is deferred).
783 For FSDP (no replicate group), skips the all-reduce and directly
784 applies gradients since there is nothing further to pipeline.
785 """
786 if comm_ctx.comm_handle is not None:
787 comm_ctx.comm_handle.wait()
788 comm_ctx.comm_handle = None
789 # Deferred div for AVG: apply after RS completes, before AR
790 if self._needs_avg_div:
791 self._reduce_output.div_(self.shard_world_size)
792 if not self._active_replicate_buckets:
793 # No replicate group — no all-reduce needed, apply grads immediately
794 self._apply_reduced_grad()
795 return
797 self._pending_all_reduce_handles = []
798 for bucket in self._active_replicate_buckets.values():
799 packed = self._pack_bucket_from_reduce_output(bucket)
800 ar_handle = dist.all_reduce(
801 packed,
802 group=bucket.group,
803 op=self._reduce_op,
804 async_op=True,
805 )
806 self._pending_all_reduce_handles.append(
807 PendingBucketAllReduce(bucket_key=bucket.key, handle=ar_handle)
808 )
809 comm_ctx.all_reduce_param_group = self
811 def wait_all_reduce_and_apply_grad(self):
812 """Phase 2 of pipelined HSDP gradient reduction.
814 Waits for the async all-reduce issued in Phase 1 and writes reduced
815 gradients back to sharded parameters.
816 """
817 for pending in self._pending_all_reduce_handles:
818 bucket = self._active_replicate_buckets[pending.bucket_key]
819 pending.handle.wait()
820 if self._needs_avg_div:
821 bucket.buffer.div_(bucket.group_size)
822 self._unpack_bucket_to_reduce_output(bucket)
823 self._pending_all_reduce_handles = []
824 comm_ctx.all_reduce_handle = None
825 self._apply_reduced_grad()
827 def apply_fusion_reduced_grad(self):
828 """Full synchronous reduction path (used for final drain and sync mode).
830 Waits for reduce-scatter, performs synchronous all-reduce if needed,
831 and applies gradients — all in one call without pipelining.
832 """
833 if comm_ctx.comm_handle is not None:
834 comm_ctx.comm_handle.wait()
835 comm_ctx.comm_handle = None
836 # Deferred div for AVG after RS
837 if self._needs_avg_div:
838 self._reduce_output.div_(self.shard_world_size)
839 for bucket in self._active_replicate_buckets.values():
840 packed = self._pack_bucket_from_reduce_output(bucket)
841 dist.all_reduce(
842 packed,
843 group=bucket.group,
844 op=self._reduce_op,
845 )
846 # Deferred div for AVG after AR
847 if self._needs_avg_div:
848 packed.div_(bucket.group_size)
849 self._unpack_bucket_to_reduce_output(bucket)
850 self._apply_reduced_grad()
852 def _apply_reduced_grad(self):
853 """Write reduced gradients from ``_reduce_output`` back to sharded parameters.
855 Slices the fused ``_reduce_output`` buffer into per-parameter sharded gradients
856 using ``torch.as_strided`` (zero-copy view), then either accumulates into the
857 existing ``.grad`` / ``.main_grad`` or assigns a new DTensor gradient.
859 Handles:
860 - Mixed-precision: casts reduced gradient to ``_orig_dtype`` if needed.
861 - CPU offload: transfers gradient to CPU (``non_blocking`` when possible).
862 - Gradient accumulation: adds to existing grad when present.
863 - Memory cleanup: nulls out unsharded grad references to free memory.
864 """
865 flat_grad_offset = 0
866 if self._reduce_hsdp_params is None:
867 return
868 # All collectives have completed; scale once on the fused buffer right
869 # before slicing it into per-parameter sharded grads.
870 for hsdp_param in self._reduce_hsdp_params:
871 # Determine target gradient tensor (regular .grad or fp32 main_grad)
872 sharded_grad = None
873 if not self.mp_policy.apply_grad_on_fp32_main_grad:
874 sharded_grad = hsdp_param.sharded_param.grad
875 else:
876 if not hasattr(hsdp_param.sharded_param, "main_grad"):
877 hsdp_param.sharded_param.main_grad = None
878 sharded_grad = hsdp_param.sharded_param.main_grad
879 shard_size = hsdp_param.sharded_size
880 # Zero-copy view into the fused reduce output for this parameter's shard
881 new_sharded_grad = torch.as_strided(
882 self._reduce_output,
883 size=shard_size,
884 stride=hsdp_param.contiguous_sharded_stride,
885 storage_offset=flat_grad_offset,
886 )
887 # Cast to original dtype if reduce was done in a different precision
888 if not self.mp_policy.apply_grad_on_fp32_main_grad and new_sharded_grad.dtype != self._orig_dtype:
889 new_sharded_grad = new_sharded_grad.to(self._orig_dtype)
890 need_synchronize = False
891 if hsdp_param.offload_to_cpu:
892 non_blocking = hsdp_param.pin_memory and sharded_grad is None
893 new_sharded_grad = new_sharded_grad.to(
894 torch.device("cpu"), non_blocking=non_blocking
895 )
896 need_synchronize = True
897 # Accumulate or assign gradient
898 if sharded_grad is not None:
899 if not self.mp_policy.apply_grad_on_fp32_main_grad:
900 hsdp_param.sharded_param.grad._local_tensor += new_sharded_grad
901 else:
902 hsdp_param.sharded_param.main_grad._local_tensor += new_sharded_grad
903 hsdp_param.sharded_param.grad = None
904 else:
905 if not self.mp_policy.apply_grad_on_fp32_main_grad:
906 hsdp_param.sharded_param.grad = hsdp_param.to_sharded_dtensor(new_sharded_grad)
907 else:
908 hsdp_param.sharded_param.main_grad = hsdp_param.to_sharded_dtensor(new_sharded_grad)
909 hsdp_param.sharded_param.grad = None
910 flat_grad_offset += shard_size.numel()
911 # Release unsharded gradient references to free memory
912 if hsdp_param.unsharded_accumulated_grad is not None:
913 hsdp_param.unsharded_accumulated_grad = None
914 elif hsdp_param.unsharded_param.grad is not None:
915 hsdp_param.unsharded_param.grad = None
917 if need_synchronize:
918 if self.device.type == "npu":
919 torch.npu.current_stream().synchronize()
920 elif self.device.type == "cuda":
921 torch.cuda.current_stream().synchronize()
922 else:
923 raise NotImplementedError(f"Unsupported device type {self.device} for "
924 f"synchronization after CPU offload.")
925 self._reduce_output = None # Release fused reduce buffer
926 self._reduce_hsdp_params = None
927 self._active_param_flat_offsets = []
928 self._active_replicate_buckets = {}
929 self._pending_all_reduce_handles = []
932class AllReduceParamGroup:
933 """Groups HSDP parameters by replicate group for fused async all-reduce.
935 This class enables zero-copy fused all-reduce by:
936 1. Pre-allocating a contiguous buffer with 512-byte alignment
937 2. Having reduce_scatter write directly into aligned views of this buffer
938 3. Performing a single all_reduce on the entire buffer
939 4. Applying gradients directly from buffer views (with manual averaging)
941 Key design decisions for numerical correctness:
942 - Uses SUM instead of AVG for all_reduce to avoid padding zeros affecting the average
943 - Manually divides by replicate_world_size when applying gradients
944 - Padding regions are initialized to zero and don't affect SUM results
946 Attributes:
947 replicate_group: Process group for the replicate dimension.
948 hsdp_params: List of HSDP parameters in this group.
949 orig_dtypes: Original dtype for each parameter (for grad casting).
950 reduce_dtype: Uniform dtype for the fused buffer.
951 reduce_op: Original reduce op (AVG or SUM), used to determine final scaling.
952 replicate_world_size: Size of the replicate group.
953 fused_buffer: Pre-allocated contiguous buffer for all params.
954 param_offsets: Element offset for each param in the fused buffer.
955 param_numels: Number of elements for each param (excluding padding).
956 all_reduce_handle: Async work handle for the in-flight all_reduce.
957 """
959 ALIGNMENT_BYTES = 512 # 512-byte alignment requirement
961 @staticmethod
962 def _resolve_reduce_dtype(
963 reduce_dtype: Optional[torch.dtype],
964 hsdp_params: List["TorchHSDPParamV2"],
965 orig_dtypes: List[torch.dtype],
966 ) -> Optional[torch.dtype]:
967 """Resolve None reduce_dtype to match ``reduce_scatter_grad``'s ``dtype or grad.dtype``."""
968 if reduce_dtype is not None:
969 return reduce_dtype
970 for hsdp_param in hsdp_params:
971 if getattr(hsdp_param, "unsharded_accumulated_grad", None) is not None:
972 return hsdp_param.unsharded_accumulated_grad_data.dtype
973 unsharded_param = getattr(hsdp_param, "unsharded_param", None)
974 if unsharded_param is not None and getattr(unsharded_param, "grad", None) is not None:
975 return hsdp_param.unsharded_grad_data.dtype
976 return orig_dtypes[0] if orig_dtypes else None
978 def __init__(
979 self,
980 replicate_group: dist.ProcessGroup,
981 hsdp_params: List["TorchHSDPParamV2"],
982 orig_dtypes: List[torch.dtype],
983 reduce_dtype: torch.dtype,
984 reduce_op: dist.ReduceOp,
985 mp_policy: Optional["MixedPrecisionPolicy"] = None,
986 ):
987 self.replicate_group = replicate_group
988 self.hsdp_params = hsdp_params
989 self.orig_dtypes = orig_dtypes
990 self.reduce_dtype = self._resolve_reduce_dtype(reduce_dtype, hsdp_params, orig_dtypes)
991 self.reduce_op = reduce_op
992 self.mp_policy = mp_policy
993 self.replicate_world_size = replicate_group.size() if replicate_group else 1
995 # Fused buffer (lazily allocated)
996 self.fused_buffer: Optional[torch.Tensor] = None
997 # Element offsets in fused_buffer (accounting for padding)
998 self.param_offsets: List[int] = []
999 # Number of elements per param (without padding)
1000 self.param_numels: List[int] = []
1002 # Async communication handle
1003 self.all_reduce_handle: Optional[dist.Work] = None
1005 def compute_aligned_layout(self) -> int:
1006 """Compute buffer layout with 512-byte alignment for total buffer size only.
1008 Parameters are packed contiguously without per-param alignment.
1009 Padding is added only at the end of the buffer to make total size
1010 512-byte aligned.
1012 Returns:
1013 Total number of elements needed for the fused buffer.
1014 """
1015 self.param_offsets = []
1016 self.param_numels = []
1018 element_size = torch.tensor([], dtype=self.reduce_dtype).element_size()
1019 current_offset = 0
1021 for hsdp_param in self.hsdp_params:
1022 # Number of elements for this param's sharded gradient
1023 numel = hsdp_param.sharded_size.numel()
1024 self.param_numels.append(numel)
1025 self.param_offsets.append(current_offset)
1026 current_offset += numel
1028 # Total buffer size in bytes (packed, no per-param padding)
1029 total_bytes = current_offset * element_size
1031 # Align total buffer size to 512 bytes (padding at end only)
1032 aligned_total_bytes = (
1033 (total_bytes + self.ALIGNMENT_BYTES - 1) // self.ALIGNMENT_BYTES
1034 ) * self.ALIGNMENT_BYTES
1035 total_numel = aligned_total_bytes // element_size
1037 return total_numel
1039 def allocate_fused_buffer(self, device: torch.device) -> None:
1040 """Allocate the fused buffer with computed layout."""
1041 total_numel = self.compute_aligned_layout()
1042 self.fused_buffer = torch.empty(total_numel, dtype=self.reduce_dtype, device=device)
1043 # Initialize to zero (important for SUM correctness with padding)
1044 self.fused_buffer.zero_()
1046 def get_param_buffer_view(self, idx: int) -> torch.Tensor:
1047 """Get a view into the fused buffer for parameter at index idx.
1049 This view can be used as the output buffer for reduce_scatter,
1050 enabling zero-copy fusion.
1052 Args:
1053 idx: Index of the parameter in hsdp_params.
1055 Returns:
1056 A 1D tensor view of size param_numels[idx].
1057 """
1058 if self.fused_buffer is None:
1059 raise RuntimeError("Fused buffer not allocated. Call allocate_fused_buffer first.")
1061 offset = self.param_offsets[idx]
1062 numel = self.param_numels[idx]
1063 return self.fused_buffer.narrow(0, offset, numel)
1065 def get_param_grad_view(self, idx: int, target_shape: torch.Size) -> torch.Tensor:
1066 """Get a reshaped view of the reduced gradient for applying to parameter.
1068 Args:
1069 idx: Index of the parameter.
1070 target_shape: Target shape (sharded_size).
1072 Returns:
1073 A view of the reduced gradient with target_shape.
1074 """
1075 flat_view = self.get_param_buffer_view(idx)
1076 return flat_view.view(target_shape)
1078 def accumulate_existing_grads_to_buffer(self) -> None:
1079 """Accumulate existing sharded_param.grad/main_grad to fused_buffer.
1081 This is called before allreduce in gradient accumulation scenario,
1082 to ensure the previously accumulated gradients (from n-1 mini steps)
1083 are included in the allreduce operation.
1085 Handles:
1086 - Mixed-precision: uses reduce_dtype for consistency
1087 - main_grad vs grad: respects mp_policy.apply_grad_on_fp32_main_grad
1088 """
1089 if self.fused_buffer is None:
1090 return
1092 for idx, hsdp_param in enumerate(self.hsdp_params):
1093 # Get existing sharded grad
1094 existing_grad = None
1095 if self.mp_policy is not None and self.mp_policy.apply_grad_on_fp32_main_grad:
1096 if hasattr(hsdp_param.sharded_param, "main_grad"):
1097 existing_grad = hsdp_param.sharded_param.main_grad
1098 else:
1099 existing_grad = hsdp_param.sharded_param.grad
1101 if existing_grad is not None and not hsdp_param.accumulated_allreduced_grad:
1102 # Get DTensor's local_tensor
1103 from hyper_parallel.core.dtensor.dtensor import DTensor
1104 if isinstance(existing_grad, DTensor):
1105 existing_grad_local = existing_grad._local_tensor
1106 else:
1107 existing_grad_local = existing_grad
1109 # Get the corresponding view in fused_buffer
1110 buffer_view = self.get_param_buffer_view(idx)
1112 # Ensure dtype consistency (convert to reduce_dtype)
1113 if existing_grad_local.dtype != self.reduce_dtype:
1114 existing_grad_local = existing_grad_local.to(self.reduce_dtype)
1116 # Accumulate to fused_buffer
1117 buffer_view.add_(existing_grad_local.view_as(buffer_view))
1118 if self.mp_policy is not None and self.mp_policy.apply_grad_on_fp32_main_grad:
1119 if hasattr(hsdp_param.sharded_param, "main_grad"):
1120 hsdp_param.sharded_param.main_grad = None
1121 else:
1122 hsdp_param.sharded_param.grad = None
1124 def issue_async_allreduce(self) -> None:
1125 """Issue async all_reduce on the fused buffer.
1127 Uses SUM operation for numerical correctness with padding.
1128 If original op was AVG, scaling is done when applying gradients.
1129 """
1130 if self.fused_buffer is None:
1131 raise RuntimeError("Fused buffer not allocated.")
1133 # Always use SUM for correctness with padding regions
1134 # If original op was AVG, we divide by world_size when applying
1135 self.all_reduce_handle = dist.all_reduce(
1136 self.fused_buffer,
1137 op=dist.ReduceOp.SUM,
1138 group=self.replicate_group,
1139 async_op=True,
1140 )
1142 def wait_and_apply_grads(self) -> bool:
1143 """Wait for all_reduce to complete and apply gradients to parameters.
1145 Returns:
1146 True if CPU synchronization is needed (for offload params).
1147 """
1148 if self.all_reduce_handle is not None:
1149 self.all_reduce_handle.wait()
1150 self.all_reduce_handle = None
1152 need_synchronize = False
1154 for idx, hsdp_param in enumerate(self.hsdp_params):
1155 # Get the reduced gradient from fused buffer
1156 reduced_grad = self.get_param_grad_view(idx, hsdp_param.sharded_size)
1158 # Apply manual averaging if original op was AVG
1159 if self.reduce_op == dist.ReduceOp.AVG and self.replicate_world_size > 1:
1160 reduced_grad = reduced_grad / self.replicate_world_size
1161 # Apply to parameter (handles dtype cast, CPU offload, accumulation)
1162 need_synchronize = hsdp_param.apply_reduced_grad(
1163 reduced_grad, self.orig_dtypes[idx]
1164 ) or need_synchronize
1165 hsdp_param.accumulated_allreduced_grad = True
1167 # Release fused buffer
1168 self.fused_buffer = None
1170 return need_synchronize