Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / hsdp_state.py: 83%
296 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"""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 mixed-precision dtypes for all managed parameters."""
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 hsdp_param.unsharded_param_buffers
348 or not hsdp_param.sharded_param.requires_grad
349 or (
350 hsdp_param.unsharded_param.grad is None
351 and hsdp_param.unsharded_accumulated_grad_data is None
352 )
353 )
354 if skip_param:
355 continue
356 params_to_reduce.append(hsdp_param)
358 if not params_to_reduce:
359 return
361 # Group by replicate process group and reduction dtype so every fused
362 # all-reduce buffer has one communication group and one element type.
363 groups_by_comm = defaultdict(list)
364 for hsdp_param in params_to_reduce:
365 if self.requires_all_reduce and hsdp_param.replicate_world_size > 1:
366 replicate_process_group = hsdp_param.mesh_info.replicate_process_group
367 group_key = (id(replicate_process_group), hsdp_param.reduce_comm_dtype())
368 groups_by_comm[group_key].append(hsdp_param)
369 else:
370 groups_by_comm[None].append(hsdp_param)
372 # Handle params that don't need all_reduce (FSDP or single replica)
373 if None in groups_by_comm:
374 for hsdp_param in groups_by_comm[None]:
375 logger.debug(
376 "post_backward module=%s launch=reduce_scatter param=%s all_reduce=False",
377 self,
378 hsdp_param,
379 )
380 hsdp_param.reduce_scatter_grad(
381 reduce_op=self.reduce_op_type,
382 )
383 self.scheduler_ctx.pre_reduce_scatter_params.append(hsdp_param)
385 # Handle params that need all_reduce (HSDP with multiple replicas)
386 for group_key, hsdp_params in groups_by_comm.items():
387 if group_key is None:
388 continue
390 # Create AllReduceParamGroup for fused all-reduce
391 group = AllReduceParamGroup(
392 replicate_group=hsdp_params[0].mesh_info.replicate_process_group,
393 hsdp_params=hsdp_params,
394 reduce_op=self.reduce_op_type,
395 )
397 # Allocate fused buffer with 512-byte alignment
398 group.allocate_fused_buffer(self.device)
400 # Issue reduce_scatter with output directly into fused buffer views
401 logger.debug(
402 "post_backward module=%s launch=fused_reduce_scatter group_params=%s",
403 self,
404 hsdp_params,
405 )
406 for idx, hsdp_param in enumerate(hsdp_params):
407 buffer_view = group.get_param_buffer_view(idx)
408 hsdp_param.reduce_scatter_grad(
409 reduce_op=self.reduce_op_type,
410 output_buffer=buffer_view,
411 )
413 # Save the group so the next module hook can wait RS and launch AR.
414 self.scheduler_ctx.pre_all_reduce_groups.append(group)
416 def _wait_prev_reduce_scatter(self) -> List[AllReduceParamGroup]:
417 """Step 1: wait prev reduce_scatter.
419 This enables overlapping:
420 - Layer N-1's reduce_scatter wait with Layer N's backward compute
422 Returns:
423 List of previous AllReduceParamGroups (one per communication group).
424 """
425 if self.scheduler_ctx.pre_all_reduce_groups:
426 prev_groups = list(self.scheduler_ctx.pre_all_reduce_groups)
427 self.scheduler_ctx.pre_all_reduce_groups.clear()
428 for prev_group in prev_groups:
429 logger.debug(
430 "post_backward module=%s wait=fused_reduce_scatter group_params=%s",
431 self,
432 prev_group.hsdp_params,
433 )
434 for hsdp_param in prev_group.hsdp_params:
435 hsdp_param.reduce_scatter_output()
436 hsdp_param.clear_reduce_scatter_output()
437 if hsdp_param.unsharded_accumulated_grad_data is not None:
438 hsdp_param.unsharded_accumulated_grad = None
439 elif hsdp_param.unsharded_param.grad is not None:
440 hsdp_param.unsharded_param.grad = None
441 return prev_groups
442 return []
444 def _issue_prev_fused_all_reduce(self, prev_groups: List[AllReduceParamGroup]) -> None:
445 """Step 4: issue the previous module's fused all-reduce asynchronously.
447 The all-reduce work is collected in ``pending_all_reduce_groups``
448 and is waited in the root backward hook.
450 Args:
451 prev_groups: Previous parameter groups whose all-reduce should be issued.
452 """
453 for prev_group in prev_groups:
454 prev_group.accumulate_reduce_partial_outputs()
455 logger.debug(
456 "post_backward module=%s launch=fused_all_reduce group_params=%s",
457 self,
458 prev_group.hsdp_params,
459 )
460 prev_group.issue_async_allreduce()
461 self.scheduler_ctx.pending_all_reduce_groups.append(prev_group)
463 def _wait_prev_reduce_scatter_without_all_reduce(self) -> None:
464 """Wait previous RS outputs that do not enter a replicate all-reduce.
466 When the current micro-step disables all-reduce, outputs accumulate in
467 ``reduce_partial_output`` without being cast or applied to the parameter.
468 On the final synchronized micro-step, the partial result is merged into
469 the current RS output and retained for root-hook finalization.
470 """
471 while self.scheduler_ctx.pre_reduce_scatter_params:
472 pre_hsdp_param = self.scheduler_ctx.pre_reduce_scatter_params.pop(0)
473 logger.debug(
474 "post_backward module=%s wait=reduce_scatter param=%s",
475 self,
476 pre_hsdp_param,
477 )
478 reduced_grad = pre_hsdp_param.reduce_scatter_output()
479 if not self.requires_all_reduce:
480 if pre_hsdp_param.reduce_partial_output is None:
481 pre_hsdp_param.reduce_partial_output = reduced_grad
482 else:
483 pre_hsdp_param.reduce_partial_output.add_(reduced_grad)
484 pre_hsdp_param.clear_reduce_scatter_output()
485 elif pre_hsdp_param.reduce_partial_output is not None:
486 reduced_grad.add_(pre_hsdp_param.reduce_partial_output)
487 pre_hsdp_param.reduce_partial_output = None
489 if pre_hsdp_param.unsharded_accumulated_grad_data is not None:
490 pre_hsdp_param.unsharded_accumulated_grad = None
491 elif pre_hsdp_param.unsharded_param.grad is not None:
492 pre_hsdp_param.unsharded_param.grad = None
494 def wait_and_split_all_reduce_work_groups(self) -> None:
495 """Wait fused all-reduce work and expose each parameter result."""
496 for group in self.scheduler_ctx.pending_all_reduce_groups:
497 logger.debug(
498 "post_backward module=%s wait=fused_all_reduce group_params=%s",
499 self,
500 group.hsdp_params,
501 )
502 group.wait_and_split_grads()
503 self.scheduler_ctx.pending_all_reduce_groups.clear()
505 def reset_iter_state(self) -> None:
506 """Clear Torch communication bookkeeping without clearing optimizer gradients."""
507 self.scheduler_ctx.pre_reduce_scatter_params.clear()
508 self.scheduler_ctx.pre_all_reduce_params.clear()
509 self.scheduler_ctx.pre_all_reduce_groups.clear()
510 self.scheduler_ctx.pending_all_reduce_groups.clear()
511 if self.param_group is not None:
512 self.param_group.reset_iter_state()
513 for hsdp_param in self.hsdp_params:
514 hsdp_param.allgather_comm_ctx.allgather_handle = None
515 hsdp_param.allgather_comm_ctx.allgather_output = None
516 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_handle = None
517 hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output = None
518 hsdp_param.all_reduce_comm_ctx.all_reduce_handle = None
519 hsdp_param.all_reduce_comm_ctx.all_reduce_output = None
520 hsdp_param.reduce_partial_output = None
521 hsdp_param.unsharded_accumulated_grad = None
522 hsdp_param._grad = None
523 if hsdp_param.unsharded_param_buffers:
524 hsdp_param.unsharded_param.grad = None
526 def set_requires_grad_sync(self, requires_grad_sync: bool) -> None:
527 """set requires grad sync flag to control gradient sync."""
528 self.reduce_grads = requires_grad_sync
530 def set_reduce_op_type(self, reduce_op_type: str) -> None:
531 """set reduce op type for gradient reduction."""
532 fsdp_support_reduce_op = {
533 "sum": torch.distributed.ReduceOp.SUM,
534 "avg": torch.distributed.ReduceOp.AVG,
535 }
536 reduce_op = reduce_op_type.lower().strip() if isinstance(reduce_op_type, str) else reduce_op_type
537 reduce_op_value = fsdp_support_reduce_op.get(reduce_op)
538 if reduce_op_value is None:
539 raise ValueError(
540 f"Unsupported reduce op type {reduce_op_type}, "
541 f"supported types are {list(fsdp_support_reduce_op.keys())}"
542 )
543 self.reduce_op_type = reduce_op_value
545 def shard(self) -> None:
546 """change parameters to sharded state"""
547 logger.debug(
548 "action=reshard module=%s params=%s",
549 self,
550 self.hsdp_params,
551 )
552 if self.is_shard:
553 return
554 for param in self.hsdp_params:
555 param.to_sharded()
556 self.is_shard = True
558 def unshard(self, async_op: bool = False) -> None:
559 """change parameters to unsharded state"""
560 logger.debug(
561 "action=unshard module=%s async_op=%s params=%s",
562 self,
563 async_op,
564 self.hsdp_params,
565 )
566 if not self.is_shard:
567 return
569 if self.comm_fusion_policy.enable_comm_fusion and self.param_group is not None:
570 self.param_group.unshard(async_op)
571 else:
572 for param in self.hsdp_params:
573 param.unshard(async_op)
574 if not async_op:
575 self.wait_for_unshard()
577 def prefetch(self) -> None:
578 """prefetch unsharded parameters"""
579 logger.debug(
580 "action=prefetch module=%s params=%s",
581 self,
582 self.hsdp_params,
583 )
584 self.unshard(async_op=True)
586 def wait_for_unshard(self) -> None:
587 """wait for all unshard parameters"""
588 logger.debug(
589 "action=wait_unshard module=%s params=%s",
590 self,
591 self.hsdp_params,
592 )
593 if not self.is_shard:
594 return
596 if self.comm_fusion_policy.enable_comm_fusion and self.param_group is not None:
597 self.param_group.wait_for_unshard()
598 else:
599 for param in self.hsdp_params:
600 param.wait_for_unshard()
601 self.is_shard = False
603 def set_gradient_scaling_factor(self, factor):
604 """Propagate the gradient scaling factor to the layer that applies it.
606 The factor is consumed on the reduce input: ``param_group.foreach_reduce``
607 for the fused (comm_fusion) path, or per-parameter ``reduce_scatter_grad``
608 / ``all_reduce_grad`` otherwise. The state does not hold a copy.
609 """
610 if self.param_group is not None:
611 self.param_group.gradient_scaling_factor = factor
612 else:
613 for hsdp_param in self.hsdp_params:
614 hsdp_param.gradient_scaling_factor = factor
616 def set_requires_all_reduce(self, requires_all_reduce: bool) -> None:
617 """Propagate the HSDP all-reduce switch to the active communication path."""
618 self.requires_all_reduce = requires_all_reduce
619 if self.param_group is not None:
620 self.param_group.requires_all_reduce = requires_all_reduce
622 def _sync_current_stream_if_needed(self, need_synchronize):
623 if need_synchronize:
624 if self.device.type == "npu":
625 torch.npu.current_stream().synchronize()
626 elif self.device.type == "cuda":
627 torch.cuda.current_stream().synchronize()
628 else:
629 raise NotImplementedError(
630 f"Unsupported device type {self.device.type} for synchronization after CPU offload."
631 )