Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / param_group.py: 92%
515 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
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.core.fully_shard.hsdp_param import HSDPParamV2
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[HSDPParamV2]
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[HSDPParamV2]
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[HSDPParamV2]:
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[HSDPParamV2]:
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[HSDPParamV2]) -> 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[HSDPParamV2],
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[HSDPParamV2],
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 @staticmethod
667 def _build_gradient_bucket_layout_from(hsdp_params):
668 param_offsets = []
669 param_numels = []
670 flat_offset = 0
671 for hsdp_param in hsdp_params:
672 param_numel = hsdp_param.padded_sharded_param_size.numel()
673 param_offsets.append(flat_offset)
674 param_numels.append(param_numel)
675 flat_offset += param_numel
676 return GradientBucketLayout(
677 hsdp_params,
678 param_offsets,
679 param_numels,
680 flat_offset
681 )
683 @staticmethod
684 def _build_all_reduce_buckets(
685 reduce_scatter_buckets: list[ReduceScatterBucket],
686 ) -> list[AllReduceBucket]:
687 """Build one optional all-reduce stage for each RS bucket."""
688 all_reduce_buckets = []
689 for reduce_scatter_bucket in reduce_scatter_buckets:
690 representative_param = reduce_scatter_bucket.hsdp_params[0]
691 if not isinstance(representative_param.mesh_info, DDPMeshInfo):
692 # FSDP, no need allreduce on replicate mesh dim. skip build relative allreduce-bucket.
693 if any(
694 isinstance(hsdp_param.mesh_info, DDPMeshInfo)
695 for hsdp_param in reduce_scatter_bucket.hsdp_params[1:]
696 ):
697 raise ValueError(
698 "A reduce-scatter bucket cannot mix parameters with and without "
699 "a subsequent all-reduce."
700 )
701 continue
702 replicate_group = representative_param.mesh_info.replicate_process_group
703 replicate_world_size = representative_param.replicate_world_size
704 for hsdp_param in reduce_scatter_bucket.hsdp_params[1:]:
705 if (
706 not isinstance(hsdp_param.mesh_info, DDPMeshInfo)
707 or hsdp_param.mesh_info.replicate_process_group is not replicate_group
708 or hsdp_param.replicate_world_size != replicate_world_size
709 ):
710 raise ValueError(
711 "All parameters in a reduce-scatter bucket must share one "
712 "subsequent all-reduce group."
713 )
714 all_reduce_buckets.append(
715 AllReduceBucket(
716 layout=reduce_scatter_bucket.layout,
717 source_reduce_scatter_bucket=reduce_scatter_bucket,
718 replicate_group=replicate_group,
719 replicate_world_size=replicate_world_size,
720 dtype=reduce_scatter_bucket.dtype,
721 reduce_op=reduce_scatter_bucket.reduce_op,
722 needs_avg_div=reduce_scatter_bucket.needs_avg_div,
723 )
724 )
725 return all_reduce_buckets
727 def _issue_reduce_scatter_buckets(self, async_op: bool) -> None:
728 """Prepare RS buffers, release source gradients, and launch communication."""
729 for reduce_scatter_bucket in self.reduce_scatter_buckets:
730 reduce_scatter_input = torch.empty(
731 reduce_scatter_bucket.layout.total_numel * reduce_scatter_bucket.shard_world_size,
732 dtype=reduce_scatter_bucket.dtype,
733 device=reduce_scatter_bucket.unsharded_grads[0].device,
734 )
735 reduce_scatter_copy_in(
736 reduce_scatter_bucket.hsdp_params,
737 reduce_scatter_bucket.unsharded_grads,
738 reduce_scatter_input,
739 reduce_scatter_bucket.shard_world_size,
740 )
741 apply_gradient_scaling_factor(reduce_scatter_input, self.gradient_scaling_factor)
742 reduce_scatter_bucket.reduce_scatter_input = reduce_scatter_input
743 if not reduce_scatter_bucket.uses_collective:
744 reduce_scatter_bucket.reduce_scatter_output = reduce_scatter_input
745 else:
746 reduce_scatter_bucket.reduce_scatter_output = reduce_scatter_input.new_empty(
747 reduce_scatter_bucket.layout.total_numel
748 )
750 for hsdp_param in reduce_scatter_bucket.hsdp_params:
751 if hsdp_param.unsharded_accumulated_grad is not None:
752 hsdp_param.unsharded_accumulated_grad = None
753 else:
754 hsdp_param.unsharded_param.grad = None
755 reduce_scatter_bucket.unsharded_grads = []
756 if not reduce_scatter_bucket.uses_collective:
757 reduce_scatter_bucket.handle = None
758 continue
759 reduce_scatter_bucket.handle = dist.reduce_scatter_tensor(
760 output=reduce_scatter_bucket.reduce_scatter_output,
761 input=reduce_scatter_bucket.reduce_scatter_input,
762 group=reduce_scatter_bucket.shard_group,
763 op=reduce_scatter_bucket.reduce_op,
764 async_op=async_op,
765 )
767 @torch.no_grad()
768 def foreach_reducescatter(
769 self,
770 reduce_scatter_reduce_op: Optional[dist.ReduceOp] = dist.ReduceOp.AVG,
771 async_op: bool = True,
772 ) -> Optional[torch.Tensor]:
773 """Launch fused reduce-scatter buckets for this module's gradients."""
774 self.reduce_scatter_buckets = self._build_reduce_scatter_buckets(
775 reduce_scatter_reduce_op,
776 )
777 if not self.reduce_scatter_buckets:
778 return
779 self.all_reduce_buckets = self._build_all_reduce_buckets(self.reduce_scatter_buckets)
780 self._issue_reduce_scatter_buckets(async_op)
781 if async_op:
782 self.comm_ctx.pre_param_group = self
783 else:
784 self.wait_reduce_scatter_and_issue_all_reduce(async_op=False)
785 self.wait_all_reduce_and_save_grad()
786 return
788 def _wait_reduce_scatter_buckets(self) -> None:
789 """Wait prepared RS buckets and finish shard-dimension averaging."""
790 for reduce_scatter_bucket in self.reduce_scatter_buckets:
791 if reduce_scatter_bucket.handle is not None:
792 reduce_scatter_bucket.handle.wait()
793 reduce_scatter_bucket.handle = None
794 reduce_scatter_bucket.reduce_scatter_input = None
795 if reduce_scatter_bucket.reduce_scatter_output is None:
796 raise RuntimeError("Reduce-scatter bucket has not been prepared.")
797 if reduce_scatter_bucket.needs_avg_div and reduce_scatter_bucket.shard_world_size > 1:
798 reduce_scatter_bucket.reduce_scatter_output.div_(
799 reduce_scatter_bucket.shard_world_size
800 )
803 @staticmethod
804 def _issue_all_reduce_buckets(
805 all_reduce_buckets: list[AllReduceBucket],
806 async_op: bool,
807 ) -> None:
808 """Launch in-place AR for bucket outputs transferred from RS."""
809 for all_reduce_bucket in all_reduce_buckets:
810 if all_reduce_bucket.all_reduce_output is None:
811 raise RuntimeError("All-reduce bucket has not received its reduce-scatter output.")
812 if not all_reduce_bucket.uses_collective:
813 all_reduce_bucket.handle = None
814 continue
815 all_reduce_bucket.handle = dist.all_reduce(
816 all_reduce_bucket.all_reduce_output,
817 group=all_reduce_bucket.replicate_group,
818 op=all_reduce_bucket.reduce_op,
819 async_op=async_op,
820 )
822 def wait_reduce_scatter_and_issue_all_reduce(self, async_op: bool = True) -> None:
823 """Wait all RS buckets and launch the resulting HSDP all-reduces."""
824 self._wait_reduce_scatter_buckets()
825 # Step1: get fianal reduce_scatter output
826 if not self.requires_all_reduce:
827 # if set_requires_all_reduce is False, accmu grad into reduce_partial_ouptut.
828 for reduce_scatter_bucket in self.reduce_scatter_buckets:
829 bucket_key = reduce_scatter_bucket.bucket_key
830 current_output = reduce_scatter_bucket.move_reducescatter_output()
831 reduce_partial_output = self.reduce_partial_outputs.get(bucket_key)
832 if reduce_partial_output is None:
833 self.reduce_partial_outputs[bucket_key] = current_output
834 else:
835 reduce_partial_output.add_(current_output)
836 self.reduce_scatter_buckets = []
837 self.all_reduce_buckets = []
838 return
840 for reduce_scatter_bucket in self.reduce_scatter_buckets:
841 reduce_partial_output = self.reduce_partial_outputs.pop(
842 reduce_scatter_bucket.bucket_key,
843 None,
844 )
845 if reduce_partial_output is not None:
846 # current reduce-scattered grad add accmulated grad
847 reduce_scatter_bucket.reduce_scatter_output.add_(reduce_partial_output)
848 # Step2: moving reducescatter bucket lifecycle
849 # Step2.1 build mapping from rs_bucket -> ar_bucket
850 all_reduce_by_source = {
851 id(all_reduce_bucket.source_reduce_scatter_bucket): all_reduce_bucket
852 for all_reduce_bucket in self.all_reduce_buckets
853 }
854 for reduce_scatter_bucket in self.reduce_scatter_buckets:
855 all_reduce_bucket = all_reduce_by_source.get(id(reduce_scatter_bucket))
856 if all_reduce_bucket is not None:
857 # Step2.2-A: need allreduce when HSDP
858 all_reduce_bucket.all_reduce_output = reduce_scatter_bucket.move_reducescatter_output()
859 continue
860 # Step2.2-B: FSDP only, scatter fused result to each hsdp_param manage.
861 reduce_scatter_output = reduce_scatter_bucket.move_reducescatter_output()
862 # reduce_scatter_bucket.reducescatter_output ref decrease in move method.
863 for hsdp_param, param_numel, param_offset in zip(
864 reduce_scatter_bucket.hsdp_params,
865 reduce_scatter_bucket.layout.param_numels,
866 reduce_scatter_bucket.param_offsets,
867 ):
868 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output = reduce_scatter_output.narrow(
869 0,
870 param_offset,
871 param_numel,
872 )
873 self.reduce_scatter_buckets = []
874 # Step3: luanch async allreduce buckets
875 self._issue_all_reduce_buckets(self.all_reduce_buckets, async_op)
876 if self.all_reduce_buckets:
877 self.comm_ctx.all_reduce_param_group = self
879 def wait_all_reduce_and_save_grad(self) -> None:
880 """Wait HSDP all-reduces and expose their per-parameter output views."""
881 for all_reduce_bucket in self.all_reduce_buckets:
882 if all_reduce_bucket.handle is not None:
883 all_reduce_bucket.handle.wait()
884 all_reduce_bucket.handle = None
885 if all_reduce_bucket.all_reduce_output is None:
886 raise RuntimeError("All-reduce output has already been released.")
887 if all_reduce_bucket.needs_avg_div and all_reduce_bucket.replicate_world_size > 1:
888 all_reduce_bucket.all_reduce_output.div_(all_reduce_bucket.replicate_world_size)
889 for hsdp_param, param_numel, param_offset in zip(
890 all_reduce_bucket.hsdp_params,
891 all_reduce_bucket.param_numels,
892 all_reduce_bucket.param_offsets,
893 ):
894 hsdp_param.all_reduce_comm_ctx.all_reduce_output = all_reduce_bucket.all_reduce_output.narrow(
895 0,
896 param_offset,
897 param_numel,
898 )
899 all_reduce_bucket.all_reduce_output = None
900 self.all_reduce_buckets = []
901 if self.comm_ctx.all_reduce_param_group is self:
902 self.comm_ctx.all_reduce_param_group = None
904 def reset_iter_state(self) -> None:
905 """Drop communication references after a completed iteration."""
906 for all_gather_bucket in self.all_gather_buckets:
907 if all_gather_bucket.all_gather_result is not None:
908 all_gather_bucket.all_gather_result.all_gather_input = None
909 all_gather_bucket.all_gather_result.all_gather_output = None
910 all_gather_bucket.all_gather_result.handle = None
911 all_gather_bucket.all_gather_result = None
912 for reduce_scatter_bucket in self.reduce_scatter_buckets:
913 reduce_scatter_bucket.unsharded_grads = []
914 reduce_scatter_bucket.reduce_scatter_input = None
915 reduce_scatter_bucket.reduce_scatter_output = None
916 reduce_scatter_bucket.handle = None
917 for all_reduce_bucket in self.all_reduce_buckets:
918 all_reduce_bucket.all_reduce_output = None
919 all_reduce_bucket.handle = None
920 self.reduce_scatter_buckets = []
921 self.all_reduce_buckets = []
922 self.reduce_partial_outputs.clear()
923 if self.comm_ctx.pre_param_group is self:
924 self.comm_ctx.pre_param_group = None
925 if self.comm_ctx.all_reduce_param_group is self:
926 self.comm_ctx.all_reduce_param_group = None
929class AllReduceParamGroup:
930 """Group HSDP parameters by replicate group for one fused async all-reduce.
932 Used by the ``comm_fusion=False`` path only: there each parameter issues its
933 own reduce-scatter, and this class re-fuses the resulting shards so the
934 replicate-dimension all-reduce is still a single collective. The
935 ``comm_fusion=True`` path never builds this group -- it fuses at the bucket
936 level via ``AllReduceBucket``.
938 Zero-copy is achieved by pre-allocating one contiguous buffer with 512-byte
939 alignment, having each ``reduce_scatter_grad`` write straight into an aligned
940 view of it, all-reducing the whole buffer once, and slicing gradients back
941 out of those views.
943 Numerical-correctness constraints (each is load-bearing, do not "simplify"):
945 - The all-reduce always uses SUM, never AVG. The buffer is padded for
946 alignment, and averaging would divide by a world size that the zero
947 padding does not participate in.
948 - AVG is therefore reconstructed manually in ``wait_and_split_grads`` by
949 dividing by ``replicate_world_size`` after the collective.
950 - The buffer is zero-initialized so the padding regions contribute nothing
951 to the SUM.
953 Attributes:
954 replicate_group: Process group for the replicate dimension.
955 hsdp_params: Parameters sharing that replicate group.
956 reduce_dtype: Uniform element type of the fused buffer.
957 reduce_op: Caller-requested op (AVG or SUM); decides the final scaling.
958 replicate_world_size: Size of the replicate group.
959 fused_buffer: Contiguous aligned buffer for all parameters; released
960 once ``wait_and_split_grads`` has handed out its views.
961 param_offsets: Element offset of each parameter inside the buffer.
962 param_numels: Padded element count of each parameter.
963 all_reduce_handle: In-flight async all-reduce work, or None.
964 """
966 ALIGNMENT_BYTES = 512
968 def __init__(
969 self,
970 replicate_group: dist.ProcessGroup,
971 hsdp_params: List[HSDPParamV2],
972 reduce_op: dist.ReduceOp,
973 ) -> None:
974 self.replicate_group = replicate_group
975 self.hsdp_params = hsdp_params
976 self.reduce_dtype = hsdp_params[0].reduce_comm_dtype()
977 self.reduce_op = reduce_op
978 self.replicate_world_size = replicate_group.size() if replicate_group else 1
979 self.fused_buffer: Optional[torch.Tensor] = None
980 self.param_offsets: List[int] = []
981 self.param_numels: List[int] = []
982 self.all_reduce_handle: Optional[dist.Work] = None
984 def compute_aligned_layout(self) -> int:
985 """Compute packed parameter offsets and align the total buffer size."""
986 self.param_offsets = []
987 self.param_numels = []
988 element_size = torch.tensor([], dtype=self.reduce_dtype).element_size()
989 current_offset = 0
990 for hsdp_param in self.hsdp_params:
991 numel = hsdp_param.padded_sharded_param_size.numel()
992 self.param_numels.append(numel)
993 self.param_offsets.append(current_offset)
994 current_offset += numel
995 total_bytes = current_offset * element_size
996 aligned_total_bytes = (
997 (total_bytes + self.ALIGNMENT_BYTES - 1) // self.ALIGNMENT_BYTES
998 ) * self.ALIGNMENT_BYTES
999 return aligned_total_bytes // element_size
1001 def allocate_fused_buffer(self, device: torch.device) -> None:
1002 """Allocate and zero the fused all-reduce buffer."""
1003 self.fused_buffer = torch.zeros(
1004 self.compute_aligned_layout(),
1005 dtype=self.reduce_dtype,
1006 device=device,
1007 )
1009 def get_param_buffer_view(self, index: int) -> torch.Tensor:
1010 """Return one parameter's padded communication view."""
1011 if self.fused_buffer is None:
1012 raise RuntimeError("Fused buffer not allocated. Call allocate_fused_buffer first.")
1013 return self.fused_buffer.narrow(
1014 0,
1015 self.param_offsets[index],
1016 self.param_numels[index],
1017 )
1019 def get_param_grad_view(self, index: int, target_shape: torch.Size) -> torch.Tensor:
1020 """Return one parameter's actual-shard gradient view."""
1021 return self.get_param_buffer_view(index).narrow(0, 0, target_shape.numel()).view(target_shape)
1023 def accumulate_reduce_partial_outputs(self) -> None:
1024 """Merge no-all-reduce micro-step outputs into the current buffer."""
1025 if self.fused_buffer is None:
1026 return
1027 for index, hsdp_param in enumerate(self.hsdp_params):
1028 if hsdp_param.reduce_partial_output is None:
1029 continue
1030 partial_output = hsdp_param.reduce_partial_output
1031 if partial_output.dtype != self.reduce_dtype:
1032 partial_output = partial_output.to(self.reduce_dtype)
1033 self.get_param_buffer_view(index).add_(partial_output.view(-1))
1034 hsdp_param.reduce_partial_output = None
1036 def issue_async_allreduce(self) -> None:
1037 """Launch SUM all-reduce; AVG division is applied when splitting."""
1038 if self.fused_buffer is None:
1039 raise RuntimeError("Fused buffer not allocated.")
1040 self.all_reduce_handle = dist.all_reduce(
1041 self.fused_buffer,
1042 op=dist.ReduceOp.SUM,
1043 group=self.replicate_group,
1044 async_op=True,
1045 )
1047 def wait_and_split_grads(self) -> None:
1048 """Wait all-reduce and expose per-parameter context views."""
1049 if self.all_reduce_handle is not None:
1050 self.all_reduce_handle.wait()
1051 self.all_reduce_handle = None
1052 for index, hsdp_param in enumerate(self.hsdp_params):
1053 reduced_grad = self.get_param_grad_view(index, hsdp_param.sharded_size)
1054 if self.reduce_op == dist.ReduceOp.AVG and self.replicate_world_size > 1:
1055 reduced_grad.div_(self.replicate_world_size)
1056 hsdp_param.all_reduce_comm_ctx.all_reduce_output = reduced_grad
1057 self.fused_buffer = None