Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / fully_shard / param_group.py: 92%
514 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15# Adapted from https://github.com/pytorch/pytorch/blob/release/2.6/torch/distributed/fsdp/_fully_shard
16# ==========================================================================
17"""Fused parameter communication for Torch fully_shard.
19Instead of issuing one collective per parameter, the managed parameters of a
20fully_shard unit are packed into contiguous buckets and communicated with a
21single collective each, which cuts kernel-launch overhead and improves
22bandwidth utilization.
24Bucketing: parameters are not fused into one global buffer. Each collective
25splits its parameters into buckets keyed by the attributes that must be uniform
26within one call -- ``(process group, dtype)`` -- because a collective cannot mix
27communication groups or element types:
29- ``AllGatherBucket``: one fused all-gather, keyed by (shard group, param dtype).
30 Owns the optional zero-copy ``flat_param_buffer`` that parameter shards are
31 rebased onto, so the gather reads parameter storage directly.
32- ``ReduceScatterBucket``: one fused reduce-scatter, keyed by (shard group,
33 reduce dtype). Packs gradients whose shard dim may differ per parameter.
34- ``AllReduceBucket``: the optional HSDP/replicate stage paired with one
35 reduce-scatter bucket; it takes over the completed RS output and reduces it
36 in place without per-parameter repacking.
38``ParamGroupCommCtx`` tracks which ``HSDPParamGroup`` owns each in-flight backward
39stage so the next module's backward can wait on it, giving the three-way
40overlap: layer N reduce-scatter with layer N-1 backward compute, and layer N
41all-reduce with layer N-1 reduce-scatter.
42"""
43from contextlib import ExitStack
44from dataclasses import dataclass
45from typing import List, Optional
47import torch
48import torch.distributed as dist
50from hyper_parallel.core.fully_shard.hsdp_utils import apply_gradient_scaling_factor
51from hyper_parallel.core.fully_shard.hsdp_scheduler import ParamGroupCommCtx
52from hyper_parallel.core.fully_shard.utils import DDPMeshInfo, FSDPMeshInfo
53from hyper_parallel.platform.torch.fully_shard.param import TorchHSDPParamV2
56@dataclass
57class AllGatherMetadata:
58 """Describe the rank-local layout of one fused all-gather bucket."""
60 param_input_dtypes: list[list[torch.dtype]]
61 param_input_numels: list[list[int]]
62 dtype: torch.dtype
63 inp_split_sizes: list[int]
64 total_input_numel: int
67@dataclass
68class AllGatherResult:
69 """Keep one all-gather input, output, and handle alive until copy-out."""
71 all_gather_input: Optional[torch.Tensor]
72 all_gather_output: Optional[torch.Tensor]
73 handle: Optional[dist.Work]
76@dataclass
77class AllGatherBucket:
78 """Own parameters and buffers sharing one all-gather group and dtype."""
80 hsdp_params: list[TorchHSDPParamV2]
81 shard_group: dist.ProcessGroup
82 shard_rank: int
83 shard_world_size: int
84 dtype: torch.dtype
85 metadata: AllGatherMetadata
86 all_gather_result: Optional[AllGatherResult] = None
87 flat_param_buffer: Optional[torch.Tensor] = None
89 def init_flat_param_buffer(self, device: Optional[torch.device]) -> None:
90 """Rebase this bucket's homogeneous shards into one persistent flat buffer.
92 Side effect: rebinds every parameter's ``_sharded_param_data`` and
93 ``sharded_param`` onto a view of the new buffer, so the original shard
94 storages are dropped. Leaves ``flat_param_buffer`` as ``None`` when the
95 bucket mixes storage dtypes or holds offloaded/meta parameters, which is
96 the signal that the zero-copy path does not apply.
97 """
98 storage_dtype = self.hsdp_params[0]._sharded_param_data.dtype
99 if any(
100 hsdp_param._sharded_param_data.dtype != storage_dtype
101 for hsdp_param in self.hsdp_params[1:]
102 ):
103 self.flat_param_buffer = None
104 return
105 if any(
106 hsdp_param.offload_to_cpu or hsdp_param.sharded_param.device.type == "meta"
107 for hsdp_param in self.hsdp_params
108 ):
109 self.flat_param_buffer = None
110 return
112 total_numel = sum(hsdp_param._sharded_param_data.numel() for hsdp_param in self.hsdp_params)
113 flat_param_buffer = torch.empty(total_numel, dtype=storage_dtype, device=device)
114 flat_offset = 0
115 for hsdp_param in self.hsdp_params:
116 param_numel = hsdp_param._sharded_param_data.numel()
117 flat_param_view = flat_param_buffer.narrow(0, flat_offset, param_numel)
118 flat_param_view.copy_(hsdp_param._sharded_param_data)
119 hsdp_param._sharded_param_data = flat_param_view
120 padded_local_param = flat_param_view.view(hsdp_param.padded_sharded_param_size)
121 local_param = padded_local_param.narrow(
122 hsdp_param.hsdp_placement.dim,
123 0,
124 hsdp_param.sharded_size[hsdp_param.hsdp_placement.dim],
125 )
126 requires_grad = hsdp_param.sharded_param.requires_grad
127 hsdp_param.sharded_param._local_tensor = local_param
128 hsdp_param.sharded_param.data = local_param
129 if requires_grad:
130 local_param.requires_grad_(True)
131 hsdp_param.sharded_param.requires_grad_(True)
132 flat_offset += param_numel
133 self.flat_param_buffer = flat_param_buffer
135 def is_flat_buffer_valid(self) -> bool:
136 """Return whether this bucket's parameter shards still use its flat storage.
138 Parameter storage can be replaced behind this bucket's back (optimizer
139 surgery, ``load_state_dict``), which silently invalidates the zero-copy
140 rebase; comparing storage pointers is what detects that.
141 """
142 if self.flat_param_buffer is None:
143 return False
144 flat_storage_ptr = self.flat_param_buffer.untyped_storage().data_ptr()
145 return all(
146 hsdp_param._sharded_param_data.untyped_storage().data_ptr() == flat_storage_ptr
147 for hsdp_param in self.hsdp_params
148 )
150 @torch.no_grad()
151 def copy_out(self) -> None:
152 """Wait this bucket's all-gather and copy it into stable unsharded buffers.
154 Blocks on the in-flight handle, then releases the bucket's all-gather
155 input/output references so the fused communication buffers can be freed
156 before the next bucket runs. No-op when the bucket has no pending result.
157 """
158 all_gather_result = self.all_gather_result
159 if all_gather_result is None or all_gather_result.all_gather_output is None:
160 return
161 if all_gather_result.handle is not None:
162 all_gather_result.handle.wait()
163 all_gather_result.handle = None
164 all_gather_output = all_gather_result.all_gather_output
165 metadata = self.metadata
166 split_with_sizes_out = []
167 for input_numels, input_dtypes, hsdp_param in zip(
168 metadata.param_input_numels,
169 metadata.param_input_dtypes,
170 self.hsdp_params,
171 ):
172 hsdp_param.init_unsharded_param_buffers(
173 input_numels,
174 input_dtypes,
175 self.shard_world_size,
176 all_gather_output.device,
177 )
178 hsdp_param.alloc_unsharded_param_buffers()
179 split_with_sizes_out.extend(hsdp_param.unsharded_param_buffers)
181 gathered_rows = all_gather_output.view(self.shard_world_size, -1)
182 output_rows = [
183 output_buffer.view(self.shard_world_size, -1)
184 for output_buffer in split_with_sizes_out
185 ]
186 non_inference_outputs = [output for output in output_rows if not output.is_inference()]
187 # PyTorch 2.6 accepts only one tensor per context manager. no_grad does
188 # not suppress version-counter bumps from the copy-out operations.
189 # pylint: disable=W0212
190 with ExitStack() as stack:
191 for output in non_inference_outputs:
192 stack.enter_context(torch.autograd._unsafe_preserve_version_counter(output))
193 if all(
194 hsdp_param.hsdp_placement.dim == 0
195 for hsdp_param in self.hsdp_params
196 ):
197 torch.split_with_sizes_copy(
198 gathered_rows,
199 metadata.inp_split_sizes,
200 dim=1,
201 out=output_rows,
202 )
203 else:
204 column_offset = 0
205 for input_numels, hsdp_param in zip(
206 metadata.param_input_numels,
207 self.hsdp_params,
208 ):
209 if hsdp_param.hsdp_placement.dim != 0 and len(input_numels) != 1:
210 raise NotImplementedError(
211 "Fused non-dim-0 all-gather expects one local shard tensor per parameter."
212 )
213 for input_numel, output_buffer in zip(
214 input_numels,
215 hsdp_param.unsharded_param_buffers,
216 ):
217 gathered_param = gathered_rows.narrow(1, column_offset, input_numel)
218 if hsdp_param.hsdp_placement.dim == 0:
219 output_buffer.view(self.shard_world_size, -1).copy_(gathered_param)
220 else:
221 packed_shape = list(hsdp_param.sharded_size)
222 packed_shape[0] *= self.shard_world_size
223 packed_param = gathered_param.contiguous().view(packed_shape)
224 param_chunks = torch.chunk(
225 packed_param,
226 self.shard_world_size,
227 dim=0,
228 )
229 torch.cat(
230 param_chunks,
231 dim=hsdp_param.hsdp_placement.dim,
232 out=output_buffer.view(hsdp_param._orig_size),
233 )
234 column_offset += input_numel
235 if column_offset != gathered_rows.size(1):
236 raise AssertionError(
237 "Fused all-gather copy-out consumed an unexpected number of elements: "
238 f"{column_offset} != {gathered_rows.size(1)}"
239 )
240 all_gather_result.all_gather_input = None
241 all_gather_result.all_gather_output = None
242 self.all_gather_result = None
245@dataclass
246class GradientBucketLayout:
247 """Describe the shared per-parameter layout of an RS-to-AR bucket chain."""
249 hsdp_params: list[TorchHSDPParamV2]
250 param_offsets: list[int]
251 param_numels: list[int]
252 total_numel: int
255@dataclass
256class ReduceScatterBucket:
257 """Own one fused reduce-scatter operation and its temporary buffers.
259 Bucket construction records only communication metadata and the source
260 gradients. Buffer preparation, gradient scaling, collective launch, and
261 final result placement are separate execution phases on ``HSDPParamGroup``.
263 The shard group plus reduce dtype is this bucket's identity across
264 micro-steps and keys its parked partial output.
265 """
267 layout: GradientBucketLayout
268 unsharded_grads: list[torch.Tensor]
269 shard_group: Optional[dist.ProcessGroup]
270 shard_world_size: int
271 dtype: torch.dtype
272 reduce_op: dist.ReduceOp
273 needs_avg_div: bool
274 reduce_scatter_input: Optional[torch.Tensor] = None
275 reduce_scatter_output: Optional[torch.Tensor] = None
276 handle: Optional[dist.Work] = None
278 @property
279 def hsdp_params(self) -> list[TorchHSDPParamV2]:
280 """Return parameters in their fused-buffer order."""
281 return self.layout.hsdp_params
283 @property
284 def param_offsets(self) -> list[int]:
285 """Return parameter offsets in the reduce-scatter output."""
286 return self.layout.param_offsets
288 @property
289 def bucket_key(self) -> tuple:
290 """Identity of this bucket's reduce-scatter fusion class."""
291 return (id(self.shard_group), self.dtype)
293 @property
294 def uses_collective(self) -> bool:
295 """Return whether this bucket needs an actual reduce-scatter call."""
296 return self.shard_group is not None and self.shard_world_size > 1
298 def move_reducescatter_output(self) -> torch.Tensor:
299 """Transfer exclusive ownership of the completed output to the next stage."""
300 if self.reduce_scatter_output is None:
301 raise RuntimeError("Reduce-scatter output has already been released.")
302 reduce_scatter_output = self.reduce_scatter_output
303 self.reduce_scatter_output = None
304 return reduce_scatter_output
307@dataclass
308class AllReduceBucket:
309 """Own the final stage of one fused RS-to-AR bucket chain.
311 The bucket shares its parameter layout with exactly one source
312 reduce-scatter bucket. Once RS completes, it takes over that bucket's whole
313 output tensor and all-reduces it in place without per-parameter repacking.
314 """
316 layout: GradientBucketLayout
317 source_reduce_scatter_bucket: ReduceScatterBucket
318 replicate_group: dist.ProcessGroup
319 replicate_world_size: int
320 dtype: torch.dtype
321 reduce_op: dist.ReduceOp
322 needs_avg_div: bool
323 all_reduce_output: Optional[torch.Tensor] = None
324 handle: Optional[dist.Work] = None
326 @property
327 def hsdp_params(self) -> list[TorchHSDPParamV2]:
328 """Return parameters in their fused-buffer order."""
329 return self.layout.hsdp_params
331 @property
332 def param_numels(self) -> list[int]:
333 """Return padded parameter sizes in the all-reduce output."""
334 return self.layout.param_numels
336 @property
337 def param_offsets(self) -> list[int]:
338 """Return parameter offsets in the all-reduce output."""
339 return self.layout.param_offsets
341 @property
342 def uses_collective(self) -> bool:
343 """Return whether this bucket needs an actual all-reduce call."""
344 return self.replicate_world_size > 1
347def get_all_gather_metadata(hsdp_params: list[TorchHSDPParamV2]) -> AllGatherMetadata:
348 """Build metadata for parameters with one all-gather communication dtype."""
349 param_input_dtypes = []
350 param_input_numels = []
351 inp_split_sizes = []
352 total_input_numel = 0
353 dtype = None
354 for hsdp_param in hsdp_params:
355 inputs = hsdp_param.all_gather_inputs
356 if dtype is None:
357 dtype = inputs[0].dtype
358 if any(tensor.dtype != dtype for tensor in inputs):
359 raise ValueError("All parameters in an all-gather bucket must have the same dtype.")
360 input_dtypes = [tensor.dtype for tensor in inputs]
361 input_numels = [tensor.numel() for tensor in inputs]
362 param_input_dtypes.append(input_dtypes)
363 param_input_numels.append(input_numels)
364 inp_split_sizes.extend(input_numels)
365 total_input_numel += sum(input_numels)
366 if dtype is None:
367 raise ValueError("Cannot build all-gather metadata for an empty parameter bucket.")
368 return AllGatherMetadata(
369 param_input_dtypes=param_input_dtypes,
370 param_input_numels=param_input_numels,
371 dtype=dtype,
372 inp_split_sizes=inp_split_sizes,
373 total_input_numel=total_input_numel,
374 )
377def all_gather_copy_in(
378 all_gather_inputs: list[torch.Tensor],
379 all_gather_output: torch.Tensor,
380 inp_split_sizes: list[int],
381 all_gather_input_numel: int,
382 rank: int,
383) -> tuple[torch.Tensor, torch.Tensor]:
384 """Copy local parameter shards into this rank's fused output slice."""
385 all_gather_input = all_gather_output.narrow(
386 0,
387 all_gather_input_numel * rank,
388 all_gather_input_numel,
389 )
390 copy_destinations = torch.split(all_gather_input, inp_split_sizes)
391 with torch.no_grad():
392 # pylint: disable=W0212
393 torch._foreach_copy_(copy_destinations, all_gather_inputs)
394 return all_gather_input, all_gather_output
397def reduce_scatter_copy_in(
398 hsdp_params: list[TorchHSDPParamV2],
399 unsharded_grads: list[torch.Tensor],
400 reduce_scatter_input: torch.Tensor,
401 world_size: int,
402) -> None:
403 """Pack gradients with mixed shard dimensions into a fused RS input."""
404 if len(hsdp_params) != len(unsharded_grads):
405 raise AssertionError(
406 "reduce_scatter_copy_in expects one hsdp_param per unsharded_grad, but got "
407 f"{len(hsdp_params)} params and {len(unsharded_grads)} grads"
408 )
409 packed_rows = reduce_scatter_input.view(world_size, -1)
410 column_offset = 0
411 with torch.no_grad():
412 for hsdp_param, unsharded_grad in zip(hsdp_params, unsharded_grads):
413 padded_shard_numel = hsdp_param.padded_sharded_param_size.numel()
414 packed_grad_slot = packed_rows.narrow(1, column_offset, padded_shard_numel)
415 shard_dim = hsdp_param.hsdp_placement.dim
416 if shard_dim == 0:
417 # pylint: disable=W0212
418 torch._chunk_cat(
419 [unsharded_grad],
420 dim=0,
421 num_chunks=world_size,
422 out=packed_grad_slot,
423 )
424 else:
425 grad_chunks = torch.chunk(unsharded_grad, world_size, dim=shard_dim)
426 packed_grad = torch.cat(grad_chunks, dim=0).contiguous().view(world_size, -1)
427 packed_grad_slot.copy_(packed_grad)
428 column_offset += padded_shard_numel
429 if column_offset != packed_rows.size(1):
430 raise AssertionError(
431 "reduce_scatter_copy_in packed an unexpected number of elements: "
432 f"{column_offset} != {packed_rows.size(1)}"
433 )
436class HSDPParamGroup:
437 """Fuse communication for the managed parameters in one fully_shard unit.
439 Packs parameter shards into contiguous buckets so communication is issued
440 once per compatible route and dtype instead of once per parameter.
442 Lifecycle within one training iteration:
443 1. Forward -- ``unshard()`` builds the all-gather buckets on first use
444 and issues one ``all_gather_into_tensor`` per bucket.
445 2. Forward (wait) -- ``wait_for_unshard()`` waits each bucket and copies
446 the gathered data out into the per-parameter unsharded buffers.
447 3. Backward -- ``foreach_reduce()`` packs the unsharded gradients and
448 issues one fused ``reduce_scatter_tensor`` per bucket, then parks
449 itself on ``ParamGroupCommCtx`` so a later module's backward waits for it.
450 4. Backward (finalize) -- ``wait_reduce_scatter_and_issue_all_reduce()``
451 waits the RS buckets and launches the HSDP all-reduces;
452 ``wait_all_reduce_and_save_grad()`` waits those and exposes each
453 parameter's reduced-gradient view.
455 Steps 3 and 4 are deliberately split across modules: that is what allows one
456 module's communication to overlap the next module's backward compute. The
457 root backward hook drains whatever is still parked when backward ends.
459 Gradient accumulation is held at bucket granularity in
460 ``reduce_partial_outputs``, keyed by the reduce-scatter bucket key. A
461 micro-step with ``requires_all_reduce=False`` keeps its whole reduce-scatter
462 output there and skips the all-reduce entirely; the next synchronizing
463 micro-step adds it back with a single whole-buffer ``add_`` and lets the
464 fused all-reduce see the total. The dict lives on the group rather than on
465 the bucket because ``foreach_reduce`` rebuilds the buckets every micro-step,
466 and holds whole buffers rather than per-parameter views so accumulation
467 costs one kernel per bucket instead of one per parameter.
468 """
470 def __init__(
471 self,
472 hsdp_params: list[TorchHSDPParamV2],
473 device: Optional[torch.device] = None,
474 enable_zero_copy: bool = True,
475 comm_ctx: Optional[ParamGroupCommCtx] = None,
476 ) -> None:
477 self.device = device
478 self.hsdp_params = hsdp_params
479 self.comm_ctx = comm_ctx or ParamGroupCommCtx()
480 self.enable_zero_copy = enable_zero_copy
481 self.gradient_scaling_factor = None
482 self.requires_all_reduce = True
483 self.all_gather_buckets: list[AllGatherBucket] = []
484 self.reduce_scatter_buckets: list[ReduceScatterBucket] = []
485 self.all_reduce_buckets: list[AllReduceBucket] = []
486 self.reduce_partial_outputs: dict[tuple, torch.Tensor] = {}
488 def _init_all_gather_buckets(self) -> None:
489 """Build ordered all-gather buckets from each parameter's communication facts."""
490 params_by_bucket = {}
491 bucket_groups = {}
492 for hsdp_param in self.hsdp_params:
493 if hsdp_param.shard_world_size <= 1:
494 continue
495 if not isinstance(hsdp_param.mesh_info, FSDPMeshInfo):
496 raise ValueError(
497 f"Fused all-gather expects FSDPMeshInfo, got {type(hsdp_param.mesh_info)}"
498 )
499 shard_group = hsdp_param.mesh_info.shard_process_group
500 communication_dtype = hsdp_param.param_dtype or hsdp_param.orig_dtype
501 bucket_key = (id(shard_group), communication_dtype)
502 if bucket_key not in params_by_bucket:
503 params_by_bucket[bucket_key] = []
504 bucket_groups[bucket_key] = shard_group
505 params_by_bucket[bucket_key].append(hsdp_param)
507 self.all_gather_buckets = []
508 for bucket_key, hsdp_params in params_by_bucket.items():
509 shard_group = bucket_groups[bucket_key]
510 param_input_numels = [
511 [hsdp_param._sharded_param_data.numel()] for hsdp_param in hsdp_params
512 ]
513 inp_split_sizes = [input_numels[0] for input_numels in param_input_numels]
514 allgather_metadata = AllGatherMetadata(
515 param_input_dtypes=[[bucket_key[1]] for _ in hsdp_params],
516 param_input_numels=param_input_numels,
517 dtype=bucket_key[1],
518 inp_split_sizes=inp_split_sizes,
519 total_input_numel=sum(inp_split_sizes),
520 )
521 self.all_gather_buckets.append(
522 AllGatherBucket(
523 hsdp_params=hsdp_params,
524 shard_group=shard_group,
525 shard_rank=shard_group.rank(),
526 shard_world_size=shard_group.size(),
527 dtype=bucket_key[1],
528 metadata=allgather_metadata,
529 )
530 )
532 def unshard(self, async_op: bool = False) -> None:
533 """Launch fused all-gathers for every communication bucket."""
534 if self.all_gather_buckets and any(
535 bucket.all_gather_result is not None for bucket in self.all_gather_buckets
536 ):
537 # already triggered by pre_module.prefetch(), return directly.
538 return
539 self.foreach_all_gather(async_op)
541 @torch.no_grad()
542 def foreach_all_gather(self, async_op: bool = False) -> None:
543 """Initialize ordered buckets and launch their all-gathers."""
544 if not self.all_gather_buckets:
545 self._init_all_gather_buckets()
546 for hsdp_param in self.hsdp_params:
547 if hsdp_param.shard_world_size <= 1:
548 # fast path, skip copy_in process.
549 hsdp_param.unshard(async_op)
550 for all_gather_bucket in self.all_gather_buckets:
551 if all_gather_bucket.all_gather_result is not None:
552 continue
553 if self.enable_zero_copy and not all_gather_bucket.is_flat_buffer_valid():
554 all_gather_bucket.init_flat_param_buffer(self.device)
556 metadata = all_gather_bucket.metadata
557 all_gather_output = torch.empty(
558 metadata.total_input_numel * all_gather_bucket.shard_world_size,
559 dtype=all_gather_bucket.dtype,
560 device=self.device,
561 )
562 if self.enable_zero_copy and all_gather_bucket.is_flat_buffer_valid():
563 if all_gather_bucket.flat_param_buffer.dtype == all_gather_bucket.dtype:
564 all_gather_input = all_gather_bucket.flat_param_buffer
565 else:
566 all_gather_input = all_gather_bucket.flat_param_buffer.to(all_gather_bucket.dtype)
567 else:
568 all_gather_inputs = []
569 for hsdp_param in all_gather_bucket.hsdp_params:
570 all_gather_inputs.extend(hsdp_param.all_gather_inputs)
571 all_gather_input, all_gather_output = all_gather_copy_in(
572 all_gather_inputs,
573 all_gather_output,
574 metadata.inp_split_sizes,
575 metadata.total_input_numel,
576 all_gather_bucket.shard_rank,
577 )
579 handle = dist.all_gather_into_tensor(
580 all_gather_output,
581 all_gather_input,
582 group=all_gather_bucket.shard_group,
583 async_op=async_op,
584 )
585 all_gather_bucket.all_gather_result = AllGatherResult(
586 all_gather_input=all_gather_input,
587 all_gather_output=all_gather_output,
588 handle=handle,
589 )
591 def wait_for_unshard(self) -> None:
592 """Wait all fused all-gathers and install stable unsharded parameters."""
593 for all_gather_bucket in self.all_gather_buckets:
594 all_gather_bucket.copy_out()
595 for hsdp_param in self.hsdp_params:
596 hsdp_param.wait_for_unshard()
598 def _build_reduce_scatter_buckets(
599 self,
600 reduce_scatter_reduce_op: dist.ReduceOp,
601 ) -> list[ReduceScatterBucket]:
602 """Build ordered reduce-scatter metadata for active gradients.
604 This method deliberately has no tensor or communication side effects:
605 it does not allocate buffers, pack gradients, apply scaling, clear
606 source gradients, or launch collectives.
607 """
608 params_by_bucket = {}
609 grads_by_bucket = {}
610 bucket_groups = {}
611 for hsdp_param in self.hsdp_params:
612 if not hsdp_param.sharded_param.requires_grad:
613 continue
614 if hsdp_param.unsharded_accumulated_grad is not None:
615 unsharded_grad = hsdp_param.unsharded_accumulated_grad_data
616 elif hsdp_param.unsharded_param.grad is not None:
617 unsharded_grad = hsdp_param.unsharded_grad_data
618 else:
619 continue
620 shard_group = (
621 hsdp_param.mesh_info.shard_process_group
622 if isinstance(hsdp_param.mesh_info, FSDPMeshInfo)
623 else None
624 )
625 reduce_dtype = hsdp_param.reduce_comm_dtype(unsharded_grad)
626 bucket_key = (id(shard_group), reduce_dtype)
627 if bucket_key not in params_by_bucket:
628 params_by_bucket[bucket_key] = []
629 grads_by_bucket[bucket_key] = []
630 bucket_groups[bucket_key] = shard_group
631 params_by_bucket[bucket_key].append(hsdp_param)
632 grads_by_bucket[bucket_key].append(unsharded_grad)
634 reduce_scatter_buckets = []
635 for bucket_key, hsdp_params in params_by_bucket.items():
636 shard_group = bucket_groups[bucket_key]
637 shard_world_size = hsdp_params[0].shard_world_size
638 if any(
639 hsdp_param.shard_world_size != shard_world_size
640 for hsdp_param in hsdp_params
641 ):
642 raise ValueError("A reduce-scatter bucket must use one shard world size.")
643 needs_avg_div = reduce_scatter_reduce_op == dist.ReduceOp.AVG
644 communication_op = dist.ReduceOp.SUM if needs_avg_div else reduce_scatter_reduce_op
645 param_offsets = []
646 param_numels = []
647 flat_offset = 0
648 for hsdp_param in hsdp_params:
649 param_numel = hsdp_param.padded_sharded_param_size.numel()
650 param_offsets.append(flat_offset)
651 param_numels.append(param_numel)
652 flat_offset += param_numel
653 reduce_scatter_buckets.append(
654 ReduceScatterBucket(
655 layout=self._build_gradient_bucket_layout_from(hsdp_params),
656 unsharded_grads=grads_by_bucket[bucket_key],
657 shard_group=shard_group,
658 shard_world_size=shard_world_size,
659 dtype=bucket_key[1],
660 reduce_op=communication_op,
661 needs_avg_div=needs_avg_div,
662 )
663 )
664 return reduce_scatter_buckets
666 def _build_gradient_bucket_layout_from(self, hsdp_params):
667 param_offsets = []
668 param_numels = []
669 flat_offset = 0
670 for hsdp_param in hsdp_params:
671 param_numel = hsdp_param.padded_sharded_param_size.numel()
672 param_offsets.append(flat_offset)
673 param_numels.append(param_numel)
674 flat_offset += param_numel
675 return GradientBucketLayout(
676 hsdp_params,
677 param_offsets,
678 param_numels,
679 flat_offset
680 )
682 @staticmethod
683 def _build_all_reduce_buckets(
684 reduce_scatter_buckets: list[ReduceScatterBucket],
685 ) -> list[AllReduceBucket]:
686 """Build one optional all-reduce stage for each RS bucket."""
687 all_reduce_buckets = []
688 for reduce_scatter_bucket in reduce_scatter_buckets:
689 representative_param = reduce_scatter_bucket.hsdp_params[0]
690 if not isinstance(representative_param.mesh_info, DDPMeshInfo):
691 # FSDP, no need allreduce on replicate mesh dim. skip build relative allreduce-bucket.
692 if any(
693 isinstance(hsdp_param.mesh_info, DDPMeshInfo)
694 for hsdp_param in reduce_scatter_bucket.hsdp_params[1:]
695 ):
696 raise ValueError(
697 "A reduce-scatter bucket cannot mix parameters with and without "
698 "a subsequent all-reduce."
699 )
700 continue
701 replicate_group = representative_param.mesh_info.replicate_process_group
702 replicate_world_size = representative_param.replicate_world_size
703 for hsdp_param in reduce_scatter_bucket.hsdp_params[1:]:
704 if (
705 not isinstance(hsdp_param.mesh_info, DDPMeshInfo)
706 or hsdp_param.mesh_info.replicate_process_group is not replicate_group
707 or hsdp_param.replicate_world_size != replicate_world_size
708 ):
709 raise ValueError(
710 "All parameters in a reduce-scatter bucket must share one "
711 "subsequent all-reduce group."
712 )
713 all_reduce_buckets.append(
714 AllReduceBucket(
715 layout=reduce_scatter_bucket.layout,
716 source_reduce_scatter_bucket=reduce_scatter_bucket,
717 replicate_group=replicate_group,
718 replicate_world_size=replicate_world_size,
719 dtype=reduce_scatter_bucket.dtype,
720 reduce_op=reduce_scatter_bucket.reduce_op,
721 needs_avg_div=reduce_scatter_bucket.needs_avg_div,
722 )
723 )
724 return all_reduce_buckets
726 def _issue_reduce_scatter_buckets(self, async_op: bool) -> None:
727 """Prepare RS buffers, release source gradients, and launch communication."""
728 for reduce_scatter_bucket in self.reduce_scatter_buckets:
729 reduce_scatter_input = torch.empty(
730 reduce_scatter_bucket.layout.total_numel * reduce_scatter_bucket.shard_world_size,
731 dtype=reduce_scatter_bucket.dtype,
732 device=reduce_scatter_bucket.unsharded_grads[0].device,
733 )
734 reduce_scatter_copy_in(
735 reduce_scatter_bucket.hsdp_params,
736 reduce_scatter_bucket.unsharded_grads,
737 reduce_scatter_input,
738 reduce_scatter_bucket.shard_world_size,
739 )
740 apply_gradient_scaling_factor(reduce_scatter_input, self.gradient_scaling_factor)
741 reduce_scatter_bucket.reduce_scatter_input = reduce_scatter_input
742 if not reduce_scatter_bucket.uses_collective:
743 reduce_scatter_bucket.reduce_scatter_output = reduce_scatter_input
744 else:
745 reduce_scatter_bucket.reduce_scatter_output = reduce_scatter_input.new_empty(
746 reduce_scatter_bucket.layout.total_numel
747 )
749 for hsdp_param in reduce_scatter_bucket.hsdp_params:
750 if hsdp_param.unsharded_accumulated_grad is not None:
751 hsdp_param.unsharded_accumulated_grad = None
752 else:
753 hsdp_param.unsharded_param.grad = None
754 reduce_scatter_bucket.unsharded_grads = []
755 if not reduce_scatter_bucket.uses_collective:
756 reduce_scatter_bucket.handle = None
757 continue
758 reduce_scatter_bucket.handle = dist.reduce_scatter_tensor(
759 output=reduce_scatter_bucket.reduce_scatter_output,
760 input=reduce_scatter_bucket.reduce_scatter_input,
761 group=reduce_scatter_bucket.shard_group,
762 op=reduce_scatter_bucket.reduce_op,
763 async_op=async_op,
764 )
766 @torch.no_grad()
767 def foreach_reducescatter(
768 self,
769 reduce_scatter_reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
770 async_op: bool = True,
771 ) -> Optional[torch.Tensor]:
772 """Launch fused reduce-scatter buckets for this module's gradients."""
773 self.reduce_scatter_buckets = self._build_reduce_scatter_buckets(
774 reduce_scatter_reduce_op,
775 )
776 if not self.reduce_scatter_buckets:
777 return
778 self.all_reduce_buckets = self._build_all_reduce_buckets(self.reduce_scatter_buckets)
779 self._issue_reduce_scatter_buckets(async_op)
780 if async_op:
781 self.comm_ctx.pre_param_group = self
782 else:
783 self.wait_reduce_scatter_and_issue_all_reduce(async_op=False)
784 self.wait_all_reduce_and_save_grad()
785 return
787 def _wait_reduce_scatter_buckets(self) -> None:
788 """Wait prepared RS buckets and finish shard-dimension averaging."""
789 for reduce_scatter_bucket in self.reduce_scatter_buckets:
790 if reduce_scatter_bucket.handle is not None:
791 reduce_scatter_bucket.handle.wait()
792 reduce_scatter_bucket.handle = None
793 reduce_scatter_bucket.reduce_scatter_input = None
794 if reduce_scatter_bucket.reduce_scatter_output is None:
795 raise RuntimeError("Reduce-scatter bucket has not been prepared.")
796 if reduce_scatter_bucket.needs_avg_div and reduce_scatter_bucket.shard_world_size > 1:
797 reduce_scatter_bucket.reduce_scatter_output.div_(
798 reduce_scatter_bucket.shard_world_size
799 )
802 @staticmethod
803 def _issue_all_reduce_buckets(
804 all_reduce_buckets: list[AllReduceBucket],
805 async_op: bool,
806 ) -> None:
807 """Launch in-place AR for bucket outputs transferred from RS."""
808 for all_reduce_bucket in all_reduce_buckets:
809 if all_reduce_bucket.all_reduce_output is None:
810 raise RuntimeError("All-reduce bucket has not received its reduce-scatter output.")
811 if not all_reduce_bucket.uses_collective:
812 all_reduce_bucket.handle = None
813 continue
814 all_reduce_bucket.handle = dist.all_reduce(
815 all_reduce_bucket.all_reduce_output,
816 group=all_reduce_bucket.replicate_group,
817 op=all_reduce_bucket.reduce_op,
818 async_op=async_op,
819 )
821 def wait_reduce_scatter_and_issue_all_reduce(self, async_op: bool = True) -> None:
822 """Wait all RS buckets and launch the resulting HSDP all-reduces."""
823 self._wait_reduce_scatter_buckets()
824 # Step1: get fianal reduce_scatter output
825 if not self.requires_all_reduce:
826 # if set_requires_all_reduce is False, accmu grad into reduce_partial_ouptut.
827 for reduce_scatter_bucket in self.reduce_scatter_buckets:
828 bucket_key = reduce_scatter_bucket.bucket_key
829 current_output = reduce_scatter_bucket.move_reducescatter_output()
830 reduce_partial_output = self.reduce_partial_outputs.get(bucket_key)
831 if reduce_partial_output is None:
832 self.reduce_partial_outputs[bucket_key] = current_output
833 else:
834 reduce_partial_output.add_(current_output)
835 self.reduce_scatter_buckets = []
836 self.all_reduce_buckets = []
837 return
839 for reduce_scatter_bucket in self.reduce_scatter_buckets:
840 reduce_partial_output = self.reduce_partial_outputs.pop(
841 reduce_scatter_bucket.bucket_key,
842 None,
843 )
844 if reduce_partial_output is not None:
845 # current reduce-scattered grad add accmulated grad
846 reduce_scatter_bucket.reduce_scatter_output.add_(reduce_partial_output)
847 # Step2: moving reducescatter bucket lifecycle
848 # Step2.1 build mapping from rs_bucket -> ar_bucket
849 all_reduce_by_source = {
850 id(all_reduce_bucket.source_reduce_scatter_bucket): all_reduce_bucket
851 for all_reduce_bucket in self.all_reduce_buckets
852 }
853 for reduce_scatter_bucket in self.reduce_scatter_buckets:
854 all_reduce_bucket = all_reduce_by_source.get(id(reduce_scatter_bucket))
855 if all_reduce_bucket is not None:
856 # Step2.2-A: need allreduce when HSDP
857 all_reduce_bucket.all_reduce_output = reduce_scatter_bucket.move_reducescatter_output()
858 continue
859 # Step2.2-B: FSDP only, scatter fused result to each hsdp_param manage.
860 reduce_scatter_output = reduce_scatter_bucket.move_reducescatter_output()
861 # reduce_scatter_bucket.reducescatter_output ref decrease in move method.
862 for hsdp_param, param_numel, param_offset in zip(
863 reduce_scatter_bucket.hsdp_params,
864 reduce_scatter_bucket.layout.param_numels,
865 reduce_scatter_bucket.param_offsets,
866 ):
867 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output = reduce_scatter_output.narrow(
868 0,
869 param_offset,
870 param_numel,
871 )
872 self.reduce_scatter_buckets = []
873 # Step3: luanch async allreduce buckets
874 self._issue_all_reduce_buckets(self.all_reduce_buckets, async_op)
875 if self.all_reduce_buckets:
876 self.comm_ctx.all_reduce_param_group = self
878 def wait_all_reduce_and_save_grad(self) -> None:
879 """Wait HSDP all-reduces and expose their per-parameter output views."""
880 for all_reduce_bucket in self.all_reduce_buckets:
881 if all_reduce_bucket.handle is not None:
882 all_reduce_bucket.handle.wait()
883 all_reduce_bucket.handle = None
884 if all_reduce_bucket.all_reduce_output is None:
885 raise RuntimeError("All-reduce output has already been released.")
886 if all_reduce_bucket.needs_avg_div and all_reduce_bucket.replicate_world_size > 1:
887 all_reduce_bucket.all_reduce_output.div_(all_reduce_bucket.replicate_world_size)
888 for hsdp_param, param_numel, param_offset in zip(
889 all_reduce_bucket.hsdp_params,
890 all_reduce_bucket.param_numels,
891 all_reduce_bucket.param_offsets,
892 ):
893 hsdp_param.all_reduce_comm_ctx.all_reduce_output = all_reduce_bucket.all_reduce_output.narrow(
894 0,
895 param_offset,
896 param_numel,
897 )
898 all_reduce_bucket.all_reduce_output = None
899 self.all_reduce_buckets = []
900 if self.comm_ctx.all_reduce_param_group is self:
901 self.comm_ctx.all_reduce_param_group = None
903 def reset_iter_state(self) -> None:
904 """Drop communication references after a completed iteration."""
905 for all_gather_bucket in self.all_gather_buckets:
906 if all_gather_bucket.all_gather_result is not None:
907 all_gather_bucket.all_gather_result.all_gather_input = None
908 all_gather_bucket.all_gather_result.all_gather_output = None
909 all_gather_bucket.all_gather_result.handle = None
910 all_gather_bucket.all_gather_result = None
911 for reduce_scatter_bucket in self.reduce_scatter_buckets:
912 reduce_scatter_bucket.unsharded_grads = []
913 reduce_scatter_bucket.reduce_scatter_input = None
914 reduce_scatter_bucket.reduce_scatter_output = None
915 reduce_scatter_bucket.handle = None
916 for all_reduce_bucket in self.all_reduce_buckets:
917 all_reduce_bucket.all_reduce_output = None
918 all_reduce_bucket.handle = None
919 self.reduce_scatter_buckets = []
920 self.all_reduce_buckets = []
921 self.reduce_partial_outputs.clear()
922 if self.comm_ctx.pre_param_group is self:
923 self.comm_ctx.pre_param_group = None
924 if self.comm_ctx.all_reduce_param_group is self:
925 self.comm_ctx.all_reduce_param_group = None
928class AllReduceParamGroup:
929 """Group HSDP parameters by replicate group for one fused async all-reduce.
931 Used by the ``comm_fusion=False`` path only: there each parameter issues its
932 own reduce-scatter, and this class re-fuses the resulting shards so the
933 replicate-dimension all-reduce is still a single collective. The
934 ``comm_fusion=True`` path never builds this group -- it fuses at the bucket
935 level via ``AllReduceBucket``.
937 Zero-copy is achieved by pre-allocating one contiguous buffer with 512-byte
938 alignment, having each ``reduce_scatter_grad`` write straight into an aligned
939 view of it, all-reducing the whole buffer once, and slicing gradients back
940 out of those views.
942 Numerical-correctness constraints (each is load-bearing, do not "simplify"):
944 - The all-reduce always uses SUM, never AVG. The buffer is padded for
945 alignment, and averaging would divide by a world size that the zero
946 padding does not participate in.
947 - AVG is therefore reconstructed manually in ``wait_and_split_grads`` by
948 dividing by ``replicate_world_size`` after the collective.
949 - The buffer is zero-initialized so the padding regions contribute nothing
950 to the SUM.
952 Attributes:
953 replicate_group: Process group for the replicate dimension.
954 hsdp_params: Parameters sharing that replicate group.
955 reduce_dtype: Uniform element type of the fused buffer.
956 reduce_op: Caller-requested op (AVG or SUM); decides the final scaling.
957 replicate_world_size: Size of the replicate group.
958 fused_buffer: Contiguous aligned buffer for all parameters; released
959 once ``wait_and_split_grads`` has handed out its views.
960 param_offsets: Element offset of each parameter inside the buffer.
961 param_numels: Padded element count of each parameter.
962 all_reduce_handle: In-flight async all-reduce work, or None.
963 """
965 ALIGNMENT_BYTES = 512
967 def __init__(
968 self,
969 replicate_group: dist.ProcessGroup,
970 hsdp_params: List[TorchHSDPParamV2],
971 reduce_op: dist.ReduceOp,
972 ) -> None:
973 self.replicate_group = replicate_group
974 self.hsdp_params = hsdp_params
975 self.reduce_dtype = hsdp_params[0].reduce_comm_dtype()
976 self.reduce_op = reduce_op
977 self.replicate_world_size = replicate_group.size() if replicate_group else 1
978 self.fused_buffer: Optional[torch.Tensor] = None
979 self.param_offsets: List[int] = []
980 self.param_numels: List[int] = []
981 self.all_reduce_handle: Optional[dist.Work] = None
983 def compute_aligned_layout(self) -> int:
984 """Compute packed parameter offsets and align the total buffer size."""
985 self.param_offsets = []
986 self.param_numels = []
987 element_size = torch.tensor([], dtype=self.reduce_dtype).element_size()
988 current_offset = 0
989 for hsdp_param in self.hsdp_params:
990 numel = hsdp_param.padded_sharded_param_size.numel()
991 self.param_numels.append(numel)
992 self.param_offsets.append(current_offset)
993 current_offset += numel
994 total_bytes = current_offset * element_size
995 aligned_total_bytes = (
996 (total_bytes + self.ALIGNMENT_BYTES - 1) // self.ALIGNMENT_BYTES
997 ) * self.ALIGNMENT_BYTES
998 return aligned_total_bytes // element_size
1000 def allocate_fused_buffer(self, device: torch.device) -> None:
1001 """Allocate and zero the fused all-reduce buffer."""
1002 self.fused_buffer = torch.zeros(
1003 self.compute_aligned_layout(),
1004 dtype=self.reduce_dtype,
1005 device=device,
1006 )
1008 def get_param_buffer_view(self, index: int) -> torch.Tensor:
1009 """Return one parameter's padded communication view."""
1010 if self.fused_buffer is None:
1011 raise RuntimeError("Fused buffer not allocated. Call allocate_fused_buffer first.")
1012 return self.fused_buffer.narrow(
1013 0,
1014 self.param_offsets[index],
1015 self.param_numels[index],
1016 )
1018 def get_param_grad_view(self, index: int, target_shape: torch.Size) -> torch.Tensor:
1019 """Return one parameter's actual-shard gradient view."""
1020 return self.get_param_buffer_view(index).narrow(0, 0, target_shape.numel()).view(target_shape)
1022 def accumulate_reduce_partial_outputs(self) -> None:
1023 """Merge no-all-reduce micro-step outputs into the current buffer."""
1024 if self.fused_buffer is None:
1025 return
1026 for index, hsdp_param in enumerate(self.hsdp_params):
1027 if hsdp_param.reduce_partial_output is None:
1028 continue
1029 partial_output = hsdp_param.reduce_partial_output
1030 if partial_output.dtype != self.reduce_dtype:
1031 partial_output = partial_output.to(self.reduce_dtype)
1032 self.get_param_buffer_view(index).add_(partial_output.view(-1))
1033 hsdp_param.reduce_partial_output = None
1035 def issue_async_allreduce(self) -> None:
1036 """Launch SUM all-reduce; AVG division is applied when splitting."""
1037 if self.fused_buffer is None:
1038 raise RuntimeError("Fused buffer not allocated.")
1039 self.all_reduce_handle = dist.all_reduce(
1040 self.fused_buffer,
1041 op=dist.ReduceOp.SUM,
1042 group=self.replicate_group,
1043 async_op=True,
1044 )
1046 def wait_and_split_grads(self) -> None:
1047 """Wait all-reduce and expose per-parameter context views."""
1048 if self.all_reduce_handle is not None:
1049 self.all_reduce_handle.wait()
1050 self.all_reduce_handle = None
1051 for index, hsdp_param in enumerate(self.hsdp_params):
1052 reduced_grad = self.get_param_grad_view(index, hsdp_param.sharded_size)
1053 if self.reduce_op == dist.ReduceOp.AVG and self.replicate_world_size > 1:
1054 reduced_grad.div_(self.replicate_world_size)
1055 hsdp_param.all_reduce_comm_ctx.all_reduce_output = reduced_grad
1056 self.fused_buffer = None