Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / fully_shard / state.py: 86%
394 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MindSpore HSDP cell state"""
16from collections import defaultdict
17from typing import List, Optional
18import mindspore as ms
19from mindspore import ops
20import mindspore.mint.distributed as dist
21from hyper_parallel.tools.logging import get_logger
22from hyper_parallel.core.fully_shard.hsdp_state import HSDPState
23from hyper_parallel.core.fully_shard.hsdp_utils import (
24 _get_param_module_infos,
25 FullyShardParamMode,
26 infer_fully_shard_param_mode,
27 apply_gradient_scaling_factor,
28)
29from hyper_parallel.platform.mindspore.fully_shard.pack_utils import build_rs_plan
30from hyper_parallel.platform.mindspore.fully_shard.param import MindSporeHSDPParamV2
31from hyper_parallel.platform.mindspore.fully_shard._version_utils import copy_without_bumping_version
32from hyper_parallel.platform.mindspore.fully_shard.param_group import (
33 AllReduceParamGroup,
34 HSDPParamGroup,
35 get_comm_ctx,
36)
37from hyper_parallel.platform.mindspore.utils import normalize_runtime_device
38from hyper_parallel.core.fully_shard.utils import CPUOffloadPolicy
40logger = get_logger("FSDP")
43def _to_dtype_if_needed(
44 tensor: ms.Tensor, dtype: Optional[ms.Type]
45) -> ms.Tensor:
46 """Cast tensor to the given dtype if it differs from current dtype.
48 Args:
49 tensor: The input tensor to potentially cast.
50 dtype: Target dtype. If None or same as tensor dtype, no-op.
51 """
52 if isinstance(dtype, ms.Type) and tensor.dtype != dtype:
53 return tensor.to(dtype)
54 return tensor
57class MindSporeHSDPStateV2(HSDPState):
58 """MindSpore HSDP cell state"""
59 # DTensor compat parameters in pure-TP mode can accumulate gradients
60 # directly on ``sharded_param.grad`` without materializing an
61 # ``_unsharded_param``. Track those async all-reduces separately from the
62 # standard unsharded-gradient queues.
63 pre_direct_all_reduce_grads = []
64 # Reserved for HSDP fused all-reduce pipeline (phase-2); kept for API parity with Torch.
65 pre_all_reduce_groups: List = []
66 pending_all_reduce_groups: List = []
68 @staticmethod
69 def _get_pending_unsharded_grad(hsdp_param):
70 """Return the pending unsharded gradient tensor for reduction paths."""
71 if hsdp_param.unsharded_accumulated_grad is not None:
72 return hsdp_param.unsharded_accumulated_grad_data
73 return hsdp_param.unsharded_grad_data
75 @staticmethod
76 def _has_pending_unsharded_grad(hsdp_param):
77 """Whether the parameter currently has a gradient waiting for reduction."""
78 if hsdp_param.unsharded_accumulated_grad is not None:
79 return True
80 if not hasattr(hsdp_param, "_unsharded_param") or hsdp_param.unsharded_param is None:
81 return False
82 return hsdp_param.unsharded_param.grad is not None
84 @staticmethod
85 def _get_local_sharded_grad(hsdp_param):
86 """Return the local gradient tensor currently stored on ``sharded_param``."""
87 grad = hsdp_param.sharded_param.grad
88 if grad is None:
89 return None
90 to_local = getattr(grad, "to_local", None)
91 if callable(to_local):
92 return to_local()
93 return grad
95 @staticmethod
96 def _synchronize_current_stream_if_needed(need_synchronize: bool) -> None:
97 """Synchronize the current device stream after non-blocking CPU offload."""
98 if not need_synchronize:
99 return
100 ms.runtime.current_stream().synchronize()
102 def _apply_pending_unsharded_grad_locally(self, hsdp_param) -> bool:
103 """Materialize pending unsharded grad onto ``sharded_param.grad`` without communication."""
104 pending_grad = self._get_pending_unsharded_grad(hsdp_param)
105 apply_gradient_scaling_factor(
106 pending_grad, hsdp_param.gradient_scaling_factor
107 )
108 return hsdp_param.apply_reduced_grad(pending_grad, hsdp_param.orig_dtype)
110 def __init__(self, cell, mesh_info, config, platform, device=None):
111 super().__init__(cell, mesh_info, config, platform, device)
112 self.comm_fusion = config.comm_fusion
113 # Do ReduceScatter/AllReduce for grad
114 self.mp_policy = config.mp_policy
115 self.offload_policy = config.offload_policy
116 self.reduce_grads = True
117 # Reshard parameter after backward
118 self.reshard_after_backward = True
119 # Requires AllReduce for grad When HSDP
120 self.requires_all_reduce = True
121 # Default reduce op is decided at the fully_shard-state level:
122 # if any managed parameter is DTensor-backed, use SUM; otherwise AVG.
123 self.reduce_op_type = self._resolve_default_reduce_op()
124 self._reset_sharded_params = False
125 self._init_param_group()
127 def _iter_managed_params(self):
128 """Return all fully_shard-managed parameters, including replicate_params."""
129 return [*self.hsdp_params, *self.replicate_params]
131 def _resolve_default_reduce_op(self):
132 """Resolve the default reduce op for the whole fully_shard state."""
133 for hsdp_param in self._iter_managed_params():
134 if hsdp_param.param_mode in (
135 FullyShardParamMode.DTENSOR_COMPAT,
136 FullyShardParamMode.DTENSOR_UNIFIED,
137 ):
138 return ops.ReduceOp.SUM
139 return ops.ReduceOp.AVG
141 def _resolve_reduce_op(self):
142 """Resolve the gradient reduction op for the current fully_shard state."""
143 return self.reduce_op_type
145 @staticmethod
146 def _comm_fusion_unsupported_reason(hsdp_param) -> Optional[str]:
147 """Return the reason why ``hsdp_param`` cannot participate in comm_fusion."""
148 if not hsdp_param.enable_fsdp_shard:
149 return "non-sharded parameters such as replicate_params are not supported"
150 if hsdp_param.param_mode not in (
151 FullyShardParamMode.LOCAL_PARAM,
152 FullyShardParamMode.DTENSOR_UNIFIED,
153 ):
154 return f"param_mode {hsdp_param.param_mode} is not supported"
155 local_shard = getattr(hsdp_param, "_sharded_local_tensor", None)
156 if local_shard is None:
157 return "missing local shard tensor for comm_fusion plan validation"
158 plan_world_size = getattr(hsdp_param, "shard_world_size", None)
159 if plan_world_size is None:
160 plan_world_size = getattr(hsdp_param, "shard_size", 1)
161 try:
162 build_rs_plan(hsdp_param, local_shard, plan_world_size)
163 except NotImplementedError as exc:
164 return str(exc)
165 except (AssertionError, ValueError) as exc:
166 return f"cannot build comm_fusion pack plan: {exc}"
167 return None
169 def _init_param_group(self):
170 """Initialize fused parameter group when comm_fusion is enabled."""
171 if self.config.comm_fusion:
172 unsupported_param = next(
173 (
174 hsdp_param
175 for hsdp_param in self.hsdp_params
176 if self._comm_fusion_unsupported_reason(hsdp_param) is not None
177 ),
178 None,
179 )
180 if unsupported_param is not None:
181 param_fqn = getattr(unsupported_param, "_param_fqn", "<unknown>")
182 reason = self._comm_fusion_unsupported_reason(unsupported_param)
183 raise NotImplementedError(
184 f"comm_fusion does not support parameter {param_fqn}: {reason}."
185 )
186 self.param_group = None
187 if self.hsdp_params:
188 self.param_group = HSDPParamGroup(
189 self.hsdp_params,
190 self.mesh_info,
191 self.device,
192 self.mp_policy,
193 self.config.comm_fusion_zero_copy,
194 )
196 def zero_grad(self):
197 """zero grad"""
198 for hsdp_param in self.hsdp_params:
199 hsdp_param.zero_grad()
200 for hsdp_param in self.replicate_params:
201 hsdp_param.zero_grad()
203 def _move_states_to_device(self):
204 """move states to device"""
205 for mod in self.modules:
206 for param in mod.get_parameters():
207 if hasattr(param, "_hsdp_param_initialized") and param._hsdp_param_initialized:
208 continue
209 param_device = normalize_runtime_device(param.device)
210 if param_device in (self.device, "meta"):
211 continue
212 param.data = param.to(self.device)
213 for buffer in mod.buffers():
214 if buffer.device in (self.device, "meta"):
215 continue
216 buffer.data = buffer.to(self.device)
218 def _init_hsdp_params(self):
219 """init hsdp parameters for cell and replicate parameters for cell."""
220 # all parameters in the module tree(s), deduplicated
221 visited_params = set()
222 replicate_params = set(self.config.replicate_params or ())
223 ignored_params = set(self.config.ignored_params or ())
224 filtered_params = []
225 for mod in self.modules:
226 for _, param in mod.parameters_and_names():
227 if hasattr(param, "_hsdp_param_initialized") and param._hsdp_param_initialized:
228 continue
229 if param in ignored_params:
230 continue
231 if param in visited_params:
232 continue
233 visited_params.add(param)
234 filtered_params.append(param)
236 module_infos = _get_param_module_infos(filtered_params, tuple(self.modules))
237 for param, module_info in zip(filtered_params, module_infos):
238 param_mode = infer_fully_shard_param_mode(self.config.mesh, [param])
239 enable_fsdp_shard = param not in replicate_params
240 hsdp_param = MindSporeHSDPParamV2(
241 param,
242 module_info,
243 self.mesh_info,
244 shard_placement_fn=self.config.shard_placement_fn,
245 mp_policy=self.mp_policy,
246 offload_policy=self.offload_policy,
247 device=self.device,
248 param_mode=param_mode,
249 enable_fsdp_shard=enable_fsdp_shard,
250 )
251 if param in replicate_params:
252 self.replicate_params.append(hsdp_param)
253 else:
254 self.hsdp_params.append(hsdp_param)
255 self.sharded_hsdp_params.append(hsdp_param)
257 def _init_mp_dtypes(self):
258 """init mp dtypes for hsdp parameters and replicate parameters"""
259 fused_trainable_params = []
260 fused_orig_dtypes = set()
261 fused_reduce_dtypes = set()
262 fused_all_gather_dtypes = set()
263 for hsdp_param in self.hsdp_params:
264 hsdp_param.init_dtype_attrs(self.mp_policy)
265 if not self.comm_fusion:
266 continue
267 all_gather_dtype = hsdp_param.orig_dtype
268 if hsdp_param.param_dtype is not None:
269 all_gather_dtype = hsdp_param.param_dtype
270 fused_all_gather_dtypes.add(all_gather_dtype)
271 if hsdp_param.sharded_param.requires_grad:
272 fused_trainable_params.append(hsdp_param)
273 fused_orig_dtypes.add(hsdp_param.orig_dtype)
274 fused_reduce_dtypes.add(hsdp_param.reduce_dtype)
275 for replicate_param in self.replicate_params:
276 replicate_param.init_dtype_attrs(self.mp_policy)
277 if not self.comm_fusion:
278 return
279 if len(fused_trainable_params) > 0 and len(fused_orig_dtypes) != 1:
280 raise AssertionError(
281 f"hsdp expects uniform original parameter dtype but got {fused_orig_dtypes}"
282 )
283 self._orig_dtype = next(iter(fused_orig_dtypes)) if fused_trainable_params else None
284 if len(fused_trainable_params) > 0 and len(fused_reduce_dtypes) != 1:
285 raise AssertionError(
286 f"hsdp expects uniform reduce dtype but got {fused_reduce_dtypes}"
287 )
288 self._reduce_dtype = next(iter(fused_reduce_dtypes)) if fused_trainable_params else None
289 if len(fused_all_gather_dtypes) > 1:
290 raise AssertionError(
291 "hsdp comm_fusion expects uniform all-gather parameter dtype "
292 f"but got {fused_all_gather_dtypes}"
293 )
295 def lazy_init(self):
296 """Refresh parameter views and validate runtime state before first execution."""
297 if self.is_shard and not self._reset_sharded_params:
298 for hsdp_param in self.hsdp_params:
299 if hsdp_param.is_sharded:
300 hsdp_param.reset_sharded_param()
301 self._reset_sharded_params = True
302 self._validate_no_meta_params()
303 self._validate_cpu_offload_params()
304 self._init_mp_dtypes()
306 def _validate_cpu_offload_params(self):
307 """Validate that all parameters are on CPU when CPU offload policy is enabled."""
308 if not isinstance(self.offload_policy, CPUOffloadPolicy):
309 return
310 hsdp_params_not_on_cpu = [
311 hsdp_param
312 for hsdp_param in self._iter_managed_params()
313 if not str(hsdp_param.sharded_param.device).lower().startswith("cpu")
314 ]
315 if hsdp_params_not_on_cpu:
316 raise RuntimeError(
317 "HSDP parameters should be materialized on CPU when enabling CPU offloading. "
318 "Found following parameters on non-CPU device: "
319 f"{[(p._param_fqn, p.sharded_param.device) for p in hsdp_params_not_on_cpu]}\n"
320 "MindSpore backend will support this feature in future version."
321 )
323 def _validate_no_meta_params(self):
324 """Validate that all parameters have been materialized from meta device."""
325 param_names_on_meta = [
326 hsdp_param._param_fqn
327 for hsdp_param in self._iter_managed_params()
328 if hsdp_param.sharded_param.device == "meta"
329 ]
330 if param_names_on_meta:
331 raise RuntimeError(
332 "HSDP parameters should be materialized from meta device before training, "
333 f"but the following were still on meta device: {param_names_on_meta}\n"
334 "For example, initialize the module weights on a real device before running training."
335 )
337 def _queue_replicate_params_allreduce(self) -> None:
338 """Queue async all-reduce for config.replicate_params (aligned with Torch)."""
339 for hsdp_param in self.replicate_params:
340 if not hasattr(hsdp_param, "_unsharded_param") or hsdp_param.unsharded_param is None:
341 continue
342 if not hsdp_param.sharded_param.requires_grad:
343 continue
344 if not self._has_pending_unsharded_grad(hsdp_param):
345 continue
346 if self._should_run_all_reduce(hsdp_param):
347 self._queue_compat_all_reduce(hsdp_param)
348 else:
349 need_synchronize = self._apply_pending_unsharded_grad_locally(hsdp_param)
350 self._synchronize_current_stream_if_needed(need_synchronize)
352 def _drain_reduce_scatter_params(self) -> bool:
353 """Wait pending reduce-scatter ops and apply sharded grads."""
354 need_synchronize = False
355 while HSDPState.pre_reduce_scatter_params:
356 hsdp_param, pre_orig_dtype = HSDPState.pre_reduce_scatter_params.pop(0)
357 logger.debug(
358 "post_backward module=%s wait=reduce_scatter param=%s",
359 self,
360 hsdp_param,
361 )
362 reduced_grad = hsdp_param.reduce_scatter_output()
363 hsdp_param.clear_reduce_scatter_output()
364 need_synchronize = (
365 hsdp_param.apply_reduced_grad(reduced_grad, pre_orig_dtype)
366 or need_synchronize
367 )
368 hsdp_param.accumulated_allreduced_grad = False
369 return need_synchronize
371 def reduce_scattered_params(self):
372 """Wait pending reduce-scatter ops and apply sharded grads (FSDP pipeline step 2)."""
373 need_synchronize = self._drain_reduce_scatter_params()
374 self._synchronize_current_stream_if_needed(need_synchronize)
376 def reduce_params(self):
377 """Apply reduced gradients from pre-staged all-reduce queues (aligned with Torch).
379 Drains ``pre_all_reduce_params`` and ``pre_direct_all_reduce_grads``. For
380 pending reduce-scatter work, call ``reduce_scattered_params()`` separately.
381 """
382 need_synchronize = False
383 while HSDPState.pre_all_reduce_params:
384 hsdp_param, pre_orig_dtype = HSDPState.pre_all_reduce_params.pop(0)
385 logger.debug(
386 "post_backward module=%s wait=all_reduce param=%s",
387 self,
388 hsdp_param,
389 )
390 reduced_grad = hsdp_param.all_reduce_output()
391 hsdp_param.clear_all_reduce_output()
392 need_synchronize = (
393 hsdp_param.apply_reduced_grad(reduced_grad, pre_orig_dtype)
394 or need_synchronize
395 )
396 while MindSporeHSDPStateV2.pre_direct_all_reduce_grads:
397 hsdp_param, handle, reduced_grad, target_grad, *_ = (
398 MindSporeHSDPStateV2.pre_direct_all_reduce_grads.pop(0)
399 )
400 if handle is not None:
401 logger.debug("post_backward module=%s wait=direct_compat_all_reduce", self)
402 handle.wait()
403 # all-reduce already applied SUM/AVG via _resolve_reduce_op(); skip legacy manual AVG div.
404 if hsdp_param.mp_policy.apply_grad_on_fp32_main_grad:
405 need_synchronize = (
406 hsdp_param.apply_reduced_grad(reduced_grad, hsdp_param.orig_dtype)
407 or need_synchronize
408 )
409 elif reduced_grad is not target_grad:
410 if reduced_grad.dtype != target_grad.dtype:
411 reduced_grad = reduced_grad.to(target_grad.dtype)
412 copy_without_bumping_version(target_grad, reduced_grad)
413 self._synchronize_current_stream_if_needed(need_synchronize)
415 def _wait_prev_reduce_scatter(self) -> List:
416 """Step 1: wait previous module RS for HSDP fused all-reduce groups."""
417 if MindSporeHSDPStateV2.pre_all_reduce_groups:
418 prev_groups = list(MindSporeHSDPStateV2.pre_all_reduce_groups)
419 MindSporeHSDPStateV2.pre_all_reduce_groups.clear()
420 for prev_group in prev_groups:
421 for hsdp_param in prev_group.hsdp_params:
422 hsdp_param.reduce_scatter_output()
423 hsdp_param.clear_reduce_scatter_output()
424 if hsdp_param.unsharded_accumulated_grad_data is not None:
425 hsdp_param.unsharded_accumulated_grad = None
426 elif hsdp_param.unsharded_param.grad is not None:
427 hsdp_param.unsharded_param.grad = None
428 return prev_groups
429 return []
431 def _wait_and_apply_prev_no_allreduce_params(self):
432 """Step 2: wait/apply previous reduce-scatter for pure FSDP params."""
433 self.reduce_scattered_params()
435 def _should_skip_reduce_scatter_issue(self, hsdp_param) -> bool:
436 """Return True when a parameter should not enter the HSDP RS/fused-AR pipeline."""
437 return (
438 not hasattr(hsdp_param, "_unsharded_param")
439 or hsdp_param.unsharded_param is None
440 or not hasattr(hsdp_param, "sharded_param")
441 or not hsdp_param.sharded_param.requires_grad
442 or hsdp_param.shard_size <= 1
443 or self._can_direct_all_reduce_compat_grad(hsdp_param)
444 or not self._has_pending_unsharded_grad(hsdp_param)
445 )
447 def _collect_params_for_reduce_scatter(self):
448 """Collect parameters that need the HSDP RS/fused-AR overlap pipeline."""
449 return [
450 hsdp_param
451 for hsdp_param in self._iter_managed_params()
452 if not self._should_skip_reduce_scatter_issue(hsdp_param)
453 ]
455 def _needs_overlap_post_backward_steps(self) -> bool:
456 """Whether the 4-step RS/AR overlap pipeline has pending work this hook."""
457 if MindSporeHSDPStateV2.pre_all_reduce_groups:
458 return True
459 if HSDPState.pre_reduce_scatter_params:
460 return True
461 return bool(self._collect_params_for_reduce_scatter())
463 def _run_overlap_post_backward_steps(self) -> None:
464 """Run the 4-step HSDP RS/AR overlap pipeline for the current module."""
465 prev_group = self._wait_prev_reduce_scatter()
466 self._wait_and_apply_prev_no_allreduce_params()
467 self._issue_reduce_scatter_for_current_module()
468 self._issue_prev_fused_allreduce(prev_group)
470 def _issue_reduce_scatter_for_current_module(self):
471 """Issue reduce_scatter for current module with fused all-reduce when needed."""
472 params_to_reduce = self._collect_params_for_reduce_scatter()
473 if not params_to_reduce:
474 return
476 groups_by_comm = defaultdict(list)
477 for hsdp_param in params_to_reduce:
478 if self._should_run_all_reduce(hsdp_param):
479 replicate_group = hsdp_param.unsharded_group_info.group
480 key = id(replicate_group) if replicate_group is not None else None
481 groups_by_comm[key].append(hsdp_param)
482 else:
483 groups_by_comm[None].append(hsdp_param)
485 if None in groups_by_comm:
486 for hsdp_param in groups_by_comm[None]:
487 hsdp_param.reduce_scatter_grad(
488 async_op=True,
489 dtype=hsdp_param.reduce_dtype,
490 reduce_op=self._resolve_reduce_op(),
491 )
492 HSDPState.pre_reduce_scatter_params.append(
493 (hsdp_param, hsdp_param.orig_dtype)
494 )
496 for key, hsdp_params in groups_by_comm.items():
497 if key is None:
498 continue
499 group_info = hsdp_params[0].unsharded_group_info
500 group = AllReduceParamGroup(
501 replicate_group=group_info.group,
502 hsdp_params=hsdp_params,
503 orig_dtypes=[hsdp_param.orig_dtype for hsdp_param in hsdp_params],
504 reduce_dtype=hsdp_params[0].reduce_dtype,
505 reduce_op=self._resolve_reduce_op(),
506 mp_policy=self.mp_policy,
507 replicate_world_size=group_info.rank_size,
508 )
509 group.allocate_fused_buffer(self.device)
510 for idx, hsdp_param in enumerate(hsdp_params):
511 buffer_view = group.get_param_buffer_view(idx)
512 hsdp_param.reduce_scatter_grad(
513 async_op=True,
514 dtype=hsdp_param.reduce_dtype,
515 reduce_op=self._resolve_reduce_op(),
516 output_buffer=buffer_view,
517 )
518 MindSporeHSDPStateV2.pre_all_reduce_groups.append(group)
520 def _issue_prev_fused_allreduce(self, prev_groups: List) -> None:
521 """Step 4: issue async all-reduce for previous HSDP groups (no-op without fusion groups)."""
522 for prev_group in prev_groups:
523 prev_group.accumulate_existing_grads_to_buffer()
524 prev_group.issue_async_allreduce()
525 MindSporeHSDPStateV2.pending_all_reduce_groups.append(prev_group)
527 @classmethod
528 def delay_apply_reduce_grads(cls) -> None:
529 """Wait pending fused all-reduce groups at root backward."""
530 need_synchronize = False
531 for group in cls.pending_all_reduce_groups:
532 need_synchronize = group.wait_and_apply_grads() or need_synchronize
533 cls.pending_all_reduce_groups.clear()
534 if need_synchronize:
535 ms.runtime.current_stream().synchronize()
537 def post_backward_for_comm_fusion(self):
538 """Drive the fused gradient-reduction pipeline for sharded params."""
539 logger.debug("post_backward module=%s mode=comm_fusion enter", self)
540 self.reduce_params()
541 comm_ctx = get_comm_ctx()
542 if comm_ctx.all_reduce_param_group is not None:
543 logger.debug("post_backward module=%s wait=comm_fusion_all_reduce", self)
544 comm_ctx.all_reduce_param_group.wait_all_reduce_and_apply_grad()
545 comm_ctx.all_reduce_param_group = None
546 if comm_ctx.pre_param_group is not None:
547 logger.debug("post_backward module=%s wait=comm_fusion_reduce_scatter", self)
548 comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce()
549 comm_ctx.pre_param_group = None
550 if self.param_group is not None:
551 logger.debug("post_backward module=%s launch=comm_fusion_reduce_scatter", self)
552 self.param_group.foreach_reduce(
553 reduce_scatter_reduce_op=self._resolve_reduce_op(),
554 )
555 self._queue_replicate_params_allreduce()
557 def _post_backward_without_reduce(self):
558 """Finish backward when gradient communication is disabled."""
559 if self.reshard_after_backward:
560 self.shard()
561 for hsdp_param in self._iter_managed_params():
562 hsdp_param.to_accumulated_grad_if_needed()
564 def _should_run_all_reduce(self, hsdp_param) -> bool:
565 """Whether the current parameter should issue an all-reduce in this backward pass."""
566 return self.requires_all_reduce and hsdp_param.dp_size > 1
568 def _queue_compat_all_reduce(self, hsdp_param):
569 """Queue the compatibility all-reduce path without FSDP sharding."""
570 if not self._should_run_all_reduce(hsdp_param):
571 return
572 # Pure all-reduce path: pass grad=None so all_reduce_grad fetches the
573 # unsharded grad itself and owns the scaling (no reduce-scatter here).
574 hsdp_param.all_reduce_grad(
575 dtype=hsdp_param.reduce_dtype,
576 async_op=True,
577 reduce_op=self._resolve_reduce_op(),
578 )
579 logger.debug(
580 "post_backward module=%s launch=compat_all_reduce param=%s",
581 self,
582 hsdp_param,
583 )
584 HSDPState.pre_all_reduce_params.append((hsdp_param, hsdp_param.orig_dtype))
586 def _can_direct_all_reduce_compat_grad(self, hsdp_param) -> bool:
587 """Whether ``hsdp_param`` should reduce its existing ``sharded_param.grad`` directly."""
588 if not hasattr(hsdp_param, "param_mode"):
589 return False
590 return (
591 hsdp_param.param_mode == FullyShardParamMode.DTENSOR_COMPAT
592 and hsdp_param.enable_fsdp_shard
593 and not hsdp_param.is_sharded
594 and hsdp_param.shard_size == 1
595 and hsdp_param.sharded_param.requires_grad
596 and self._should_run_all_reduce(hsdp_param)
597 and self._get_local_sharded_grad(hsdp_param) is not None
598 )
600 def _queue_direct_compat_all_reduce(self, hsdp_param):
601 """Queue all-reduce for DTENSOR_COMPAT params whose grad stays on ``sharded_param``."""
602 grad = self._get_local_sharded_grad(hsdp_param)
603 if grad is None:
604 return
605 reduced_grad = _to_dtype_if_needed(grad, hsdp_param.reduce_dtype)
606 # All-reduce needs a contiguous buffer; the local sharded grad may be a
607 # non-contiguous view. No-op when already contiguous; the copy is written
608 # back to grad in reduce_params().
609 reduced_grad = reduced_grad.contiguous()
610 # Pure all-reduce path (no reduce-scatter): this leg owns the scaling.
611 # all-reduce below is in-place, so scale in-place before it.
612 apply_gradient_scaling_factor(reduced_grad, hsdp_param.gradient_scaling_factor)
613 reduce_group_info = getattr(hsdp_param, "unsharded_group_info", None)
614 reduce_group = reduce_group_info.group if reduce_group_info is not None else None
615 reduce_group_size = reduce_group_info.rank_size if reduce_group_info is not None else 1
616 handle = None
617 if reduce_group_size > 1:
618 if reduce_group is None:
619 raise RuntimeError("Expected a valid unsharded all-reduce group when rank_size > 1")
620 handle = dist.all_reduce(
621 reduced_grad,
622 group=reduce_group,
623 op=self._resolve_reduce_op(),
624 async_op=True,
625 )
626 MindSporeHSDPStateV2.pre_direct_all_reduce_grads.append(
627 (hsdp_param, handle, reduced_grad, grad, reduce_group_size, False)
628 )
630 def post_backward(self, *_):
631 """Post-backward hook that accumulates, reduces, and reshards gradients for all managed parameters."""
632 for hsdp_param in self._iter_managed_params():
633 hsdp_param.accumulate_unsharded_grad_if_needed()
634 if not self.reduce_grads:
635 self._post_backward_without_reduce()
636 return
637 if not self.comm_fusion:
638 self.reduce_params()
639 for hsdp_param in self._iter_managed_params():
640 # replicate_params are queued once by _queue_replicate_params_allreduce().
641 if not getattr(hsdp_param, "enable_fsdp_shard", True):
642 continue
643 if not hasattr(hsdp_param, "_unsharded_param") or hsdp_param.unsharded_param is None:
644 if self._can_direct_all_reduce_compat_grad(hsdp_param):
645 self._queue_direct_compat_all_reduce(hsdp_param)
646 continue
647 if not hasattr(hsdp_param, "sharded_param") or not hsdp_param.sharded_param.requires_grad:
648 continue
649 if not self._has_pending_unsharded_grad(hsdp_param):
650 continue
651 if hsdp_param.shard_size <= 1:
652 if self._should_run_all_reduce(hsdp_param):
653 self._queue_compat_all_reduce(hsdp_param)
654 else:
655 logger.debug(
656 "post_backward module=%s apply=no_comm_grad param=%s",
657 self,
658 hsdp_param,
659 )
660 # No-communication path (shard_size == 1, no all-reduce):
661 # this leg owns the scaling since the grad never goes through
662 # reduce_scatter_grad / all_reduce_grad.
663 need_synchronize = self._apply_pending_unsharded_grad_locally(
664 hsdp_param
665 )
666 self._synchronize_current_stream_if_needed(need_synchronize)
668 if self._needs_overlap_post_backward_steps():
669 self._run_overlap_post_backward_steps()
670 self._queue_replicate_params_allreduce()
671 else:
672 self.post_backward_for_comm_fusion()
673 if self.reshard_after_backward:
674 self.shard()
676 def set_requires_grad_sync(self, requires_grad_sync):
677 """set requires grad sync flag to control gradient sync."""
678 self.reduce_grads = requires_grad_sync
680 def set_reduce_op_type(self, reduce_op_type: str):
681 """set reduce op type for gradient reduction."""
682 fsdp_support_reduce_op = {
683 "sum": ops.ReduceOp.SUM,
684 "avg": ops.ReduceOp.AVG,
685 }
686 if reduce_op_type not in fsdp_support_reduce_op:
687 raise ValueError(
688 f"Unsupported reduce op type {reduce_op_type}, "
689 f"supported types are {list(fsdp_support_reduce_op.keys())}")
690 reduce_op: str = reduce_op_type.lower().strip()
691 self.reduce_op_type = fsdp_support_reduce_op.get(reduce_op)