Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / fully_shard / param_group.py: 80%
451 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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"""MindSpore HSDP parameter groups with fused mint collectives."""
17from __future__ import annotations
19import math
20from dataclasses import dataclass, field
21from typing import Any, List, NamedTuple, Optional
23import mindspore as ms
24from mindspore.common.api import _no_grad
25import mindspore.mint.distributed as dist
27from hyper_parallel.core.fully_shard.hsdp_scheduler import ParamGroupCommCtx
28from hyper_parallel.core.fully_shard.hsdp_utils import apply_gradient_scaling_factor
29from hyper_parallel.core.fully_shard.utils import DDPMeshInfo, FSDPMeshInfo, MixedPrecisionPolicy
30from hyper_parallel.platform.mindspore.fully_shard._version_utils import copy_without_bumping_version
31from hyper_parallel.platform.mindspore.fully_shard.pack_utils import build_rs_plan, pack_for_reduce_scatter
32from hyper_parallel.platform.mindspore.fully_shard.param import MindSporeHSDPParamV2
35def _normalize_device(device: Any) -> str:
36 """Normalize a runtime device for mint allocation APIs."""
37 return str(device).split(":", 1)[0]
40def _shape_numel(shape) -> int:
41 """Return the element count of a MindSpore shape."""
42 return math.prod(int(dim) for dim in shape)
45@dataclass
46class AllGatherMetadata:
47 """Describe the rank-local layout of one fused all-gather bucket."""
49 param_input_dtypes: list[list[Any]]
50 param_input_numels: list[list[int]]
51 dtype: Any
52 inp_split_sizes: list[int]
53 total_input_numel: int
54 hash_key: int = field(init=False)
56 def __post_init__(self) -> None:
57 self.hash_key = hash(
58 (
59 tuple(tuple(dtypes) for dtypes in self.param_input_dtypes),
60 tuple(tuple(numels) for numels in self.param_input_numels),
61 self.dtype,
62 tuple(self.inp_split_sizes),
63 self.total_input_numel,
64 )
65 )
68class AllGatherResult(NamedTuple):
69 """Keep fused all-gather buffers and handle alive until copy-out."""
71 all_gather_input: Optional[ms.Tensor]
72 all_gather_output: Optional[ms.Tensor]
73 metadata: Optional[AllGatherMetadata]
74 handle: Optional[Any]
77class AllGatherMetadataCache:
78 """Cache all-gather metadata across iterations."""
80 _cache: dict[int, AllGatherMetadata] = {}
82 @classmethod
83 def get_metadata(cls, hsdp_params, fn):
84 """Retrieve or compute metadata keyed by parameter identity and version."""
85 param_key = tuple((id(param), getattr(param, "version", 0)) for param in hsdp_params)
86 key = hash(param_key)
87 if key not in cls._cache:
88 cls._cache[key] = fn(hsdp_params)
89 return cls._cache[key]
92@dataclass
93class AllGatherBucket:
94 """Own parameters sharing one all-gather group and dtype."""
96 hsdp_params: list[MindSporeHSDPParamV2]
97 shard_group: Any
98 shard_rank: int
99 shard_world_size: int
100 dtype: Any
101 metadata: AllGatherMetadata
102 all_gather_result: Optional[AllGatherResult] = None
104 @_no_grad()
105 def copy_out(self) -> None:
106 """Wait the all-gather and copy each result into stable parameter buffers."""
107 result = self.all_gather_result
108 if result is None or result.all_gather_output is None:
109 return
110 if result.handle is not None:
111 result.handle.wait()
112 all_gather_output = result.all_gather_output
113 output_buffers = []
114 for input_numels, input_dtypes, hsdp_param in zip(
115 self.metadata.param_input_numels,
116 self.metadata.param_input_dtypes,
117 self.hsdp_params,
118 ):
119 hsdp_param.init_unsharded_param_buffers(
120 input_numels,
121 input_dtypes,
122 self.shard_world_size,
123 _normalize_device(all_gather_output.device),
124 )
125 hsdp_param.alloc_unsharded_param_buffers()
126 output_buffers.extend(hsdp_param.unsharded_param_buffers)
127 split_with_sizes_copy(
128 all_gather_output.view(self.shard_world_size, -1),
129 self.metadata.inp_split_sizes,
130 dim=1,
131 out=[tensor.view(self.shard_world_size, -1) for tensor in output_buffers],
132 )
133 self.all_gather_result = None
136@dataclass
137class GradientBucketLayout:
138 """Describe the per-parameter layout shared by an RS-to-AR bucket chain."""
140 hsdp_params: list[MindSporeHSDPParamV2]
141 param_offsets: list[int]
142 param_numels: list[int]
143 total_numel: int
146@dataclass
147class ReduceScatterBucket:
148 """Own one fused reduce-scatter and its temporary buffers."""
150 layout: GradientBucketLayout
151 unsharded_grads: list[ms.Tensor]
152 shard_group: Any
153 shard_world_size: int
154 dtype: Any
155 reduce_op: str
156 needs_avg_div: bool
157 reduce_scatter_input: Optional[ms.Tensor] = None
158 reduce_scatter_output: Optional[ms.Tensor] = None
159 handle: Optional[Any] = None
161 @property
162 def hsdp_params(self) -> list[MindSporeHSDPParamV2]:
163 """Return parameters in fused-buffer order."""
164 return self.layout.hsdp_params
166 @property
167 def param_offsets(self) -> list[int]:
168 """Return parameter offsets in the reduce-scatter output."""
169 return self.layout.param_offsets
171 @property
172 def bucket_key(self) -> tuple:
173 """Return the identity of this fusion class across micro-steps."""
174 return (self.shard_group, self.dtype)
176 @property
177 def uses_collective(self) -> bool:
178 """Whether this bucket needs an actual reduce-scatter collective."""
179 return self.shard_group is not None and self.shard_world_size > 1
181 def move_reduce_scatter_output(self) -> ms.Tensor:
182 """Transfer exclusive ownership of the completed output."""
183 if self.reduce_scatter_output is None:
184 raise RuntimeError("Reduce-scatter output has already been released.")
185 output = self.reduce_scatter_output
186 self.reduce_scatter_output = None
187 return output
190@dataclass
191class AllReduceBucket:
192 """Own the all-reduce stage following one fused reduce-scatter bucket."""
194 layout: GradientBucketLayout
195 source_reduce_scatter_bucket: ReduceScatterBucket
196 replicate_group: Any
197 replicate_world_size: int
198 dtype: Any
199 needs_avg_div: bool
200 all_reduce_output: Optional[ms.Tensor] = None
201 handle: Optional[Any] = None
203 @property
204 def hsdp_params(self) -> list[MindSporeHSDPParamV2]:
205 """Return parameters in fused-buffer order."""
206 return self.layout.hsdp_params
208 @property
209 def uses_collective(self) -> bool:
210 """Whether this bucket needs an actual all-reduce collective."""
211 return self.replicate_group is not None and self.replicate_world_size > 1
214def get_all_gather_metadata(hsdp_params) -> AllGatherMetadata:
215 """Collect fused all-gather metadata for one communication dtype."""
216 param_input_dtypes = []
217 param_input_numels = []
218 inp_split_sizes = []
219 total_input_numel = 0
220 dtype = None
221 for hsdp_param in hsdp_params:
222 inputs = hsdp_param.all_gather_inputs
223 if dtype is None:
224 dtype = inputs[0].dtype
225 if any(tensor.dtype != dtype for tensor in inputs):
226 raise ValueError("All parameters in an all-gather bucket must have the same dtype.")
227 input_dtypes = [tensor.dtype for tensor in inputs]
228 input_numels = [tensor.numel() for tensor in inputs]
229 param_input_dtypes.append(input_dtypes)
230 param_input_numels.append(input_numels)
231 inp_split_sizes.extend(input_numels)
232 total_input_numel += sum(input_numels)
233 if dtype is None:
234 raise ValueError("Cannot build all-gather metadata for an empty parameter bucket.")
235 return AllGatherMetadata(
236 param_input_dtypes,
237 param_input_numels,
238 dtype,
239 inp_split_sizes,
240 total_input_numel,
241 )
244@_no_grad()
245def all_gather_copy_in(
246 all_gather_inputs,
247 all_gather_output,
248 inp_split_sizes,
249 all_gather_input_numel,
250 rank,
251):
252 """Build a contiguous fused all-gather input without writing through views."""
253 del inp_split_sizes, all_gather_input_numel, rank
254 all_gather_input = ms.mint.cat(
255 [tensor.reshape(-1) for tensor in all_gather_inputs],
256 dim=0,
257 )
258 return all_gather_input, all_gather_output
261@_no_grad()
262def split_with_sizes_copy(all_gather_output, split_sizes, dim, out):
263 """Copy dim-1 slices from a fused all-gather into stable buffers."""
264 if dim != 1:
265 raise NotImplementedError("split_with_sizes_copy currently only supports dim=1")
266 offset = 0
267 for destination, size in zip(out, split_sizes):
268 copy_without_bumping_version(
269 destination,
270 all_gather_output.narrow(dim, offset, size),
271 )
272 offset += size
275@_no_grad()
276def reduce_scatter_copy_in(
277 hsdp_params: List[MindSporeHSDPParamV2],
278 unsharded_grads: List[ms.Tensor],
279 reduce_scatter_input: ms.Tensor,
280 world_size: int,
281) -> None:
282 """Pack gradients with mixed shard dimensions without writing through views."""
283 if len(hsdp_params) != len(unsharded_grads):
284 raise AssertionError(
285 "reduce_scatter_copy_in expects one hsdp_param per unsharded_grad, but got "
286 f"{len(hsdp_params)} params and {len(unsharded_grads)} grads"
287 )
288 packed_grads = []
289 for hsdp_param, unsharded_grad in zip(hsdp_params, unsharded_grads):
290 plan = build_rs_plan(hsdp_param, unsharded_grad.contiguous(), world_size)
291 packed_grads.append(pack_for_reduce_scatter(unsharded_grad.contiguous(), plan))
292 packed_rows = ms.mint.cat(packed_grads, dim=1)
293 if packed_rows.numel() != reduce_scatter_input.numel():
294 raise AssertionError(
295 "reduce_scatter_copy_in packed an unexpected number of elements: "
296 f"{packed_rows.numel()} != {reduce_scatter_input.numel()}"
297 )
298 copy_without_bumping_version(reduce_scatter_input, packed_rows.reshape(-1))
301class HSDPParamGroup:
302 """Fuse compatible collectives for parameters in one fully_shard unit."""
304 def __init__(
305 self,
306 hsdp_params,
307 device: Optional[str] = None,
308 mp_policy: Optional[MixedPrecisionPolicy] = None,
309 enable_zero_copy: bool = False,
310 comm_ctx: Optional[ParamGroupCommCtx] = None,
311 ):
312 self.device = device
313 self.hsdp_params = hsdp_params
314 self.mp_policy = mp_policy
315 self.enable_zero_copy = enable_zero_copy
316 self.comm_ctx = comm_ctx or ParamGroupCommCtx()
317 self.gradient_scaling_factor = None
318 self.requires_all_reduce = True
319 self.all_gather_buckets: list[AllGatherBucket] = []
320 self.reduce_scatter_buckets: list[ReduceScatterBucket] = []
321 self.all_reduce_buckets: list[AllReduceBucket] = []
322 self.reduce_partial_outputs: dict[tuple, ms.Tensor] = {}
324 def _init_all_gather_buckets(self) -> None:
325 """Build ordered all-gather buckets from per-parameter routes and dtypes."""
326 params_by_bucket = {}
327 bucket_groups = {}
328 for hsdp_param in self.hsdp_params:
329 if hsdp_param.shard_world_size <= 1:
330 continue
331 if not isinstance(hsdp_param.mesh_info, FSDPMeshInfo):
332 raise ValueError(
333 f"Fused all-gather expects FSDPMeshInfo, got {type(hsdp_param.mesh_info)}"
334 )
335 shard_group = hsdp_param.mesh_info.shard_process_group
336 communication_dtype = hsdp_param.param_dtype or hsdp_param.orig_dtype
337 bucket_key = (shard_group, communication_dtype)
338 params_by_bucket.setdefault(bucket_key, []).append(hsdp_param)
339 bucket_groups[bucket_key] = shard_group
341 self.all_gather_buckets = []
342 for bucket_key, hsdp_params in params_by_bucket.items():
343 metadata = get_all_gather_metadata(hsdp_params)
344 self.all_gather_buckets.append(
345 AllGatherBucket(
346 hsdp_params=hsdp_params,
347 shard_group=bucket_groups[bucket_key],
348 shard_rank=hsdp_params[0].shard_rank,
349 shard_world_size=hsdp_params[0].shard_world_size,
350 dtype=bucket_key[1],
351 metadata=metadata,
352 )
353 )
355 def unshard(self, async_op: bool = False) -> None:
356 """Launch fused all-gathers for every compatible bucket."""
357 if self.all_gather_buckets and any(
358 bucket.all_gather_result is not None for bucket in self.all_gather_buckets
359 ):
360 return
361 self.foreach_all_gather(async_op)
363 @_no_grad()
364 def foreach_all_gather(self, async_op: bool = False) -> None:
365 """Initialize ordered buckets and launch their all-gathers."""
366 if not self.all_gather_buckets:
367 self._init_all_gather_buckets()
368 for hsdp_param in self.hsdp_params:
369 if hsdp_param.shard_world_size <= 1:
370 hsdp_param.unshard(async_op)
371 for bucket in self.all_gather_buckets:
372 if bucket.all_gather_result is not None:
373 continue
374 metadata = bucket.metadata
375 all_gather_output = ms.mint.empty(
376 (metadata.total_input_numel * bucket.shard_world_size,),
377 dtype=bucket.dtype,
378 device=_normalize_device(self.device),
379 )
380 all_gather_inputs = []
381 for hsdp_param in bucket.hsdp_params:
382 hsdp_param.reset_sharded_param()
383 all_gather_inputs.extend(hsdp_param.all_gather_inputs)
384 all_gather_input, all_gather_output = all_gather_copy_in(
385 all_gather_inputs,
386 all_gather_output,
387 metadata.inp_split_sizes,
388 metadata.total_input_numel,
389 bucket.shard_rank,
390 )
391 handle = dist.all_gather_into_tensor(
392 all_gather_output,
393 all_gather_input,
394 group=bucket.shard_group,
395 async_op=async_op,
396 )
397 bucket.all_gather_result = AllGatherResult(
398 all_gather_input=all_gather_input,
399 all_gather_output=all_gather_output,
400 metadata=metadata,
401 handle=handle,
402 )
404 def wait_for_unshard(self) -> None:
405 """Wait all all-gathers and install stable unsharded parameters."""
406 for bucket in self.all_gather_buckets:
407 bucket.copy_out()
408 for hsdp_param in self.hsdp_params:
409 hsdp_param.wait_for_unshard()
411 @staticmethod
412 def _build_gradient_bucket_layout(hsdp_params) -> GradientBucketLayout:
413 """Build one compact output layout for the supplied parameter order."""
414 param_offsets = []
415 param_numels = []
416 total_numel = 0
417 for hsdp_param in hsdp_params:
418 param_offsets.append(total_numel)
419 param_numel = _shape_numel(hsdp_param.sharded_size)
420 param_numels.append(param_numel)
421 total_numel += param_numel
422 return GradientBucketLayout(
423 hsdp_params,
424 param_offsets,
425 param_numels,
426 total_numel,
427 )
429 def _build_reduce_scatter_buckets(self, reduce_op: str) -> list[ReduceScatterBucket]:
430 """Build ordered reduce-scatter buckets without tensor side effects."""
431 params_by_bucket = {}
432 grads_by_bucket = {}
433 groups_by_bucket = {}
434 for hsdp_param in self.hsdp_params:
435 if not hsdp_param.sharded_param.requires_grad:
436 continue
437 if hsdp_param.unsharded_accumulated_grad is not None:
438 unsharded_grad = hsdp_param.unsharded_accumulated_grad_data
439 elif hsdp_param.unsharded_param.grad is not None:
440 unsharded_grad = hsdp_param.unsharded_grad_data
441 else:
442 continue
443 shard_group = (
444 hsdp_param.mesh_info.shard_process_group
445 if isinstance(hsdp_param.mesh_info, FSDPMeshInfo)
446 else None
447 )
448 reduce_dtype = hsdp_param.reduce_comm_dtype(unsharded_grad)
449 bucket_key = (shard_group, reduce_dtype)
450 params_by_bucket.setdefault(bucket_key, []).append(hsdp_param)
451 grads_by_bucket.setdefault(bucket_key, []).append(unsharded_grad)
452 groups_by_bucket[bucket_key] = shard_group
454 buckets = []
455 for bucket_key, hsdp_params in params_by_bucket.items():
456 shard_world_size = hsdp_params[0].shard_world_size
457 if any(param.shard_world_size != shard_world_size for param in hsdp_params):
458 raise ValueError("A reduce-scatter bucket must use one shard world size.")
459 needs_avg_div = reduce_op == "avg"
460 buckets.append(
461 ReduceScatterBucket(
462 layout=self._build_gradient_bucket_layout(hsdp_params),
463 unsharded_grads=grads_by_bucket[bucket_key],
464 shard_group=groups_by_bucket[bucket_key],
465 shard_world_size=shard_world_size,
466 dtype=bucket_key[1],
467 reduce_op="sum" if needs_avg_div else reduce_op,
468 needs_avg_div=needs_avg_div,
469 )
470 )
471 return buckets
473 @staticmethod
474 def _build_all_reduce_buckets(
475 reduce_scatter_buckets: list[ReduceScatterBucket],
476 ) -> list[AllReduceBucket]:
477 """Build an optional all-reduce stage for each RS bucket."""
478 buckets = []
479 for rs_bucket in reduce_scatter_buckets:
480 representative = rs_bucket.hsdp_params[0]
481 if not isinstance(representative.mesh_info, DDPMeshInfo):
482 if any(
483 isinstance(hsdp_param.mesh_info, DDPMeshInfo)
484 for hsdp_param in rs_bucket.hsdp_params[1:]
485 ):
486 raise ValueError(
487 "A reduce-scatter bucket cannot mix parameters with and without "
488 "a subsequent all-reduce."
489 )
490 continue
491 replicate_group = representative.mesh_info.replicate_process_group
492 replicate_world_size = representative.replicate_world_size
493 for hsdp_param in rs_bucket.hsdp_params[1:]:
494 if (
495 not isinstance(hsdp_param.mesh_info, DDPMeshInfo)
496 or hsdp_param.mesh_info.replicate_process_group != replicate_group
497 or hsdp_param.replicate_world_size != replicate_world_size
498 ):
499 raise ValueError(
500 "All parameters in a reduce-scatter bucket must share one "
501 "subsequent all-reduce group."
502 )
503 buckets.append(
504 AllReduceBucket(
505 layout=rs_bucket.layout,
506 source_reduce_scatter_bucket=rs_bucket,
507 replicate_group=replicate_group,
508 replicate_world_size=replicate_world_size,
509 dtype=rs_bucket.dtype,
510 needs_avg_div=rs_bucket.needs_avg_div,
511 )
512 )
513 return buckets
515 def _issue_reduce_scatter_buckets(self, async_op: bool) -> None:
516 """Prepare RS buffers, release source gradients, and launch collectives."""
517 for bucket in self.reduce_scatter_buckets:
518 reduce_scatter_input = ms.mint.empty(
519 (bucket.layout.total_numel * bucket.shard_world_size,),
520 dtype=bucket.dtype,
521 device=_normalize_device(bucket.unsharded_grads[0].device),
522 )
523 reduce_scatter_copy_in(
524 bucket.hsdp_params,
525 bucket.unsharded_grads,
526 reduce_scatter_input,
527 bucket.shard_world_size,
528 )
529 apply_gradient_scaling_factor(reduce_scatter_input, self.gradient_scaling_factor)
530 bucket.reduce_scatter_input = reduce_scatter_input
531 if bucket.uses_collective:
532 bucket.reduce_scatter_output = ms.mint.empty(
533 (bucket.layout.total_numel,),
534 dtype=bucket.dtype,
535 device=_normalize_device(reduce_scatter_input.device),
536 )
537 else:
538 bucket.reduce_scatter_output = reduce_scatter_input
539 for hsdp_param in bucket.hsdp_params:
540 hsdp_param.clear_unsharded_source_grad()
541 bucket.unsharded_grads = []
542 if not bucket.uses_collective:
543 continue
544 bucket.handle = dist.reduce_scatter_tensor(
545 bucket.reduce_scatter_output,
546 bucket.reduce_scatter_input,
547 group=bucket.shard_group,
548 op=bucket.reduce_op,
549 async_op=async_op,
550 )
552 @_no_grad()
553 def foreach_reducescatter(
554 self,
555 reduce_scatter_reduce_op: str = "avg",
556 async_op: bool = True,
557 ) -> None:
558 """Launch fused reduce-scatter buckets for this module's gradients."""
559 self.reduce_scatter_buckets = self._build_reduce_scatter_buckets(
560 reduce_scatter_reduce_op,
561 )
562 if not self.reduce_scatter_buckets:
563 return
564 self.all_reduce_buckets = self._build_all_reduce_buckets(self.reduce_scatter_buckets)
565 self._issue_reduce_scatter_buckets(async_op)
566 if async_op:
567 self.comm_ctx.pre_param_group = self
568 else:
569 self.wait_reduce_scatter_and_issue_all_reduce(async_op=False)
570 self.wait_all_reduce_and_save_grad()
572 def _wait_reduce_scatter_buckets(self) -> None:
573 """Wait RS buckets and finish shard-dimension averaging."""
574 for bucket in self.reduce_scatter_buckets:
575 if bucket.handle is not None:
576 bucket.handle.wait()
577 bucket.handle = None
578 bucket.reduce_scatter_input = None
579 if bucket.reduce_scatter_output is None:
580 raise RuntimeError("Reduce-scatter bucket has not been prepared.")
581 if bucket.needs_avg_div and bucket.shard_world_size > 1:
582 bucket.reduce_scatter_output = ms.mint.div(
583 bucket.reduce_scatter_output,
584 bucket.shard_world_size,
585 )
587 @staticmethod
588 def _issue_all_reduce_buckets(
589 all_reduce_buckets: list[AllReduceBucket],
590 async_op: bool,
591 ) -> None:
592 """Launch SUM all-reduces for completed RS outputs."""
593 for bucket in all_reduce_buckets:
594 if bucket.all_reduce_output is None:
595 raise RuntimeError("All-reduce bucket has not received its reduce-scatter output.")
596 if not bucket.uses_collective:
597 continue
598 bucket.handle = dist.all_reduce(
599 bucket.all_reduce_output,
600 group=bucket.replicate_group,
601 op="sum",
602 async_op=async_op,
603 )
605 def wait_reduce_scatter_and_issue_all_reduce(self, async_op: bool = True) -> None:
606 """Wait all RS buckets and launch the resulting HSDP all-reduces."""
607 self._wait_reduce_scatter_buckets()
608 if not self.requires_all_reduce:
609 for bucket in self.reduce_scatter_buckets:
610 current_output = bucket.move_reduce_scatter_output()
611 partial_output = self.reduce_partial_outputs.get(bucket.bucket_key)
612 self.reduce_partial_outputs[bucket.bucket_key] = (
613 current_output
614 if partial_output is None
615 else ms.mint.add(partial_output, current_output)
616 )
617 self.reduce_scatter_buckets = []
618 self.all_reduce_buckets = []
619 return
621 for bucket in self.reduce_scatter_buckets:
622 partial_output = self.reduce_partial_outputs.pop(bucket.bucket_key, None)
623 if partial_output is not None:
624 bucket.reduce_scatter_output = ms.mint.add(
625 bucket.reduce_scatter_output,
626 partial_output,
627 )
628 all_reduce_by_source = {
629 id(bucket.source_reduce_scatter_bucket): bucket
630 for bucket in self.all_reduce_buckets
631 }
632 for rs_bucket in self.reduce_scatter_buckets:
633 all_reduce_bucket = all_reduce_by_source.get(id(rs_bucket))
634 if all_reduce_bucket is not None:
635 all_reduce_bucket.all_reduce_output = rs_bucket.move_reduce_scatter_output()
636 continue
637 reduce_scatter_output = rs_bucket.move_reduce_scatter_output()
638 for hsdp_param, param_numel, param_offset in zip(
639 rs_bucket.hsdp_params,
640 rs_bucket.layout.param_numels,
641 rs_bucket.param_offsets,
642 ):
643 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output = (
644 reduce_scatter_output.narrow(0, param_offset, param_numel)
645 )
646 self.reduce_scatter_buckets = []
647 self._issue_all_reduce_buckets(self.all_reduce_buckets, async_op)
648 if self.all_reduce_buckets:
649 self.comm_ctx.all_reduce_param_group = self
651 def wait_all_reduce_and_save_grad(self) -> None:
652 """Wait all-reduces and expose their per-parameter output views."""
653 for bucket in self.all_reduce_buckets:
654 if bucket.handle is not None:
655 bucket.handle.wait()
656 bucket.handle = None
657 if bucket.all_reduce_output is None:
658 raise RuntimeError("All-reduce output has already been released.")
659 output = bucket.all_reduce_output
660 if bucket.needs_avg_div and bucket.replicate_world_size > 1:
661 output = ms.mint.div(output, bucket.replicate_world_size)
662 for hsdp_param, param_numel, param_offset in zip(
663 bucket.hsdp_params,
664 bucket.layout.param_numels,
665 bucket.layout.param_offsets,
666 ):
667 hsdp_param.all_reduce_comm_ctx.all_reduce_output = output.narrow(
668 0,
669 param_offset,
670 param_numel,
671 )
672 bucket.all_reduce_output = None
673 self.all_reduce_buckets = []
674 if self.comm_ctx.all_reduce_param_group is self:
675 self.comm_ctx.all_reduce_param_group = None
677 def reset_iter_state(self) -> None:
678 """Drop communication references after a completed iteration."""
679 for bucket in self.all_gather_buckets:
680 bucket.all_gather_result = None
681 for bucket in self.reduce_scatter_buckets:
682 bucket.unsharded_grads = []
683 bucket.reduce_scatter_input = None
684 bucket.reduce_scatter_output = None
685 bucket.handle = None
686 for bucket in self.all_reduce_buckets:
687 bucket.all_reduce_output = None
688 bucket.handle = None
689 self.reduce_scatter_buckets = []
690 self.all_reduce_buckets = []
691 self.reduce_partial_outputs.clear()
692 if self.comm_ctx.pre_param_group is self:
693 self.comm_ctx.pre_param_group = None
694 if self.comm_ctx.all_reduce_param_group is self:
695 self.comm_ctx.all_reduce_param_group = None
698class AllReduceParamGroup:
699 """Fuse per-parameter RS outputs for one HSDP replicate all-reduce."""
701 ALIGNMENT_BYTES = 512
703 def __init__(
704 self,
705 replicate_group,
706 hsdp_params: List[MindSporeHSDPParamV2],
707 reduce_op: str,
708 ):
709 self.replicate_group = replicate_group
710 self.hsdp_params = hsdp_params
711 self.reduce_dtype = hsdp_params[0].reduce_comm_dtype()
712 self.reduce_op = reduce_op
713 self.replicate_world_size = hsdp_params[0].replicate_world_size
714 self.fused_buffer: Optional[ms.Tensor] = None
715 self.param_offsets: List[int] = []
716 self.param_numels: List[int] = []
717 self.all_reduce_handle: Optional[Any] = None
719 def compute_aligned_layout(self) -> int:
720 """Compute packed parameter offsets and align total buffer size."""
721 self.param_offsets = []
722 self.param_numels = []
723 element_size = int(ms.Tensor([], dtype=self.reduce_dtype).itemsize)
724 current_offset = 0
725 for hsdp_param in self.hsdp_params:
726 numel = _shape_numel(hsdp_param.sharded_size)
727 self.param_offsets.append(current_offset)
728 self.param_numels.append(numel)
729 current_offset += numel
730 total_bytes = current_offset * element_size
731 aligned_total_bytes = (
732 (total_bytes + self.ALIGNMENT_BYTES - 1) // self.ALIGNMENT_BYTES
733 ) * self.ALIGNMENT_BYTES
734 return aligned_total_bytes // element_size
736 def allocate_fused_buffer(self, device: Any) -> None:
737 """Allocate and zero the fused all-reduce buffer."""
738 del device
739 self.fused_buffer = ms.mint.zeros(
740 (self.compute_aligned_layout(),),
741 dtype=self.reduce_dtype,
742 )
744 def get_param_buffer_view(self, index: int) -> ms.Tensor:
745 """Return one parameter's communication view."""
746 if self.fused_buffer is None:
747 raise RuntimeError("Fused buffer not allocated. Call allocate_fused_buffer first.")
748 return self.fused_buffer.narrow(
749 0,
750 self.param_offsets[index],
751 self.param_numels[index],
752 )
754 def accumulate_reduce_partial_outputs(self) -> None:
755 """Pack RS outputs and prior micro-step partials into one aligned AR buffer."""
756 param_outputs = []
757 for hsdp_param in self.hsdp_params:
758 reduced_output = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output
759 if reduced_output is None:
760 raise RuntimeError("All-reduce group requires one completed reduce-scatter output per parameter.")
761 if reduced_output.dtype != self.reduce_dtype:
762 reduced_output = reduced_output.to(self.reduce_dtype)
763 partial_output = hsdp_param.reduce_partial_output
764 if partial_output is not None:
765 if partial_output.dtype != self.reduce_dtype:
766 partial_output = partial_output.to(self.reduce_dtype)
767 reduced_output = ms.mint.add(reduced_output, partial_output)
768 hsdp_param.reduce_partial_output = None
769 param_outputs.append(reduced_output.reshape(-1))
770 hsdp_param.clear_reduce_scatter_output()
771 hsdp_param.clear_unsharded_source_grad()
772 packed_output = ms.mint.cat(param_outputs, dim=0)
773 total_numel = self.compute_aligned_layout()
774 if packed_output.numel() < total_numel:
775 padding = ms.mint.zeros(
776 (total_numel - packed_output.numel(),),
777 dtype=self.reduce_dtype,
778 )
779 packed_output = ms.mint.cat((packed_output, padding), dim=0)
780 self.fused_buffer = packed_output
782 def issue_async_allreduce(self) -> None:
783 """Launch SUM all-reduce; AVG division is applied when splitting."""
784 if self.fused_buffer is None:
785 raise RuntimeError("Fused buffer not allocated.")
786 self.all_reduce_handle = dist.all_reduce(
787 self.fused_buffer,
788 op="sum",
789 group=self.replicate_group,
790 async_op=True,
791 )
793 def wait_and_split_grads(self) -> None:
794 """Wait all-reduce and expose per-parameter context views."""
795 if self.all_reduce_handle is not None:
796 self.all_reduce_handle.wait()
797 self.all_reduce_handle = None
798 if self.fused_buffer is None:
799 raise RuntimeError("Fused buffer has already been released.")
800 output = self.fused_buffer
801 if self.reduce_op == "avg" and self.replicate_world_size > 1:
802 output = ms.mint.div(output, self.replicate_world_size)
803 for index, hsdp_param in enumerate(self.hsdp_params):
804 hsdp_param.all_reduce_comm_ctx.all_reduce_output = output.narrow(
805 0,
806 self.param_offsets[index],
807 self.param_numels[index],
808 )
809 self.fused_buffer = None
812__all__ = [
813 "AllGatherBucket",
814 "AllGatherMetadata",
815 "AllGatherMetadataCache",
816 "AllGatherResult",
817 "AllReduceBucket",
818 "AllReduceParamGroup",
819 "GradientBucketLayout",
820 "HSDPParamGroup",
821 "ReduceScatterBucket",
822 "all_gather_copy_in",
823 "get_all_gather_metadata",
824 "reduce_scatter_copy_in",
825 "split_with_sizes_copy",
826]