Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / hsdp_scheduler.py: 70%
388 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 scheduler"""
16import functools
17import inspect
18from typing import Any, Callable, List, Mapping, Optional, ParamSpec, Tuple, TypeVar, Union
20import torch
21from torch import nn
22from torch.autograd import Variable
23from torch.utils._pytree import tree_flatten, tree_unflatten
25from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
26from hyper_parallel.core.fully_shard.hook_function import PostBackwardFunction
27from hyper_parallel.core.fully_shard.hsdp_utils import (
28 FSDPSchedulerState,
29 get_managed_modules_parameters,
30)
31from hyper_parallel.core.fully_shard.utils import (
32 CommFusionPolicy,
33 SourceShardMetaInfo,
34 apply_to_tensors,
35 cast_fp_tensor,
36 parameters_dict,
37 profiler_record,
38)
39from hyper_parallel.tools.logging import get_logger
41logger = get_logger("FSDP")
43nn.Module = nn.Module
45_P = ParamSpec("_P")
46_R = TypeVar("_R")
49def _dynamo_disable(func: Callable[_P, _R]) -> Callable[_P, _R]:
50 """Disable Dynamo tracing while an FSDP runtime hook executes."""
52 @functools.wraps(func)
53 def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
54 return torch._dynamo.disable(
55 func,
56 recursive=True,
57 )(*args, **kwargs)
59 return wrapper
62class ParamGroupCommCtx:
63 """Track the in-flight parameter-group communication for one module tree."""
65 def __init__(self) -> None:
66 self.pre_param_group = None
67 self.all_reduce_param_group = None
68 self.comm_handle = None
71class HSDPSchedulerContext:
72 """Share scheduler and backward-pipeline state within one HSDP module tree."""
74 def __init__(self) -> None:
75 self.is_last_backward: bool = True
76 self.root_module = None
77 # Compile tracing may enter the root more than once; initialize shared
78 # parameter state only on the first real forward.
79 self.lazy_init_done: bool = False
80 self.root_bp_state = False
81 # all_hsdp_schedulers (for one module tree structure), deduplicated at
82 # registration time because ``fully_shard`` may be given a module list
83 # whose modules share one scheduler.
84 self.all_hsdp_schedulers = []
85 # Parameter FQNs are initialized once after all schedulers share this context.
86 self._param_fqn_initialized = False
87 # Backward pipeline queues shared only by schedulers in this module tree.
88 self.pre_reduce_scatter_params = []
89 self.pre_all_reduce_params = []
90 self.pre_direct_all_reduce_grads = []
91 self.pre_all_reduce_groups = []
92 self.pending_all_reduce_groups = []
93 self.param_group_comm_ctx = ParamGroupCommCtx()
96class HSDPSchedulerV2:
97 """HSDPScheduler is used to scheduler hsdp"""
99 def __init__(
100 self,
101 cell: Union[nn.Module, Tuple[nn.Module, ...]],
102 mesh,
103 reshard_after_forward,
104 shard_placement_fn,
105 mp_policy,
106 offload_policy,
107 ignored_params,
108 replicate_params,
109 device,
110 comm_fusion,
111 comm_fusion_zero_copy=False,
112 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]] = None,
113 ):
114 """init hsdp scheduler.
116 Args:
117 cell: A single module or tuple of modules to manage as one FSDP unit.
118 mesh: Explicit data-parallel device mesh.
119 reshard_after_forward: Whether to reshard parameters after forward.
120 shard_placement_fn: Optional function selecting the parameter shard dimension.
121 mp_policy: Mixed-precision policy.
122 offload_policy: Parameter offload policy.
123 ignored_params: Parameters excluded from fully_shard management.
124 replicate_params: Managed parameters that remain replicated.
125 device: Target device for parameters.
126 comm_fusion: Communication fusion configuration.
127 comm_fusion_zero_copy: Whether fused communication writes into views.
128 source_shard_infos: Optional per-parameter TP/EP source metadata.
129 """
130 self.modules = (cell,) if isinstance(cell, nn.Module) else tuple(cell)
131 self.cell = self.modules[0]
132 self.mesh: DeviceMesh = mesh
133 self.shard_placement_fn = shard_placement_fn
134 self.mp_policy = mp_policy
135 self.offload_policy = offload_policy
136 self.comm_fusion_policy = CommFusionPolicy(comm_fusion, comm_fusion_zero_copy)
137 self.ignored_params = ignored_params
138 self.replicate_params = replicate_params
139 self.device = device
140 self.reshard_after_forward = reshard_after_forward
141 self.source_shard_infos = source_shard_infos
142 self.scheduler_state = None
143 self.forward_prefetch_cells = []
144 self.backward_prefetch_cells = []
145 self._backup_forward_fetch = None
146 # Flag to identify root module.
147 self._is_root = True
148 # module and its all sub-modules share one same 'HSDPSchedulerContext'
149 self.scheduler_ctx = HSDPSchedulerContext()
150 # When ``fully_shard`` is given multiple root modules, forward pre/post hooks coordinate
151 # so unshard / PostBackward / reshard run once per forward (aligned with PyTorch FSDP2).
152 self._fsdp_group_post_pending: Optional[set] = set() if len(self.modules) > 1 else None
153 self._new_cell_state()
154 self._register_hooks()
156 def _new_cell_state(self):
157 """Create the HSDP state that owns this scheduler's managed parameters."""
158 # Imported lazily: ``state`` pulls in ``param_group``, which imports this
159 # module for ``ParamGroupCommCtx``.
160 from hyper_parallel.core.fully_shard.hsdp_state import ( # pylint: disable=C0415
161 HSDPStateV2,
162 )
164 self.hsdp_state = HSDPStateV2(
165 self.modules,
166 self.mesh,
167 self.shard_placement_fn,
168 self.comm_fusion_policy,
169 self.mp_policy,
170 self.offload_policy,
171 self.ignored_params,
172 self.replicate_params,
173 self.scheduler_ctx,
174 self.device,
175 source_shard_infos=self.source_shard_infos,
176 )
178 def _register_hooks(self):
179 """Register the forward and backward hooks driving shard/unshard."""
180 self._register_forward_backward_hooks()
182 def _get_managed_params(self):
183 """Return deduplicated parameters from all managed modules."""
184 return get_managed_modules_parameters(self.modules, self.ignored_params)
186 def set_reshard_after_forward(self, reshard_after_forward: bool) -> None:
187 """Set reshard_after_forward flag.
189 Args:
190 reshard_after_forward: Whether to reshard parameters after forward.
191 """
192 if not isinstance(reshard_after_forward, bool):
193 raise ValueError(f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}")
194 self.reshard_after_forward = reshard_after_forward
196 def set_reshard_after_backward(self, reshard_after_backward: bool) -> None:
197 """Set reshard_after_backward flag.
199 Args:
200 reshard_after_backward: Whether to reshard after backward completes.
201 """
202 if not isinstance(reshard_after_backward, bool):
203 raise ValueError(f"reshard_after_backward should be a bool, got {type(reshard_after_backward)}")
204 if self.hsdp_state is not None:
205 self.hsdp_state.reshard_after_backward = reshard_after_backward
207 def set_requires_all_reduce(self, requires_all_reduce: bool) -> None:
208 """Set requires_all_reduce flag.
210 Args:
211 requires_all_reduce: Whether this unit participates in all-reduce.
212 """
213 if not isinstance(requires_all_reduce, bool):
214 raise ValueError(f"requires_all_reduce should be a bool, got {type(requires_all_reduce)}")
215 if self.hsdp_state is not None:
216 self.hsdp_state.set_requires_all_reduce(requires_all_reduce)
218 @_dynamo_disable
219 def reset_iter_state(self) -> None:
220 """Reset scheduler bookkeeping after a completed iteration."""
221 self.scheduler_ctx.root_bp_state = False
222 self.scheduler_ctx.pre_reduce_scatter_params.clear()
223 self.scheduler_ctx.pre_all_reduce_params.clear()
224 self.scheduler_ctx.pre_direct_all_reduce_grads.clear()
225 self.scheduler_ctx.pre_all_reduce_groups.clear()
226 self.scheduler_ctx.pending_all_reduce_groups.clear()
227 self.scheduler_ctx.param_group_comm_ctx.pre_param_group = None
228 self.scheduler_ctx.param_group_comm_ctx.all_reduce_param_group = None
229 self.scheduler_ctx.param_group_comm_ctx.comm_handle = None
230 self.scheduler_state = None
231 if self._fsdp_group_post_pending is not None:
232 self._fsdp_group_post_pending.clear()
233 self._restore_forward_prefetch_after_recompute()
234 self.hsdp_state.reset_iter_state()
236 def set_requires_grad_sync(self, requires_grad_sync: bool) -> None:
237 """Set flag controlling whether gradients are synchronized.
239 Args:
240 requires_grad_sync: When True, enable grad sync for this scheduler.
241 """
242 if not isinstance(requires_grad_sync, bool):
243 raise ValueError(f"requires_grad_sync should be a bool, got {type(requires_grad_sync)}")
244 self.hsdp_state.set_requires_grad_sync(requires_grad_sync)
246 # pylint: disable=W0613
247 def _hsdp_forward_pre_hook(self, cell, args, kwargs):
248 """Forward pre hook to unsharded parameter for forward process."""
249 logger.debug("hook=forward_pre enter module=%s", self.hsdp_state)
250 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
251 logger.debug("hook=forward_pre skip module=%s reason=pre_backward", self.hsdp_state)
252 return args, kwargs
253 if self.scheduler_ctx.root_bp_state:
254 self._disable_forward_prefetch_for_recompute()
255 if self.scheduler_ctx.root_module is None:
256 tree_ctx = self.scheduler_ctx
257 tree_ctx.root_module = self.cell
258 registered_schedulers = set()
259 for module_name, module in tree_ctx.root_module.named_modules():
260 from hyper_parallel.core.fully_shard.api import HSDPModule # pylint: disable=C0415
261 if isinstance(module, HSDPModule):
262 submod_scheduler = module.hsdp_scheduler
263 if submod_scheduler is None or id(submod_scheduler) in registered_schedulers:
264 continue
265 registered_schedulers.add(id(submod_scheduler))
266 if submod_scheduler.scheduler_ctx is not tree_ctx:
267 if submod_scheduler.scheduler_ctx.root_module is not None:
268 raise ValueError(
269 "HSDP scheduler already belongs to another initialized module tree"
270 )
271 submod_scheduler.scheduler_ctx = tree_ctx
272 submod_scheduler.hsdp_state.scheduler_ctx = tree_ctx
273 if submod_scheduler.hsdp_state.param_group is not None:
274 submod_scheduler.hsdp_state.param_group.comm_ctx = tree_ctx.param_group_comm_ctx
275 submod_scheduler._is_root = submod_scheduler is self # pylint: disable=protected-access
276 submod_scheduler.hsdp_state.module_name = module_name
277 tree_ctx.all_hsdp_schedulers.append(submod_scheduler)
279 self.scheduler_state = FSDPSchedulerState.PRE_FORWARD
280 if self._is_root and not self.scheduler_ctx.lazy_init_done:
281 self._init_params_fqn()
282 self._lazy_init_all_states()
283 self.scheduler_ctx.lazy_init_done = True
284 if self.mp_policy.cast_forward_inputs and self.mp_policy.param_dtype:
285 cast_fn = functools.partial(cast_fp_tensor, self.mp_policy.param_dtype)
286 args = apply_to_tensors(cast_fn, args)
287 kwargs = apply_to_tensors(cast_fn, kwargs)
288 with profiler_record(f"pre_forward unshard:{self.hsdp_state.module_name}"):
289 logger.debug("hook=forward_pre action=unshard module=%s", self.hsdp_state)
290 self.hsdp_state.unshard()
291 for prefetch_cell in self.forward_prefetch_cells:
292 prefetch_state = prefetch_cell.hsdp_scheduler.hsdp_state
293 with profiler_record(f"pre_forward prefetch:"
294 f"{prefetch_state.module_name}"):
295 logger.debug(
296 "hook=forward_pre action=prefetch module=%s target=%s",
297 self.hsdp_state,
298 prefetch_state,
299 )
300 prefetch_state.prefetch()
301 return args, kwargs
303 def _lazy_init_all_states(self):
304 if self._is_root and self.scheduler_ctx.root_module is not None:
305 for submod_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
306 hsdp_state = submod_scheduler.hsdp_state
307 if hsdp_state:
308 hsdp_state.lazy_init()
310 def _init_params_fqn(self): # pylint: disable=W0212
311 if not self._is_root or self.scheduler_ctx.root_module is None:
312 return
313 if self.scheduler_ctx._param_fqn_initialized: # pylint: disable=protected-access
314 return
315 # Build a map from original (sharded) parameter tensor to its HSDPParam wrapper.
316 param_to_hsdp_param = {}
317 for submod_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
318 hsdp_state = submod_scheduler.hsdp_state
319 if hsdp_state is None:
320 continue
321 for hsdp_param in hsdp_state.hsdp_params:
322 orig_param = hsdp_param.sharded_param
323 # Shared parameters: keep only the first mapping to preserve the
324 # first-seen FQN (consistent with the deduplication in _init_hsdp_params).
325 if orig_param not in param_to_hsdp_param:
326 param_to_hsdp_param[orig_param] = hsdp_param
328 # Walk the full parameter tree and assign FQNs; skip params already seen
329 # (shared-parameter deduplication: first name wins).
330 visited_params = set()
331 for param_name, parameter in parameters_dict(self.scheduler_ctx.root_module):
332 if parameter in visited_params:
333 continue
334 visited_params.add(parameter)
335 hsdp_param = param_to_hsdp_param.get(parameter)
336 if hsdp_param is not None:
337 hsdp_param._param_fqn = param_name # pylint: disable=W0212
338 self.scheduler_ctx._param_fqn_initialized = True # pylint: disable=protected-access
340 # pylint: disable=W0613, R1710
341 def _hsdp_forward_hook(self, cell, inputs, outputs):
342 """Forward hook to shard parameter for saving memory."""
343 logger.debug("hook=forward enter module=%s", self.hsdp_state)
344 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
345 logger.debug("hook=forward skip module=%s reason=pre_backward", self.hsdp_state)
346 return
347 self.scheduler_state = FSDPSchedulerState.FORWARD
348 if self.reshard_after_forward:
349 with profiler_record(f"forward reshard:{self.hsdp_state.module_name}"):
350 logger.debug("hook=forward action=reshard module=%s", self.hsdp_state)
351 self.hsdp_state.shard()
352 if self.mp_policy.output_dtype is not None:
353 outputs = apply_to_tensors(
354 functools.partial(cast_fp_tensor, self.mp_policy.output_dtype),
355 outputs,
356 )
357 return outputs
359 # pylint: disable=W0613
360 def _hsdp_backward_pre_hook(self, cell, grad_outputs):
361 """Backward pre hook to unsharded parameter for backward process."""
362 logger.debug("hook=backward_pre enter module=%s", self.hsdp_state)
363 self.scheduler_state = FSDPSchedulerState.PRE_BACKWARD
364 if self.reshard_after_forward:
365 with profiler_record(f"pre_backward unshard:{self.hsdp_state.module_name}"):
366 logger.debug("hook=backward_pre action=unshard module=%s", self.hsdp_state)
367 self.hsdp_state.unshard()
368 for prefetch_cell in self.backward_prefetch_cells:
369 prefetch_state = prefetch_cell.hsdp_scheduler.hsdp_state
370 with profiler_record(f"pre_backward prefetch:"
371 f"{prefetch_state.module_name}"):
372 logger.debug(
373 "hook=backward_pre action=prefetch module=%s target=%s",
374 self.hsdp_state,
375 prefetch_state,
376 )
377 prefetch_state.prefetch()
379 # pylint: disable=W0613
380 def _hsdp_backward_hook(self, cell, grad_inputs, grad_outputs):
381 """Backward hook to shard parameter for optimizer process or saving memory."""
382 logger.debug("hook=backward_hook enter module=%s", self.hsdp_state)
383 self.scheduler_state = FSDPSchedulerState.BACKWARD
384 with profiler_record(f"post_backward:{self.hsdp_state.module_name}"):
385 logger.debug("hook=backward_hook action=post_backward module=%s", self.hsdp_state)
386 self.hsdp_state.post_backward()
387 if self._fsdp_group_post_pending is not None:
388 self._fsdp_group_post_pending.clear()
390 # pylint: disable=W0613
391 @staticmethod
392 def _grouped_forward_pre_hook_skip(cell, args, kwargs) -> None:
393 """Return value when grouped pre-forward should not run (first module already did).
395 ``nn.Module`` pre-hooks use ``None`` for a no-op.
396 """
397 return None
399 @staticmethod
400 def _grouped_forward_post_hook_skip(outputs) -> None:
401 """Return value when grouped post-forward is deferred to a later module.
403 ``nn.Module`` forward hooks use ``None`` for a no-op.
404 """
405 return None
407 @_dynamo_disable
408 def _grouped_forward_pre_hook(self, cell, args, kwargs):
409 """Run FSDP pre-forward only for the first module in the group (PyTorch FSDP2-aligned)."""
410 pending = self._fsdp_group_post_pending
411 if pending is None:
412 return self._forward_pre_hook(cell, args, kwargs)
413 if len(pending) == 0:
414 pending.update(self.modules)
415 return self._forward_pre_hook(cell, args, kwargs)
416 return self._grouped_forward_pre_hook_skip(cell, args, kwargs)
418 def _make_grouped_forward_post_hook(self, mod):
419 """Build post-forward hook: last module in the group runs reshard + output backward hooks."""
421 def grouped_post_hook(cell, inputs, outputs):
422 pending = self._fsdp_group_post_pending
423 if pending is None:
424 return self._forward_hook(cell, inputs, outputs)
425 pending.discard(mod)
426 if len(pending) == 0:
427 return self._forward_hook(cell, inputs, outputs)
428 return self._grouped_forward_post_hook_skip(outputs)
430 return grouped_post_hook
432 def set_forward_prefetch_cells(self, hsdp_cell_list: List[Any]) -> None:
433 """Set cells prefetched during forward.
435 Args:
436 hsdp_cell_list: HSDP cells to prefetch ahead of forward.
437 """
438 self.forward_prefetch_cells = hsdp_cell_list
440 def set_backward_prefetch_cells(self, hsdp_cell_list: List[Any]) -> None:
441 """Set cells prefetched during backward.
443 Args:
444 hsdp_cell_list: HSDP cells to prefetch ahead of backward.
445 """
446 self.backward_prefetch_cells = hsdp_cell_list
448 def _disable_forward_prefetch_for_recompute(self) -> None:
449 """Temporarily disable forward prefetch during activation recompute."""
450 self._backup_forward_fetch = self.forward_prefetch_cells
451 self.forward_prefetch_cells = []
453 def _restore_forward_prefetch_after_recompute(self) -> None:
454 """Restore the saved forward prefetch list after activation recompute."""
455 if self._backup_forward_fetch is None:
456 return
457 self.forward_prefetch_cells = self._backup_forward_fetch
458 self._backup_forward_fetch = None
460 def _register_post_backward_hook(self, args, kwargs):
461 """Wrap forward args/kwargs through PostBackwardFunction to register backward hook."""
462 if not torch.is_grad_enabled():
463 return args, kwargs
464 args_list, args_spec = tree_flatten(args)
465 kwargs_list, kwargs_spec = tree_flatten(kwargs)
466 args_kwargs_list = list(args_list) + list(kwargs_list)
467 inp_tensor_indices: List[int] = []
468 inp_tensors: List[torch.Tensor] = []
469 for i, obj in enumerate(args_kwargs_list):
470 if torch.is_tensor(obj) and obj.requires_grad:
471 inp_tensor_indices.append(i)
472 inp_tensors.append(obj)
473 if len(inp_tensors) == 0:
474 return args, kwargs # no tensors that require gradients
475 processed_tensors = PostBackwardFunction.apply(self, *inp_tensors)
476 for inp_tensor_idx, processed_tensor in zip(inp_tensor_indices, processed_tensors):
477 args_kwargs_list[inp_tensor_idx] = processed_tensor
478 args_list = args_kwargs_list[: len(args_list)]
479 kwargs_list = args_kwargs_list[len(args_list) :]
480 args = tree_unflatten(args_list, args_spec)
481 kwargs = tree_unflatten(kwargs_list, kwargs_spec)
482 return args, kwargs
484 @_dynamo_disable
485 def _forward_pre_hook(self, cell, args, kwargs):
486 """Execute forward pre hook and set up backward hook."""
487 args, kwargs = self._hsdp_forward_pre_hook(cell, args, kwargs)
488 return self._register_post_backward_hook(args, kwargs)
490 def _register_backward_pre_hook(self, outputs) -> None:
491 """Register gradient hooks on all requires-grad outputs to trigger backward pre hook."""
492 flat_outputs, _ = tree_flatten(outputs)
493 for output in flat_outputs:
494 if isinstance(output, torch.Tensor) and output.requires_grad:
495 handle_ref = [None]
496 # pylint: disable=C0103, W0102
498 def wrapper_for_backward_pre_hook(grad, _handle_ref=handle_ref):
499 """Remove this hook after it fires to prevent accmulation"""
500 handle = _handle_ref[0]
501 if handle is not None:
502 handle.remove()
503 return self._backward_pre_hook(grad)
504 # pylint: enable=C0103, W0102
505 handle = output.register_hook(wrapper_for_backward_pre_hook)
506 handle_ref[0] = handle
508 @_dynamo_disable
509 def _forward_hook(self, cell, inputs, outputs): # pylint: disable=R1710
510 """Execute forward hook."""
511 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
512 return
513 self._register_backward_pre_hook(outputs)
514 if self.scheduler_ctx.root_bp_state:
515 self._restore_forward_prefetch_after_recompute()
516 return
517 return self._hsdp_forward_hook(cell, inputs, outputs)
519 # pylint: disable=W0212
520 @_dynamo_disable
521 def _backward_pre_hook(self, grad):
522 """Execute backward pre hook."""
523 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
524 return grad
525 if self._is_root:
526 Variable._execution_engine.queue_callback(self._root_backward_hook)
527 self.scheduler_ctx.root_bp_state = True
528 self._hsdp_backward_pre_hook(self.cell, None)
529 return grad
531 @_dynamo_disable
532 def _root_backward_hook(self):
533 """Drain all DP pipelines, then run final TP reduction and apply gradients."""
534 logger.debug("hook=root_backward_hook enter module=%s", self.hsdp_state)
535 self.check_state_prefetch_handles_consumed()
536 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
537 # let modules which are not triggered backward_hook launch backward communication.
538 hsdp_scheduler._backward_hook()
539 self.scheduler_ctx.root_bp_state = False
540 with torch.profiler.record_function(f"root_backward reduce:{self.hsdp_state.module_name}"):
541 logger.debug(
542 "hook=root_backward_hook action=final_reduce module=%s",
543 self.hsdp_state,
544 )
545 self._finalize_comm_fusion_reductions()
546 self._finalize_per_param_reductions()
547 self.launch_tp_replicate_reduce_and_apply()
549 def check_state_prefetch_handles_consumed(self) -> None:
550 """Check that all parameters in the module tree consumed their prefetch work.
552 Raises:
553 RuntimeError: If a parameter or fused bucket retains an all-gather
554 handle because the configured prefetch order is incorrect.
555 """
556 unconsumed_param_fqns = []
557 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
558 hsdp_state = hsdp_scheduler.hsdp_state
559 if hsdp_state is None:
560 continue
561 unconsumed_param_fqns.extend(
562 str(param) for param in hsdp_state.hsdp_params
563 if param.allgather_comm_ctx.allgather_handle is not None
564 )
565 if hsdp_state.param_group is not None:
566 for bucket in hsdp_state.param_group.all_gather_buckets:
567 result = bucket.all_gather_result
568 if result is not None and result.handle is not None:
569 unconsumed_param_fqns.extend(str(param) for param in bucket.hsdp_params)
570 if unconsumed_param_fqns:
571 param_fqns = ", ".join(dict.fromkeys(unconsumed_param_fqns))
572 raise RuntimeError(
573 f"Unconsumed prefetch handles for parameters: {param_fqns}. "
574 "The prefetch order is incorrect. Set forward/backward prefetch order "
575 "according to the model's actual execution order."
576 )
578 def _finalize_comm_fusion_reductions(self) -> None:
579 """Drain the comm_fusion=True RS/AR pipeline."""
580 comm_ctx = self.scheduler_ctx.param_group_comm_ctx
581 if comm_ctx.all_reduce_param_group is not None:
582 logger.debug(
583 "hook=root_backward_hook wait=comm_fusion_all_reduce module=%s",
584 self.hsdp_state,
585 )
586 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad()
587 comm_ctx.all_reduce_param_group = None
588 if comm_ctx.pre_param_group is not None:
589 logger.debug(
590 "hook=root_backward_hook wait=comm_fusion_reduce_scatter module=%s",
591 self.hsdp_state,
592 )
593 comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce()
594 comm_ctx.pre_param_group = None
595 if comm_ctx.all_reduce_param_group is not None:
596 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad()
597 comm_ctx.all_reduce_param_group = None
599 def _finalize_per_param_reductions(self) -> None:
600 """Drain the module-tree-local comm_fusion=False RS/AR queues."""
601 # A fused root may own non-fused children, so always drain the tree queues.
602 last_all_reduce_groups = self.hsdp_state._wait_prev_reduce_scatter()
603 self.hsdp_state._wait_prev_reduce_scatter_without_all_reduce()
604 self.hsdp_state._issue_prev_fused_all_reduce(last_all_reduce_groups)
605 self.hsdp_state.wait_and_split_all_reduce_work_groups()
607 def launch_tp_replicate_reduce_and_apply(self) -> None:
608 """Run final TP replicate reductions and apply gradients for all states."""
609 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
610 hsdp_state = hsdp_scheduler.hsdp_state
611 if hsdp_state is None:
612 continue
613 need_synchronize = False
614 for hsdp_param in hsdp_state.hsdp_params:
615 reduced_grad = hsdp_param.all_reduce_comm_ctx.all_reduce_output
616 if reduced_grad is None:
617 reduced_grad = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output
618 if reduced_grad is None:
619 continue
620 hsdp_param.all_reduce_source_replicate_grad_inplace(
621 reduced_grad,
622 hsdp_state.reduce_op_type,
623 )
624 need_synchronize = hsdp_param.apply_reduced_grad(reduced_grad) or need_synchronize
625 hsdp_param.clear_all_reduce_output()
626 hsdp_param.clear_reduce_scatter_output()
627 hsdp_state._sync_current_stream_if_needed(need_synchronize)
629 @_dynamo_disable
630 def _backward_hook(self):
631 """Execute backward hook."""
632 if self.scheduler_state == FSDPSchedulerState.BACKWARD:
633 return
634 self._hsdp_backward_hook(self.cell, None, None)
636 @staticmethod
637 def _register_forward_module_hook(mod, hook) -> None:
638 """Register forward hook; use ``always_call=True`` when supported (matches PyTorch FSDP)."""
639 sig = inspect.signature(mod.register_forward_hook)
640 if "always_call" in sig.parameters:
641 mod.register_forward_hook(hook, prepend=False, always_call=True)
642 else:
643 mod.register_forward_hook(hook, prepend=False)
645 def _register_forward_backward_hooks(self):
646 """Register module forward and backward hook on all managed modules."""
647 if self._fsdp_group_post_pending is None:
648 for mod in self.modules:
649 mod.register_forward_pre_hook(self._forward_pre_hook, with_kwargs=True)
650 mod.register_forward_hook(self._forward_hook)
651 return
652 for mod in self.modules:
653 mod.register_forward_pre_hook(self._grouped_forward_pre_hook, with_kwargs=True)
654 grouped_forward_hook = _dynamo_disable(self._make_grouped_forward_post_hook(mod))
655 self._register_forward_module_hook(mod, grouped_forward_hook)