Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / hsdp_scheduler.py: 73%
424 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 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 self.post_backward_final_callback_queued = False
82 # all_hsdp_schedulers (for one module tree structure), deduplicated at
83 # registration time because ``fully_shard`` may be given a module list
84 # whose modules share one scheduler.
85 self.all_hsdp_schedulers = []
86 # Parameter FQNs are initialized once after all schedulers share this context.
87 self._param_fqn_initialized = False
88 # Backward pipeline queues shared only by schedulers in this module tree.
89 self.pre_reduce_scatter_params = []
90 self.pre_all_reduce_params = []
91 self.pre_direct_all_reduce_grads = []
92 self.pre_all_reduce_groups = []
93 self.pending_all_reduce_groups = []
94 self.param_group_comm_ctx = ParamGroupCommCtx()
97class HSDPSchedulerV2:
98 """HSDPScheduler is used to scheduler hsdp"""
100 def __init__(
101 self,
102 cell: Union[nn.Module, Tuple[nn.Module, ...]],
103 mesh,
104 reshard_after_forward,
105 shard_placement_fn,
106 mp_policy,
107 offload_policy,
108 ignored_params,
109 replicate_params,
110 device,
111 comm_fusion,
112 comm_fusion_zero_copy=False,
113 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]] = None,
114 ):
115 """init hsdp scheduler.
117 Args:
118 cell: A single module or tuple of modules to manage as one FSDP unit.
119 mesh: Explicit data-parallel device mesh.
120 reshard_after_forward: Whether to reshard parameters after forward.
121 shard_placement_fn: Optional function selecting the parameter shard dimension.
122 mp_policy: Mixed-precision policy.
123 offload_policy: Parameter offload policy.
124 ignored_params: Parameters excluded from fully_shard management.
125 replicate_params: Managed parameters that remain replicated.
126 device: Target device for parameters.
127 comm_fusion: Communication fusion configuration.
128 comm_fusion_zero_copy: Whether fused communication writes into views.
129 source_shard_infos: Optional per-parameter TP/EP source metadata.
130 """
131 self.modules = (cell,) if isinstance(cell, nn.Module) else tuple(cell)
132 self.cell = self.modules[0]
133 self.mesh: DeviceMesh = mesh
134 self.shard_placement_fn = shard_placement_fn
135 self.mp_policy = mp_policy
136 self.offload_policy = offload_policy
137 self.comm_fusion_policy = CommFusionPolicy(comm_fusion, comm_fusion_zero_copy)
138 self.ignored_params = ignored_params
139 self.replicate_params = replicate_params
140 self.device = device
141 self.reshard_after_forward = reshard_after_forward
142 self.source_shard_infos = source_shard_infos
143 self.scheduler_state = None
144 self.forward_prefetch_cells = []
145 self.backward_prefetch_cells = []
146 self._backup_forward_fetch = None
147 # Flag to identify root module.
148 self._is_root = True
149 # module and its all sub-modules share one same 'HSDPSchedulerContext'
150 self.scheduler_ctx = HSDPSchedulerContext()
151 # When ``fully_shard`` is given multiple root modules, forward pre/post hooks coordinate
152 # so unshard / PostBackward / reshard run once per forward (aligned with PyTorch FSDP2).
153 self._fsdp_group_post_pending: Optional[set] = set() if len(self.modules) > 1 else None
154 self._new_cell_state()
155 self._register_hooks()
157 def _new_cell_state(self):
158 """Create the HSDP state that owns this scheduler's managed parameters."""
159 # Imported lazily: ``state`` pulls in ``param_group``, which imports this
160 # module for ``ParamGroupCommCtx``.
161 from hyper_parallel.core.fully_shard.hsdp_state import ( # pylint: disable=C0415
162 HSDPStateV2,
163 )
165 self.hsdp_state = HSDPStateV2(
166 self.modules,
167 self.mesh,
168 self.shard_placement_fn,
169 self.comm_fusion_policy,
170 self.mp_policy,
171 self.offload_policy,
172 self.ignored_params,
173 self.replicate_params,
174 self.scheduler_ctx,
175 self.device,
176 source_shard_infos=self.source_shard_infos,
177 )
179 def _register_hooks(self):
180 """Register the forward and backward hooks driving shard/unshard."""
181 self._register_forward_backward_hooks()
183 def _get_managed_params(self):
184 """Return deduplicated parameters from all managed modules."""
185 return get_managed_modules_parameters(self.modules, self.ignored_params)
187 def set_reshard_after_forward(self, reshard_after_forward: bool) -> None:
188 """Set reshard_after_forward flag.
190 Args:
191 reshard_after_forward: Whether to reshard parameters after forward.
192 """
193 if not isinstance(reshard_after_forward, bool):
194 raise ValueError(f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}")
195 self.reshard_after_forward = reshard_after_forward
197 def set_reshard_after_backward(self, reshard_after_backward: bool) -> None:
198 """Set reshard_after_backward flag.
200 Args:
201 reshard_after_backward: Whether to reshard after backward completes.
202 """
203 if not isinstance(reshard_after_backward, bool):
204 raise ValueError(f"reshard_after_backward should be a bool, got {type(reshard_after_backward)}")
205 if self.hsdp_state is not None:
206 self.hsdp_state.reshard_after_backward = reshard_after_backward
208 def set_requires_all_reduce(self, requires_all_reduce: bool) -> None:
209 """Set requires_all_reduce flag.
211 Args:
212 requires_all_reduce: Whether this unit participates in all-reduce.
213 """
214 if not isinstance(requires_all_reduce, bool):
215 raise ValueError(f"requires_all_reduce should be a bool, got {type(requires_all_reduce)}")
216 if self.hsdp_state is not None:
217 self.hsdp_state.set_requires_all_reduce(requires_all_reduce)
219 @_dynamo_disable
220 def reset_iter_state(self) -> None:
221 """Reset scheduler bookkeeping after a completed iteration."""
222 self.scheduler_ctx.root_bp_state = False
223 self.scheduler_ctx.post_backward_final_callback_queued = False
224 self.scheduler_ctx.pre_reduce_scatter_params.clear()
225 self.scheduler_ctx.pre_all_reduce_params.clear()
226 self.scheduler_ctx.pre_direct_all_reduce_grads.clear()
227 self.scheduler_ctx.pre_all_reduce_groups.clear()
228 self.scheduler_ctx.pending_all_reduce_groups.clear()
229 self.scheduler_ctx.param_group_comm_ctx.pre_param_group = None
230 self.scheduler_ctx.param_group_comm_ctx.all_reduce_param_group = None
231 self.scheduler_ctx.param_group_comm_ctx.comm_handle = None
232 self.scheduler_state = None
233 if self._fsdp_group_post_pending is not None:
234 self._fsdp_group_post_pending.clear()
235 self._restore_forward_prefetch_after_recompute()
236 self.hsdp_state.reset_iter_state()
238 def set_requires_grad_sync(self, requires_grad_sync: bool) -> None:
239 """Set flag controlling whether gradients are synchronized.
241 Args:
242 requires_grad_sync: When True, enable grad sync for this scheduler.
243 """
244 if not isinstance(requires_grad_sync, bool):
245 raise ValueError(f"requires_grad_sync should be a bool, got {type(requires_grad_sync)}")
246 self.hsdp_state.set_requires_grad_sync(requires_grad_sync)
248 # pylint: disable=W0613
249 def _hsdp_forward_pre_hook(self, cell, args, kwargs):
250 """Forward pre hook to unsharded parameter for forward process."""
251 logger.debug("hook=forward_pre enter module=%s", self.hsdp_state)
252 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
253 logger.debug("hook=forward_pre skip module=%s reason=pre_backward", self.hsdp_state)
254 return args, kwargs
255 if self.scheduler_ctx.root_bp_state:
256 self._disable_forward_prefetch_for_recompute()
257 if self.scheduler_ctx.root_module is None:
258 tree_ctx = self.scheduler_ctx
259 tree_ctx.root_module = self.cell
260 registered_schedulers = set()
261 for module_name, module in tree_ctx.root_module.named_modules():
262 from hyper_parallel.core.fully_shard.api import HSDPModule # pylint: disable=C0415
263 if isinstance(module, HSDPModule):
264 submod_scheduler = module.hsdp_scheduler
265 if submod_scheduler is None or id(submod_scheduler) in registered_schedulers:
266 continue
267 registered_schedulers.add(id(submod_scheduler))
268 if submod_scheduler.scheduler_ctx is not tree_ctx:
269 if submod_scheduler.scheduler_ctx.root_module is not None:
270 raise ValueError(
271 "HSDP scheduler already belongs to another initialized module tree"
272 )
273 submod_scheduler.scheduler_ctx = tree_ctx
274 submod_scheduler.hsdp_state.scheduler_ctx = tree_ctx
275 if submod_scheduler.hsdp_state.param_group is not None:
276 submod_scheduler.hsdp_state.param_group.comm_ctx = tree_ctx.param_group_comm_ctx
277 submod_scheduler._is_root = submod_scheduler is self # pylint: disable=protected-access
278 submod_scheduler.hsdp_state.module_name = module_name
279 tree_ctx.all_hsdp_schedulers.append(submod_scheduler)
281 self.scheduler_state = FSDPSchedulerState.PRE_FORWARD
282 if self._is_root and not self.scheduler_ctx.lazy_init_done:
283 self._init_params_fqn()
284 self._lazy_init_all_states()
285 self.scheduler_ctx.lazy_init_done = True
286 if self.mp_policy.cast_forward_inputs and self.mp_policy.param_dtype:
287 cast_fn = functools.partial(cast_fp_tensor, self.mp_policy.param_dtype)
288 args = apply_to_tensors(cast_fn, args)
289 kwargs = apply_to_tensors(cast_fn, kwargs)
290 with profiler_record(f"pre_forward unshard:{self.hsdp_state.module_name}"):
291 logger.debug("hook=forward_pre action=unshard module=%s", self.hsdp_state)
292 self.hsdp_state.unshard()
293 for prefetch_cell in self.forward_prefetch_cells:
294 prefetch_state = prefetch_cell.hsdp_scheduler.hsdp_state
295 with profiler_record(f"pre_forward prefetch:"
296 f"{prefetch_state.module_name}"):
297 logger.debug(
298 "hook=forward_pre action=prefetch module=%s target=%s",
299 self.hsdp_state,
300 prefetch_state,
301 )
302 prefetch_state.prefetch()
303 return args, kwargs
305 def _lazy_init_all_states(self):
306 if self._is_root and self.scheduler_ctx.root_module is not None:
307 for submod_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
308 hsdp_state = submod_scheduler.hsdp_state
309 if hsdp_state:
310 hsdp_state.lazy_init()
312 def _init_params_fqn(self): # pylint: disable=W0212
313 if not self._is_root or self.scheduler_ctx.root_module is None:
314 return
315 if self.scheduler_ctx._param_fqn_initialized: # pylint: disable=protected-access
316 return
317 # Build a map from original (sharded) parameter tensor to its HSDPParam wrapper.
318 param_to_hsdp_param = {}
319 for submod_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
320 hsdp_state = submod_scheduler.hsdp_state
321 if hsdp_state is None:
322 continue
323 for hsdp_param in hsdp_state.hsdp_params:
324 orig_param = hsdp_param.sharded_param
325 # Shared parameters: keep only the first mapping to preserve the
326 # first-seen FQN (consistent with the deduplication in _init_hsdp_params).
327 if orig_param not in param_to_hsdp_param:
328 param_to_hsdp_param[orig_param] = hsdp_param
330 # Walk the full parameter tree and assign FQNs; skip params already seen
331 # (shared-parameter deduplication: first name wins).
332 visited_params = set()
333 for param_name, parameter in parameters_dict(self.scheduler_ctx.root_module):
334 if parameter in visited_params:
335 continue
336 visited_params.add(parameter)
337 hsdp_param = param_to_hsdp_param.get(parameter)
338 if hsdp_param is not None:
339 hsdp_param._param_fqn = param_name # pylint: disable=W0212
340 self.scheduler_ctx._param_fqn_initialized = True # pylint: disable=protected-access
342 # pylint: disable=W0613, R1710
343 def _hsdp_forward_hook(self, cell, inputs, outputs):
344 """Forward hook to shard parameter for saving memory."""
345 logger.debug("hook=forward enter module=%s", self.hsdp_state)
346 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
347 logger.debug("hook=forward skip module=%s reason=pre_backward", self.hsdp_state)
348 return
349 self.scheduler_state = FSDPSchedulerState.FORWARD
350 if self.reshard_after_forward:
351 with profiler_record(f"forward reshard:{self.hsdp_state.module_name}"):
352 logger.debug("hook=forward action=reshard module=%s", self.hsdp_state)
353 self.hsdp_state.shard()
354 if self.mp_policy.output_dtype is not None:
355 outputs = apply_to_tensors(
356 functools.partial(cast_fp_tensor, self.mp_policy.output_dtype),
357 outputs,
358 )
359 return outputs
361 # pylint: disable=W0613
362 def _hsdp_backward_pre_hook(self, cell, grad_outputs):
363 """Backward pre hook to unsharded parameter for backward process."""
364 logger.debug("hook=backward_pre enter module=%s", self.hsdp_state)
365 self.scheduler_state = FSDPSchedulerState.PRE_BACKWARD
366 if self.reshard_after_forward:
367 with profiler_record(f"pre_backward unshard:{self.hsdp_state.module_name}"):
368 logger.debug("hook=backward_pre action=unshard module=%s", self.hsdp_state)
369 self.hsdp_state.unshard()
370 for prefetch_cell in self.backward_prefetch_cells:
371 prefetch_state = prefetch_cell.hsdp_scheduler.hsdp_state
372 with profiler_record(f"pre_backward prefetch:"
373 f"{prefetch_state.module_name}"):
374 logger.debug(
375 "hook=backward_pre action=prefetch module=%s target=%s",
376 self.hsdp_state,
377 prefetch_state,
378 )
379 prefetch_state.prefetch()
381 # pylint: disable=W0613
382 def _hsdp_backward_hook(self, cell, grad_inputs, grad_outputs):
383 """Backward hook to shard parameter for optimizer process or saving memory."""
384 logger.debug("hook=backward_hook enter module=%s", self.hsdp_state)
385 self.scheduler_state = FSDPSchedulerState.BACKWARD
386 with profiler_record(f"post_backward:{self.hsdp_state.module_name}"):
387 logger.debug("hook=backward_hook action=post_backward module=%s", self.hsdp_state)
388 self.hsdp_state.post_backward()
389 if self._fsdp_group_post_pending is not None:
390 self._fsdp_group_post_pending.clear()
392 # pylint: disable=W0613
393 @staticmethod
394 def _grouped_forward_pre_hook_skip(cell, args, kwargs) -> None:
395 """Return value when grouped pre-forward should not run (first module already did).
397 ``nn.Module`` pre-hooks use ``None`` for a no-op.
398 """
399 return None
401 @staticmethod
402 def _grouped_forward_post_hook_skip(outputs) -> None:
403 """Return value when grouped post-forward is deferred to a later module.
405 ``nn.Module`` forward hooks use ``None`` for a no-op.
406 """
407 return None
409 @_dynamo_disable
410 def _grouped_forward_pre_hook(self, cell, args, kwargs):
411 """Run FSDP pre-forward only for the first module in the group (PyTorch FSDP2-aligned)."""
412 pending = self._fsdp_group_post_pending
413 if pending is None:
414 return self._forward_pre_hook(cell, args, kwargs)
415 if len(pending) == 0:
416 pending.update(self.modules)
417 return self._forward_pre_hook(cell, args, kwargs)
418 return self._grouped_forward_pre_hook_skip(cell, args, kwargs)
420 def _make_grouped_forward_post_hook(self, mod):
421 """Build post-forward hook: last module in the group runs reshard + output backward hooks."""
423 def grouped_post_hook(cell, inputs, outputs):
424 pending = self._fsdp_group_post_pending
425 if pending is None:
426 return self._forward_hook(cell, inputs, outputs)
427 pending.discard(mod)
428 if len(pending) == 0:
429 return self._forward_hook(cell, inputs, outputs)
430 return self._grouped_forward_post_hook_skip(outputs)
432 return grouped_post_hook
434 def set_forward_prefetch_cells(self, hsdp_cell_list: List[Any]) -> None:
435 """Set cells prefetched during forward.
437 Args:
438 hsdp_cell_list: HSDP cells to prefetch ahead of forward.
439 """
440 self.forward_prefetch_cells = hsdp_cell_list
442 def set_backward_prefetch_cells(self, hsdp_cell_list: List[Any]) -> None:
443 """Set cells prefetched during backward.
445 Args:
446 hsdp_cell_list: HSDP cells to prefetch ahead of backward.
447 """
448 self.backward_prefetch_cells = hsdp_cell_list
450 def _disable_forward_prefetch_for_recompute(self) -> None:
451 """Temporarily disable forward prefetch during activation recompute."""
452 self._backup_forward_fetch = self.forward_prefetch_cells
453 self.forward_prefetch_cells = []
455 def _restore_forward_prefetch_after_recompute(self) -> None:
456 """Restore the saved forward prefetch list after activation recompute."""
457 if self._backup_forward_fetch is None:
458 return
459 self.forward_prefetch_cells = self._backup_forward_fetch
460 self._backup_forward_fetch = None
462 def _register_post_backward_hook(self, args, kwargs):
463 """Wrap forward args/kwargs through PostBackwardFunction to register backward hook."""
464 if not torch.is_grad_enabled():
465 return args, kwargs
466 args_list, args_spec = tree_flatten(args)
467 kwargs_list, kwargs_spec = tree_flatten(kwargs)
468 args_kwargs_list = list(args_list) + list(kwargs_list)
469 inp_tensor_indices: List[int] = []
470 inp_tensors: List[torch.Tensor] = []
471 for i, obj in enumerate(args_kwargs_list):
472 if torch.is_tensor(obj) and obj.requires_grad:
473 inp_tensor_indices.append(i)
474 inp_tensors.append(obj)
475 if len(inp_tensors) == 0:
476 return args, kwargs # no tensors that require gradients
477 processed_tensors = PostBackwardFunction.apply(self, *inp_tensors)
478 for inp_tensor_idx, processed_tensor in zip(inp_tensor_indices, processed_tensors):
479 args_kwargs_list[inp_tensor_idx] = processed_tensor
480 args_list = args_kwargs_list[: len(args_list)]
481 kwargs_list = args_kwargs_list[len(args_list) :]
482 args = tree_unflatten(args_list, args_spec)
483 kwargs = tree_unflatten(kwargs_list, kwargs_spec)
484 return args, kwargs
486 @_dynamo_disable
487 def _forward_pre_hook(self, cell, args, kwargs):
488 """Execute forward pre hook and set up backward hook."""
489 args, kwargs = self._hsdp_forward_pre_hook(cell, args, kwargs)
490 return self._register_post_backward_hook(args, kwargs)
492 def _register_backward_pre_hook(self, outputs) -> None:
493 """Register gradient hooks on all requires-grad outputs to trigger backward pre hook."""
494 flat_outputs, _ = tree_flatten(outputs)
495 for output in flat_outputs:
496 if isinstance(output, torch.Tensor) and output.requires_grad:
497 handle_ref = [None]
498 # pylint: disable=C0103, W0102
500 def wrapper_for_backward_pre_hook(grad, _handle_ref=handle_ref):
501 """Remove this hook after it fires to prevent accmulation"""
502 handle = _handle_ref[0]
503 if handle is not None:
504 handle.remove()
505 return self._backward_pre_hook(grad)
506 # pylint: enable=C0103, W0102
507 handle = output.register_hook(wrapper_for_backward_pre_hook)
508 handle_ref[0] = handle
510 @_dynamo_disable
511 def _forward_hook(self, cell, inputs, outputs): # pylint: disable=R1710
512 """Execute forward hook."""
513 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
514 return
515 self._register_backward_pre_hook(outputs)
516 if self.scheduler_ctx.root_bp_state:
517 self._restore_forward_prefetch_after_recompute()
518 return
519 return self._hsdp_forward_hook(cell, inputs, outputs)
521 # pylint: disable=W0212
522 @_dynamo_disable
523 def _backward_pre_hook(self, grad):
524 """Execute backward pre hook."""
525 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD:
526 return grad
527 if not self.scheduler_ctx.post_backward_final_callback_queued:
528 self.scheduler_ctx.post_backward_final_callback_queued = True
529 root_scheduler = self.scheduler_ctx.root_module.hsdp_scheduler
530 Variable._execution_engine.queue_callback(root_scheduler._root_backward_hook)
531 self.scheduler_ctx.root_bp_state = True
532 self._hsdp_backward_pre_hook(self.cell, None)
533 return grad
535 @_dynamo_disable
536 def _root_backward_hook(self):
537 """Drain all DP pipelines, then run final TP reduction and apply gradients."""
538 logger.debug("hook=root_backward_hook enter module=%s", self.hsdp_state)
539 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
540 # let modules which are not triggered backward_hook launch backward communication.
541 hsdp_scheduler._backward_hook()
542 # Keep pending gradient reductions in flight while unused parameter
543 # prefetches finish, rather than delaying their launch on those waits.
544 self._finalize_unused_prefetches()
545 self.scheduler_ctx.root_bp_state = False
546 with torch.profiler.record_function(f"root_backward reduce:{self.hsdp_state.module_name}"):
547 logger.debug(
548 "hook=root_backward_hook action=final_reduce module=%s",
549 self.hsdp_state,
550 )
551 self._finalize_comm_fusion_reductions()
552 self._finalize_per_param_reductions()
553 self.launch_tp_replicate_reduce_and_apply()
554 # BACKWARD only deduplicates hooks within one autograd execution;
555 # accumulated no-sync gradients must outlive this state reset.
556 for scheduler in self.scheduler_ctx.all_hsdp_schedulers:
557 scheduler.scheduler_state = None
558 self.scheduler_ctx.post_backward_final_callback_queued = False
560 def _finalize_unused_prefetches(self) -> None:
561 """Wait and release unused all-gathers after autograd has finished.
563 A module used in forward need not participate in this backward. Its
564 speculative prefetch may therefore have no pre-backward consumer.
565 """
566 for scheduler in self.scheduler_ctx.all_hsdp_schedulers:
567 state = scheduler.hsdp_state
568 if state is None:
569 continue
570 for param in state.hsdp_params:
571 context = param.allgather_comm_ctx
572 if context.allgather_handle is None and context.allgather_output is None:
573 continue
574 with profiler_record(f"finalize_unused_prefetch param:{param}"):
575 if context.allgather_handle is not None:
576 context.allgather_handle.wait()
577 context.allgather_handle = None
578 context.allgather_output = None
579 if state.is_shard and param.shard_world_size > 1:
580 param.free_unsharded_param()
581 if state.param_group is not None:
582 for bucket in state.param_group.all_gather_buckets:
583 result = bucket.all_gather_result
584 if result is None:
585 continue
586 param_fqns = ", ".join(str(param) for param in bucket.hsdp_params)
587 with profiler_record(f"finalize_unused_prefetch bucket:{state.module_name}:{param_fqns}"):
588 if result.handle is not None:
589 result.handle.wait()
590 result.handle = None
591 result.all_gather_input = None
592 result.all_gather_output = None
593 bucket.all_gather_result = None
595 def check_state_prefetch_handles_consumed(self) -> None:
596 """Check that all parameters in the module tree consumed their prefetch work.
598 Raises:
599 RuntimeError: If a parameter or fused bucket retains an all-gather
600 handle because the configured prefetch order is incorrect.
601 """
602 unconsumed_param_fqns = []
603 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
604 hsdp_state = hsdp_scheduler.hsdp_state
605 if hsdp_state is None:
606 continue
607 unconsumed_param_fqns.extend(
608 str(param) for param in hsdp_state.hsdp_params
609 if param.allgather_comm_ctx.allgather_handle is not None
610 )
611 if hsdp_state.param_group is not None:
612 for bucket in hsdp_state.param_group.all_gather_buckets:
613 result = bucket.all_gather_result
614 if result is not None and result.handle is not None:
615 unconsumed_param_fqns.extend(str(param) for param in bucket.hsdp_params)
616 if unconsumed_param_fqns:
617 param_fqns = ", ".join(dict.fromkeys(unconsumed_param_fqns))
618 raise RuntimeError(
619 f"Unconsumed prefetch handles for parameters: {param_fqns}. "
620 "The prefetch order is incorrect. Set forward/backward prefetch order "
621 "according to the model's actual execution order."
622 )
624 def _finalize_comm_fusion_reductions(self) -> None:
625 """Drain the comm_fusion=True RS/AR pipeline."""
626 comm_ctx = self.scheduler_ctx.param_group_comm_ctx
627 if comm_ctx.all_reduce_param_group is not None:
628 logger.debug(
629 "hook=root_backward_hook wait=comm_fusion_all_reduce module=%s",
630 self.hsdp_state,
631 )
632 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad()
633 comm_ctx.all_reduce_param_group = None
634 if comm_ctx.pre_param_group is not None:
635 logger.debug(
636 "hook=root_backward_hook wait=comm_fusion_reduce_scatter module=%s",
637 self.hsdp_state,
638 )
639 comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce()
640 comm_ctx.pre_param_group = None
641 if comm_ctx.all_reduce_param_group is not None:
642 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad()
643 comm_ctx.all_reduce_param_group = None
645 def _finalize_per_param_reductions(self) -> None:
646 """Drain the module-tree-local comm_fusion=False RS/AR queues."""
647 # A fused root may own non-fused children, so always drain the tree queues.
648 last_all_reduce_groups = self.hsdp_state._wait_prev_reduce_scatter()
649 self.hsdp_state._wait_prev_reduce_scatter_without_all_reduce()
650 self.hsdp_state._issue_prev_fused_all_reduce(last_all_reduce_groups)
651 self.hsdp_state.wait_and_split_all_reduce_work_groups()
653 def launch_tp_replicate_reduce_and_apply(self) -> None:
654 """Run final TP replicate reductions and apply gradients for all states."""
655 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
656 hsdp_state = hsdp_scheduler.hsdp_state
657 if hsdp_state is None:
658 continue
659 need_synchronize = False
660 for hsdp_param in hsdp_state.hsdp_params:
661 reduced_grad = hsdp_param.all_reduce_comm_ctx.all_reduce_output
662 if reduced_grad is None:
663 reduced_grad = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output
664 if reduced_grad is None:
665 continue
666 hsdp_param.all_reduce_source_replicate_grad_inplace(
667 reduced_grad,
668 hsdp_state.reduce_op_type,
669 )
670 need_synchronize = hsdp_param.apply_reduced_grad(reduced_grad) or need_synchronize
671 hsdp_param.clear_all_reduce_output()
672 hsdp_param.clear_reduce_scatter_output()
673 hsdp_state._sync_current_stream_if_needed(need_synchronize)
675 @_dynamo_disable
676 def _backward_hook(self):
677 """Execute backward hook."""
678 if self.scheduler_state == FSDPSchedulerState.BACKWARD:
679 return
680 self._hsdp_backward_hook(self.cell, None, None)
682 @staticmethod
683 def _register_forward_module_hook(mod, hook) -> None:
684 """Register forward hook; use ``always_call=True`` when supported (matches PyTorch FSDP)."""
685 sig = inspect.signature(mod.register_forward_hook)
686 if "always_call" in sig.parameters:
687 mod.register_forward_hook(hook, prepend=False, always_call=True)
688 else:
689 mod.register_forward_hook(hook, prepend=False)
691 def _register_forward_backward_hooks(self):
692 """Register module forward and backward hook on all managed modules."""
693 if self._fsdp_group_post_pending is None:
694 for mod in self.modules:
695 mod.register_forward_pre_hook(self._forward_pre_hook, with_kwargs=True)
696 mod.register_forward_hook(self._forward_hook)
697 return
698 for mod in self.modules:
699 mod.register_forward_pre_hook(self._grouped_forward_pre_hook, with_kwargs=True)
700 grouped_forward_hook = _dynamo_disable(self._make_grouped_forward_post_hook(mod))
701 self._register_forward_module_hook(mod, grouped_forward_hook)