Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / hsdp_state.py: 84%
296 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07: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"""HSDP cell state"""
16# pylint: disable=protected-access
18from collections import defaultdict
19from typing import List, Mapping, Optional
21import torch
22from torch import nn
24from hyper_parallel.tools.logging import get_logger
25from hyper_parallel.core.dtensor.dtensor import DTensor
26from hyper_parallel.core.fully_shard.hsdp_utils import (
27 _get_param_module_infos,
28)
29from hyper_parallel.core.fully_shard.utils import (
30 CPUOffloadPolicy,
31 DDPMeshInfo,
32 FSDPMeshInfo,
33 HSDPMeshInfo,
34 SourceShardMetaInfo,
35)
36from hyper_parallel.core.fully_shard.hsdp_param import HSDPParamV2
37from hyper_parallel.core.fully_shard.param_group import HSDPParamGroup, AllReduceParamGroup
39logger = get_logger("FSDP")
41nn.Module = nn.Module
44def _to_dtype_if_needed(
45 tensor: torch.Tensor, dtype: Optional[torch.dtype]
46) -> torch.Tensor:
47 """Cast tensor to the given dtype if it differs from current dtype.
49 Args:
50 tensor: The input tensor to potentially cast.
51 dtype: Target dtype. If None or same as tensor dtype, no-op.
52 """
53 if dtype is not None and tensor.dtype != dtype:
54 return tensor.to(dtype)
55 return tensor
58class HSDPStateV2:
59 """HSDP state for cell"""
61 def __init__(
62 self,
63 cell,
64 mesh,
65 shard_placement_fn,
66 comm_fusion_policy,
67 mp_policy,
68 offload_policy,
69 raw_ignored_params,
70 raw_replicate_params,
71 scheduler_ctx,
72 device,
73 source_shard_infos: Optional[Mapping[torch.nn.Parameter, SourceShardMetaInfo]] = None,
74 ):
75 """
76 Initialize HSDPStateV2.
78 Args:
79 cell (nn.Module): The module whose parameters are managed by this state.
80 mesh: Mesh topology for shard/replicate dimensions.
81 shard_placement_fn: Optional function selecting the parameter shard dimension.
82 comm_fusion_policy: Communication fusion configuration.
83 mp_policy: Mixed-precision policy.
84 offload_policy: Parameter offload policy.
85 raw_ignored_params: Parameters excluded from fully_shard management.
86 raw_replicate_params: Managed parameters that remain replicated.
87 scheduler_ctx: Scheduler context shared by this module tree.
88 device (torch.device): Target device.
89 source_shard_infos: Optional per-parameter TP/EP source metadata.
90 """
91 self.source_shard_infos = source_shard_infos
92 self.modules = (cell,) if isinstance(cell, nn.Module) else tuple(cell)
93 self.cell = self.modules[0]
94 self.mesh = mesh
95 self.shard_placement_fn = shard_placement_fn
96 self.mp_policy = mp_policy
97 self.offload_policy = offload_policy
98 self.comm_fusion_policy = comm_fusion_policy
99 self.raw_ignored_params = set(raw_ignored_params or ())
100 self.raw_replicate_params = set(raw_replicate_params or ())
101 self.scheduler_ctx = scheduler_ctx
102 self.device = device
103 self.hsdp_params: List[HSDPParamV2] = []
104 self.param_group = None
105 self._move_states_to_device()
106 self._init_hsdp_params()
107 self.is_shard = True
108 self.module_name = None
109 # requires_gradient_sync
110 self.reduce_grads = True
111 # Reshard parameter after backward
112 self.reshard_after_backward = True
113 # Requires AllReduce for grad When HSDP
114 self.requires_all_reduce = True
115 self.set_reduce_op_type("avg")
116 self._reset_sharded_params = False
117 self._init_param_group()
119 def __repr__(self) -> str:
120 """Stable debug name used in log lines.
122 ``module_name`` is only assigned by the root forward pre-hook, so fall
123 back to the managed module class and object id before names are set.
124 Logging's ``%s`` calls this lazily -- only when a record is emitted.
125 """
126 if self.module_name:
127 return str(self.module_name)
128 return f"{self.cell.__class__.__name__}@{id(self.cell):x}"
130 def _init_param_group(self):
131 """Initialize fused parameter group for communication fusion.
133 All managed parameters enter one ``HSDPParamGroup``. Parameters without
134 an FSDP shard dimension use the group's local all-gather and
135 reduce-scatter paths before entering replicate all-reduce buckets.
136 """
137 self.param_group = None
138 if not self.comm_fusion_policy.enable_comm_fusion:
139 return
140 if self.hsdp_params:
141 # pylint: disable=E1128
142 self.param_group = HSDPParamGroup(
143 self.hsdp_params,
144 self.device,
145 self.comm_fusion_policy.comm_fusion_zero_copy,
146 comm_ctx=self.scheduler_ctx.param_group_comm_ctx,
147 )
149 def _move_states_to_device(self):
150 """move states to device"""
151 for mod in self.modules:
152 for param in mod.parameters():
153 if hasattr(param, "_hsdp_param_initialized") and param._hsdp_param_initialized:
154 continue
155 if param.device == self.device or param.device.type == "meta":
156 continue
157 param.data = param.to(self.device)
158 for buffer in mod.buffers():
159 if buffer.device == self.device or buffer.device.type == "meta":
160 continue
161 buffer.data = buffer.to(self.device)
163 def _build_param_source_shard_info(
164 self, param: torch.nn.Parameter
165 ) -> Optional[SourceShardMetaInfo]:
166 """Build normalized source-layout metadata for one managed parameter."""
167 if isinstance(param, DTensor):
168 if self.source_shard_infos is not None:
169 raise ValueError(
170 "source_shard_infos cannot be provided when fully_shard manages a native DTensor parameter"
171 )
172 return SourceShardMetaInfo(
173 mesh=param.device_mesh,
174 placements=tuple(param.placements),
175 origin_is_dtensor=True,
176 )
177 if self.source_shard_infos is None:
178 return None
179 return self.source_shard_infos.get(param)
181 def _init_hsdp_params(self):
182 """Initialize all fully_shard-managed parameters for the module."""
183 # all parameters in the module tree(s), deduplicated
184 visited_params = set()
185 filtered_params = []
186 for mod in self.modules:
187 for _, param in mod.named_parameters():
188 if param in self.raw_ignored_params:
189 continue
190 if hasattr(param, "_hsdp_param_initialized") and param._hsdp_param_initialized:
191 continue
192 if param in visited_params:
193 continue
194 visited_params.add(param)
195 filtered_params.append(param)
197 module_infos = _get_param_module_infos(filtered_params, tuple(self.modules))
198 for param, module_info in zip(filtered_params, module_infos):
199 self.hsdp_params.append(
200 HSDPParamV2(
201 param,
202 module_info,
203 self._build_param_mesh_info(param),
204 shard_placement_fn=self.shard_placement_fn,
205 mp_policy=self.mp_policy,
206 offload_policy=self.offload_policy,
207 device=self.device,
208 source_shard_info=self._build_param_source_shard_info(param),
209 )
210 )
212 def _build_param_mesh_info(self, parameter):
213 if self.mesh.ndim not in (1, 2):
214 raise ValueError(
215 "fully_shard only supports explicit 1D DP/FSDP meshes or 2D HSDP meshes. "
216 f"Got mesh.ndim={self.mesh.ndim}."
217 )
218 if parameter in self.raw_replicate_params:
219 return DDPMeshInfo(
220 mesh=self.mesh if self.mesh.ndim == 1 else self.mesh.flatten(),
221 replicate_mesh_dim=0,
222 )
223 if self.mesh.ndim == 1:
224 return FSDPMeshInfo(mesh=self.mesh, shard_mesh_dim=0)
225 return HSDPMeshInfo(
226 mesh=self.mesh,
227 shard_mesh_dim=1,
228 replicate_mesh_dim=0,
229 )
231 def _init_mp_dtypes(self):
232 """Initialize dtype attributes for each parameter from the mixed-precision policy."""
233 for hsdp_param in self.hsdp_params:
234 hsdp_param.init_dtype_attrs(self.mp_policy)
236 def _validate_cpu_offload_params(self):
237 """Validate that all parameters are on CPU when CPU offload policy is enabled."""
238 if not isinstance(self.offload_policy, CPUOffloadPolicy):
239 return
240 hsdp_params_not_on_cpu = [
241 hsdp_param
242 for hsdp_param in self.hsdp_params
243 if hsdp_param.sharded_param.device.type != "cpu"
244 ]
245 if hsdp_params_not_on_cpu:
246 raise RuntimeError(
247 "HSDP parameters should be materialized on CPU when enabling CPU offloading. "
248 'For example, load a CPU state dict or call module.to_empty(device="cpu"). '
249 "Found following parameters on non-CPU device: "
250 f"{[(p._param_fqn, p.sharded_param.device) for p in hsdp_params_not_on_cpu]}\n"
251 )
253 def lazy_init(self):
254 """Deferred initialization: reset sharded params, validate devices, and set mixed-precision dtypes."""
255 if self.is_shard and not self._reset_sharded_params:
256 for hsdp_param in self.hsdp_params:
257 hsdp_param.reset_sharded_param()
258 self._reset_sharded_params = True
259 self._validate_no_meta_params()
260 self._validate_cpu_offload_params()
261 self._init_mp_dtypes()
263 def _validate_no_meta_params(self):
264 param_names_on_meta = [
265 hsdp_param._param_fqn
266 for hsdp_param in self.hsdp_params
267 if hsdp_param.sharded_param.device.type == "meta"
268 ]
269 if param_names_on_meta:
270 raise RuntimeError(
271 "HSDP parameters should be materialized from meta device before training, "
272 f"but the following were still on meta device: {param_names_on_meta}\n"
273 "For example, call module.to_empty(device) to materialize to device and "
274 "call module.reset_parameters() on each module to initialize values."
275 )
277 def post_backward_for_comm_fusion(self):
278 """post_backward_for_comm_fusion."""
279 logger.debug("post_backward module=%s mode=comm_fusion enter", self)
280 # Fused gradient reduction path: first apply any pending async reduction
281 # from the previous module's backward (pipelined overlap), then issue
282 # this module's fused reduce-scatter (+ all-reduce for HSDP).
283 comm_ctx = self.scheduler_ctx.param_group_comm_ctx
284 # Phase 2: save gradients for the param group whose all-reduce is done.
285 if comm_ctx.all_reduce_param_group is not None:
286 logger.debug("post_backward module=%s wait=comm_fusion_all_reduce", self)
287 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad()
288 comm_ctx.all_reduce_param_group = None
289 # Phase 1: wait reduce_scatter, issue async all_reduce for previous layer
290 if comm_ctx.pre_param_group is not None:
291 logger.debug("post_backward module=%s wait=comm_fusion_reduce_scatter", self)
292 comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce()
293 comm_ctx.pre_param_group = None
294 if self.param_group is not None:
295 logger.debug("post_backward module=%s launch=comm_fusion_reduce_scatter", self)
296 self.param_group.foreach_reducescatter(
297 reduce_scatter_reduce_op=self.reduce_op_type,
298 )
300 def post_backward(self, *unused): # pylint: disable=unused-argument
301 """Reduce gradients and reshard parameters after backward."""
302 logger.debug(
303 "post_backward module=%s enter reduce_grads=%s comm_fusion=%s reshard_after_backward=%s",
304 self,
305 self.reduce_grads,
306 self.comm_fusion_policy.enable_comm_fusion,
307 self.reshard_after_backward,
308 )
309 for hsdp_param in self.hsdp_params:
310 hsdp_param.accumulate_unsharded_grad_if_needed()
311 if not self.reduce_grads:
312 if self.reshard_after_backward:
313 self.shard()
314 for hsdp_param in self.hsdp_params:
315 hsdp_param.to_accumulated_grad_if_needed()
316 return
317 if self.reshard_after_backward:
318 # Reshard before gradient communication to reduce backward memory peak.
319 self.shard()
320 if not self.comm_fusion_policy.enable_comm_fusion:
321 # Step 1: wait previous reduce-scatter (for params needing all-reduce)
322 prev_group = self._wait_prev_reduce_scatter()
324 # Step 2: wait previous reduce-scatter outputs that skip replicate all-reduce
325 self._wait_prev_reduce_scatter_without_all_reduce()
327 # Step 3: issue current reduce_scatter
328 self._issue_reduce_scatter_for_current_module()
330 # Step 4: issue previous fused all-reduce asynchronously
331 self._issue_prev_fused_all_reduce(prev_group)
332 else:
333 self.post_backward_for_comm_fusion()
335 def _issue_reduce_scatter_for_current_module(self):
336 """Issue reduce_scatter for current module's parameters with fused all-reduce support.
338 This method groups parameters by their replicate_process_group and:
339 1. For params without all_reduce needs: issue reduce_scatter directly
340 2. For params with all_reduce needs: allocate fused buffer and issue reduce_scatter
341 into aligned views, enabling zero-copy fused all_reduce later.
342 """
343 # Collect parameters that need gradient reduction
344 params_to_reduce = []
345 for hsdp_param in self.hsdp_params:
346 skip_param = (
347 not hasattr(hsdp_param, "_unsharded_param")
348 or not hsdp_param.unsharded_param_buffers
349 or not hsdp_param.sharded_param.requires_grad
350 or (
351 hsdp_param.unsharded_param.grad is None
352 and hsdp_param.unsharded_accumulated_grad_data is None
353 )
354 )
355 if skip_param:
356 continue
357 params_to_reduce.append(hsdp_param)
359 if not params_to_reduce:
360 return
362 # Group by replicate process group and reduction dtype so every fused
363 # all-reduce buffer has one communication group and one element type.
364 groups_by_comm = defaultdict(list)
365 for hsdp_param in params_to_reduce:
366 if self.requires_all_reduce and hsdp_param.replicate_world_size > 1:
367 replicate_process_group = hsdp_param.mesh_info.replicate_process_group
368 group_key = (id(replicate_process_group), hsdp_param.reduce_comm_dtype())
369 groups_by_comm[group_key].append(hsdp_param)
370 else:
371 groups_by_comm[None].append(hsdp_param)
373 # Handle params that don't need all_reduce (FSDP or single replica)
374 if None in groups_by_comm:
375 for hsdp_param in groups_by_comm[None]:
376 logger.debug(
377 "post_backward module=%s launch=reduce_scatter param=%s all_reduce=False",
378 self,
379 hsdp_param,
380 )
381 hsdp_param.reduce_scatter_grad(
382 reduce_op=self.reduce_op_type,
383 )
384 self.scheduler_ctx.pre_reduce_scatter_params.append(hsdp_param)
386 # Handle params that need all_reduce (HSDP with multiple replicas)
387 for group_key, hsdp_params in groups_by_comm.items():
388 if group_key is None:
389 continue
391 # Create AllReduceParamGroup for fused all-reduce
392 group = AllReduceParamGroup(
393 replicate_group=hsdp_params[0].mesh_info.replicate_process_group,
394 hsdp_params=hsdp_params,
395 reduce_op=self.reduce_op_type,
396 )
398 # Allocate fused buffer with 512-byte alignment
399 group.allocate_fused_buffer(self.device)
401 # Issue reduce_scatter with output directly into fused buffer views
402 logger.debug(
403 "post_backward module=%s launch=fused_reduce_scatter group_params=%s",
404 self,
405 hsdp_params,
406 )
407 for idx, hsdp_param in enumerate(hsdp_params):
408 buffer_view = group.get_param_buffer_view(idx)
409 hsdp_param.reduce_scatter_grad(
410 reduce_op=self.reduce_op_type,
411 output_buffer=buffer_view,
412 )
414 # Save the group so the next module hook can wait RS and launch AR.
415 self.scheduler_ctx.pre_all_reduce_groups.append(group)
417 def _wait_prev_reduce_scatter(self) -> List[AllReduceParamGroup]:
418 """Step 1: wait prev reduce_scatter.
420 This enables overlapping:
421 - Layer N-1's reduce_scatter wait with Layer N's backward compute
423 Returns:
424 List of previous AllReduceParamGroups (one per communication group).
425 """
426 if self.scheduler_ctx.pre_all_reduce_groups:
427 prev_groups = list(self.scheduler_ctx.pre_all_reduce_groups)
428 self.scheduler_ctx.pre_all_reduce_groups.clear()
429 for prev_group in prev_groups:
430 logger.debug(
431 "post_backward module=%s wait=fused_reduce_scatter group_params=%s",
432 self,
433 prev_group.hsdp_params,
434 )
435 for hsdp_param in prev_group.hsdp_params:
436 hsdp_param.reduce_scatter_output()
437 hsdp_param.clear_reduce_scatter_output()
438 if hsdp_param.unsharded_accumulated_grad_data is not None:
439 hsdp_param.unsharded_accumulated_grad = None
440 elif hsdp_param.unsharded_param.grad is not None:
441 hsdp_param.unsharded_param.grad = None
442 return prev_groups
443 return []
445 def _issue_prev_fused_all_reduce(self, prev_groups: List[AllReduceParamGroup]) -> None:
446 """Step 4: issue the previous module's fused all-reduce asynchronously.
448 The all-reduce work is collected in ``pending_all_reduce_groups``
449 and is waited in the root backward hook.
451 Args:
452 prev_groups: Previous parameter groups whose all-reduce should be issued.
453 """
454 for prev_group in prev_groups:
455 prev_group.accumulate_reduce_partial_outputs()
456 logger.debug(
457 "post_backward module=%s launch=fused_all_reduce group_params=%s",
458 self,
459 prev_group.hsdp_params,
460 )
461 prev_group.issue_async_allreduce()
462 self.scheduler_ctx.pending_all_reduce_groups.append(prev_group)
464 def _wait_prev_reduce_scatter_without_all_reduce(self) -> None:
465 """Wait previous RS outputs that do not enter a replicate all-reduce.
467 When the current micro-step disables all-reduce, outputs accumulate in
468 ``reduce_partial_output`` without being cast or applied to the parameter.
469 On the final synchronized micro-step, the partial result is merged into
470 the current RS output and retained for root-hook finalization.
471 """
472 while self.scheduler_ctx.pre_reduce_scatter_params:
473 pre_hsdp_param = self.scheduler_ctx.pre_reduce_scatter_params.pop(0)
474 logger.debug(
475 "post_backward module=%s wait=reduce_scatter param=%s",
476 self,
477 pre_hsdp_param,
478 )
479 reduced_grad = pre_hsdp_param.reduce_scatter_output()
480 if not self.requires_all_reduce:
481 if pre_hsdp_param.reduce_partial_output is None:
482 pre_hsdp_param.reduce_partial_output = reduced_grad
483 else:
484 pre_hsdp_param.reduce_partial_output.add_(reduced_grad)
485 pre_hsdp_param.clear_reduce_scatter_output()
486 elif pre_hsdp_param.reduce_partial_output is not None:
487 reduced_grad.add_(pre_hsdp_param.reduce_partial_output)
488 pre_hsdp_param.reduce_partial_output = None
490 if pre_hsdp_param.unsharded_accumulated_grad_data is not None:
491 pre_hsdp_param.unsharded_accumulated_grad = None
492 elif pre_hsdp_param.unsharded_param.grad is not None:
493 pre_hsdp_param.unsharded_param.grad = None
495 def wait_and_split_all_reduce_work_groups(self) -> None:
496 """Wait fused all-reduce work and expose each parameter result."""
497 for group in self.scheduler_ctx.pending_all_reduce_groups:
498 logger.debug(
499 "post_backward module=%s wait=fused_all_reduce group_params=%s",
500 self,
501 group.hsdp_params,
502 )
503 group.wait_and_split_grads()
504 self.scheduler_ctx.pending_all_reduce_groups.clear()
506 def reset_iter_state(self) -> None:
507 """Clear Torch communication bookkeeping without clearing optimizer gradients."""
508 self.scheduler_ctx.pre_reduce_scatter_params.clear()
509 self.scheduler_ctx.pre_all_reduce_params.clear()
510 self.scheduler_ctx.pre_all_reduce_groups.clear()
511 self.scheduler_ctx.pending_all_reduce_groups.clear()
512 if self.param_group is not None:
513 self.param_group.reset_iter_state()
514 for hsdp_param in self.hsdp_params:
515 hsdp_param.allgather_comm_ctx.allgather_handle = None
516 hsdp_param.allgather_comm_ctx.allgather_output = None
517 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_handle = None
518 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output = None
519 hsdp_param.all_reduce_comm_ctx.all_reduce_handle = None
520 hsdp_param.all_reduce_comm_ctx.all_reduce_output = None
521 hsdp_param.reduce_partial_output = None
522 hsdp_param.unsharded_accumulated_grad = None
523 hsdp_param._grad = None
524 if hsdp_param.unsharded_param_buffers:
525 hsdp_param.unsharded_param.grad = None
527 def set_requires_grad_sync(self, requires_grad_sync: bool) -> None:
528 """set requires grad sync flag to control gradient sync."""
529 self.reduce_grads = requires_grad_sync
531 def set_reduce_op_type(self, reduce_op_type: str) -> None:
532 """set reduce op type for gradient reduction."""
533 fsdp_support_reduce_op = {
534 "sum": torch.distributed.ReduceOp.SUM,
535 "avg": torch.distributed.ReduceOp.AVG,
536 }
537 reduce_op = reduce_op_type.lower().strip() if isinstance(reduce_op_type, str) else reduce_op_type
538 reduce_op_value = fsdp_support_reduce_op.get(reduce_op)
539 if reduce_op_value is None:
540 raise ValueError(
541 f"Unsupported reduce op type {reduce_op_type}, "
542 f"supported types are {list(fsdp_support_reduce_op.keys())}"
543 )
544 self.reduce_op_type = reduce_op_value
546 def shard(self) -> None:
547 """change parameters to sharded state"""
548 logger.debug(
549 "action=reshard module=%s params=%s",
550 self,
551 self.hsdp_params,
552 )
553 if self.is_shard:
554 return
555 for param in self.hsdp_params:
556 param.to_sharded()
557 self.is_shard = True
559 def unshard(self, async_op: bool = False) -> None:
560 """change parameters to unsharded state"""
561 logger.debug(
562 "action=unshard module=%s async_op=%s params=%s",
563 self,
564 async_op,
565 self.hsdp_params,
566 )
567 if not self.is_shard:
568 return
570 if self.comm_fusion_policy.enable_comm_fusion and self.param_group is not None:
571 self.param_group.unshard(async_op)
572 else:
573 for param in self.hsdp_params:
574 param.unshard(async_op)
575 if not async_op:
576 self.wait_for_unshard()
578 def prefetch(self) -> None:
579 """prefetch unsharded parameters"""
580 logger.debug(
581 "action=prefetch module=%s params=%s",
582 self,
583 self.hsdp_params,
584 )
585 self.unshard(async_op=True)
587 def wait_for_unshard(self) -> None:
588 """wait for all unshard parameters"""
589 logger.debug(
590 "action=wait_unshard module=%s params=%s",
591 self,
592 self.hsdp_params,
593 )
594 if not self.is_shard:
595 return
597 if self.comm_fusion_policy.enable_comm_fusion and self.param_group is not None:
598 self.param_group.wait_for_unshard()
599 else:
600 for param in self.hsdp_params:
601 param.wait_for_unshard()
602 self.is_shard = False
604 def set_gradient_scaling_factor(self, factor):
605 """Propagate the gradient scaling factor to the layer that applies it.
607 The factor is consumed on the reduce input: ``param_group.foreach_reduce``
608 for the fused (comm_fusion) path, or per-parameter ``reduce_scatter_grad``
609 / ``all_reduce_grad`` otherwise. The state does not hold a copy.
610 """
611 if self.param_group is not None:
612 self.param_group.gradient_scaling_factor = factor
613 else:
614 for hsdp_param in self.hsdp_params:
615 hsdp_param.gradient_scaling_factor = factor
617 def set_requires_all_reduce(self, requires_all_reduce: bool) -> None:
618 """Propagate the HSDP all-reduce switch to the active communication path."""
619 self.requires_all_reduce = requires_all_reduce
620 if self.param_group is not None:
621 self.param_group.requires_all_reduce = requires_all_reduce
623 def _sync_current_stream_if_needed(self, need_synchronize):
624 if need_synchronize:
625 if self.device.type == "npu":
626 torch.npu.current_stream().synchronize()
627 elif self.device.type == "cuda":
628 torch.cuda.current_stream().synchronize()
629 else:
630 raise NotImplementedError(
631 f"Unsupported device type {self.device.type} for synchronization after CPU offload."
632 )