Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / fsdp.py: 84%
215 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates
2# Copyright 2026 Huawei Technologies Co., Ltd
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
16"""Trainer-side FSDP lifecycle policies and runtime diagnostics.
18Split out of the former ``auto_models/trainer/base.py`` in stage 7
19(05 §15.11 step 3). ``BaseTrainer`` keeps the same-named methods as thin
20delegating subclass hooks; the policy itself lives here. The FSDP config is
21duck-typed (``reshard_after_backward`` / ``dp_shard_size`` /
22``requires_grad_sync`` attributes) so this module does not import Trainer DTOs.
23"""
24# pylint: disable=forbidden-backend-import
26import logging
27from collections import defaultdict
28from dataclasses import dataclass, field
29from typing import Any, List, Optional
31import torch
32from torch.utils._pytree import tree_flatten
35logger = logging.getLogger(__name__)
38@dataclass
39class _FSDPMicroStepTrace:
40 """Observed FSDP unit order for one forward/backward micro-step."""
42 global_step: int
43 micro_step: int
44 forward: list[str] = field(default_factory=list)
45 backward: list[str] = field(default_factory=list)
48class FSDPRuntimeDiagnostics:
49 """Validate Torch FSDP prefetch order at the optimizer-step boundary.
51 This observer does not launch, wait, release, or otherwise mutate FSDP
52 communication. It records the actual module execution order and inspects
53 outstanding all-gather state immediately before the optimizer may update
54 parameter shards.
55 """
57 def __init__(self, hsdp_model_parts: List[Any]) -> None:
58 """Attach diagnostic hooks to all HSDP units in the model parts."""
59 self._module_by_name: dict[str, Any] = {}
60 self._name_by_module_id: dict[int, str] = {}
61 self._scheduler_by_name: dict[str, Any] = {}
62 self._hook_handles: list[Any] = []
63 self._traces: list[_FSDPMicroStepTrace] = []
64 self._active_trace: Optional[_FSDPMicroStepTrace] = None
65 self._discover_units(hsdp_model_parts)
66 self._register_hooks()
68 @property
69 def unit_count(self) -> int:
70 """Return the number of observed HSDP modules."""
71 return len(self._module_by_name)
73 def _discover_units(self, hsdp_model_parts: List[Any]) -> None:
74 """Build stable names for HSDP modules without importing backend internals."""
75 multiple_parts = len(hsdp_model_parts) > 1
76 for part_index, model_part in enumerate(hsdp_model_parts):
77 part_prefix = f"part[{part_index}]" if multiple_parts else ""
78 for module_fqn, module in model_part.named_modules():
79 scheduler = getattr(module, "hsdp_scheduler", None)
80 if scheduler is None or id(module) in self._name_by_module_id:
81 continue
82 if part_prefix and module_fqn:
83 module_name = f"{part_prefix}.{module_fqn}"
84 elif part_prefix:
85 module_name = part_prefix
86 else:
87 module_name = module_fqn or "<root>"
88 self._module_by_name[module_name] = module
89 self._name_by_module_id[id(module)] = module_name
90 self._scheduler_by_name[module_name] = scheduler
92 def _register_hooks(self) -> None:
93 """Register hooks on every unit for order and child-bypass diagnostics."""
94 for module_name, module in self._module_by_name.items():
95 self._hook_handles.append(
96 module.register_forward_pre_hook(
97 self._make_forward_pre_hook(module_name)
98 )
99 )
100 self._hook_handles.append(
101 module.register_forward_hook(
102 self._make_backward_registration_hook(module_name)
103 )
104 )
106 def _make_forward_pre_hook(self, module_name: str):
107 """Create a hook recording non-recompute forward execution."""
109 def record_forward(module: Any, inputs: Any) -> None:
110 """Record one observed module forward entry."""
111 del inputs
112 if self._active_trace is None:
113 return
114 scheduler = getattr(module, "hsdp_scheduler", None)
115 scheduler_ctx = getattr(scheduler, "scheduler_ctx", None)
116 if getattr(scheduler_ctx, "root_bp_state", False):
117 return
118 self._active_trace.forward.append(module_name)
120 return record_forward
122 def _make_backward_registration_hook(self, module_name: str):
123 """Create a forward hook attaching a one-shot output-gradient observer."""
125 def register_backward_observer(module: Any, inputs: Any, outputs: Any) -> None:
126 """Attach a one-shot gradient observer to differentiable outputs."""
127 del inputs
128 scheduler = getattr(module, "hsdp_scheduler", None)
129 scheduler_ctx = getattr(scheduler, "scheduler_ctx", None)
130 if getattr(scheduler_ctx, "root_bp_state", False):
131 return
132 output_tensors = [
133 output
134 for output in tree_flatten(outputs)[0]
135 if torch.is_tensor(output) and output.requires_grad
136 ]
137 if not output_tensors:
138 return
140 fired = [False]
141 hook_handles = []
143 def record_backward(gradient: Any) -> Any:
144 """Record the module on its first output-gradient callback."""
145 if not fired[0]:
146 fired[0] = True
147 if self._active_trace is not None:
148 self._active_trace.backward.append(module_name)
149 for hook_handle in hook_handles:
150 hook_handle.remove()
151 return gradient
153 for output_tensor in output_tensors:
154 hook_handles.append(output_tensor.register_hook(record_backward))
156 return register_backward_observer
158 def begin_micro_step(self, global_step: int, micro_step: int) -> None:
159 """Start tracing one micro-step in the current optimizer step."""
160 if micro_step == 0:
161 self._traces.clear()
162 trace = _FSDPMicroStepTrace(global_step, micro_step)
163 self._traces.append(trace)
164 self._active_trace = trace
166 @staticmethod
167 def _first_positions(execution_order: list[str]) -> dict[str, int]:
168 """Return first-observed positions while preserving repeated calls in logs."""
169 positions = {}
170 for position, module_name in enumerate(execution_order):
171 positions.setdefault(module_name, position)
172 return positions
174 def _prefetch_target_names(self, module_name: str, phase: str) -> list[str]:
175 """Resolve configured prefetch targets to diagnostic module names."""
176 scheduler = self._scheduler_by_name[module_name]
177 attribute = (
178 "forward_prefetch_cells"
179 if phase == "forward"
180 else "backward_prefetch_cells"
181 )
182 target_names = []
183 for target in getattr(scheduler, attribute, ()):
184 target_names.append(
185 self._name_by_module_id.get(
186 id(target),
187 f"<untracked:{type(target).__name__}>",
188 )
189 )
190 return target_names
192 def _find_order_violations(self) -> list[str]:
193 """Compare configured prefetch edges with each observed execution trace."""
194 violations = []
195 for trace in self._traces:
196 for phase in ("forward", "backward"):
197 execution_order = getattr(trace, phase)
198 positions = self._first_positions(execution_order)
199 for source_name, source_position in positions.items():
200 for target_name in self._prefetch_target_names(source_name, phase):
201 target_position = positions.get(target_name)
202 if target_position is None:
203 violations.append(
204 f"step={trace.global_step} micro_step={trace.micro_step} "
205 f"phase={phase}: {source_name} prefetched {target_name}, "
206 "but the target did not execute"
207 )
208 elif target_position <= source_position:
209 violations.append(
210 f"step={trace.global_step} micro_step={trace.micro_step} "
211 f"phase={phase}: {source_name} prefetched {target_name} "
212 f"at positions {source_position}->{target_position}"
213 )
214 return violations
216 @staticmethod
217 def _parameter_name(hsdp_param: Any, param_index: int) -> str:
218 """Return a useful parameter name for diagnostics."""
219 return getattr(hsdp_param, "_param_fqn", None) or f"param[{param_index}]"
221 def _find_pending_all_gathers(self) -> list[str]:
222 """Find unconsumed fused and non-fused all-gathers without mutating them."""
223 pending = []
224 seen_schedulers = set()
225 for module_name, scheduler in self._scheduler_by_name.items():
226 if id(scheduler) in seen_schedulers:
227 continue
228 seen_schedulers.add(id(scheduler))
229 state = getattr(scheduler, "hsdp_state", None)
230 if state is None:
231 continue
233 param_group = getattr(state, "param_group", None)
234 for bucket_index, bucket in enumerate(
235 getattr(param_group, "all_gather_buckets", ())
236 ):
237 result = getattr(bucket, "all_gather_result", None)
238 if result is not None:
239 handle = getattr(result, "handle", None)
240 pending.append(
241 f"{module_name}: fused_bucket[{bucket_index}] "
242 f"handle={'set' if handle is not None else 'completed-not-consumed'}"
243 )
245 for param_index, hsdp_param in enumerate(getattr(state, "hsdp_params", ())):
246 comm_ctx = getattr(hsdp_param, "allgather_comm_ctx", None)
247 handle = getattr(comm_ctx, "allgather_handle", None)
248 output = getattr(comm_ctx, "allgather_output", None)
249 if handle is not None or output is not None:
250 state_details = []
251 if handle is not None:
252 state_details.append("handle=set")
253 if output is not None:
254 state_details.append("output=not-consumed")
255 pending.append(
256 f"{module_name}: "
257 f"{self._parameter_name(hsdp_param, param_index)} "
258 + ", ".join(state_details)
259 )
260 return pending
262 @staticmethod
263 def _hsdp_parameter_has_gradient(hsdp_param: Any) -> bool:
264 """Return whether a managed parameter participated in backward."""
265 sharded_param = getattr(hsdp_param, "sharded_param", None)
266 candidates = (
267 getattr(sharded_param, "grad", None),
268 getattr(sharded_param, "main_grad", None),
269 getattr(hsdp_param, "_grad", None),
270 getattr(hsdp_param, "unsharded_accumulated_grad", None),
271 getattr(hsdp_param, "unsharded_accumulated_grad_data", None),
272 )
273 return any(value is not None for value in candidates)
275 def _find_bypassed_child_units(self) -> list[str]:
276 """Find child FSDP parameters used without executing their owner module."""
277 if not self._traces:
278 return []
279 forward_counts = defaultdict(int)
280 for trace in self._traces:
281 for module_name in trace.forward:
282 forward_counts[module_name] += 1
283 bypassed = []
284 root_names = {"<root>"}
285 root_names.update(
286 name for name in self._module_by_name
287 if name.startswith("part[") and "." not in name
288 )
289 for module_name, scheduler in self._scheduler_by_name.items():
290 if module_name in root_names or forward_counts[module_name] > 0:
291 continue
292 state = getattr(scheduler, "hsdp_state", None)
293 gradient_parameters = [
294 self._parameter_name(hsdp_param, param_index)
295 for param_index, hsdp_param in enumerate(getattr(state, "hsdp_params", ()))
296 if self._hsdp_parameter_has_gradient(hsdp_param)
297 ]
298 if gradient_parameters:
299 bypassed.append(
300 f"{module_name}: forward_count=0, gradient_parameters={gradient_parameters}"
301 )
302 return bypassed
304 def validate_before_optimizer_step(self) -> None:
305 """Fail before parameter updates if the traced FSDP lifecycle is unsafe."""
306 self._active_trace = None
307 order_violations = self._find_order_violations()
308 pending_all_gathers = self._find_pending_all_gathers()
309 bypassed_child_units = self._find_bypassed_child_units()
310 error_sections = []
311 if bypassed_child_units:
312 error_sections.append(
313 "child FSDP module was bypassed while its parameter "
314 "participated in backward:\n - " + "\n - ".join(bypassed_child_units)
315 )
316 if order_violations:
317 error_sections.append(
318 "configured prefetch targets do not follow "
319 "the observed runtime order:\n - " + "\n - ".join(order_violations)
320 )
321 if pending_all_gathers:
322 error_sections.append(
323 "unconsumed parameter all-gather exists at "
324 "the optimizer boundary:\n - " + "\n - ".join(pending_all_gathers)
325 )
326 if error_sections:
327 raise RuntimeError(
328 "\n".join(error_sections)
329 + "\nThe optimizer was not executed. Configure FSDP prefetch from the "
330 "model's actual forward/backward execution order; do not carry an "
331 "all-gather across optimizer.step()."
332 )
334 def trace_evidence(self) -> dict[str, Any]:
335 """Return stable execution-order evidence for structured reports."""
336 return {
337 "units": sorted(self._module_by_name),
338 "micro_steps": [
339 {
340 "global_step": trace.global_step,
341 "micro_step": trace.micro_step,
342 "forward": list(trace.forward),
343 "backward": list(trace.backward),
344 }
345 for trace in self._traces
346 ],
347 }
349 def close(self) -> None:
350 """Remove diagnostic hooks when a caller no longer needs the observer."""
351 for hook_handle in self._hook_handles:
352 hook_handle.remove()
353 self._hook_handles.clear()
356def build_fsdp_runtime_diagnostics(
357 hsdp_model_parts: List[Any],
358) -> Optional[FSDPRuntimeDiagnostics]:
359 """Create diagnostics for Torch Trainer model parts, or return ``None``."""
360 if not hsdp_model_parts:
361 return None
362 diagnostics = FSDPRuntimeDiagnostics(hsdp_model_parts)
363 if diagnostics.unit_count == 0:
364 diagnostics.close()
365 return None
366 logger.info(
367 "Enabled Trainer FSDP runtime diagnostics for %d modules",
368 diagnostics.unit_count,
369 )
370 return diagnostics
373def model_reshard(
374 hsdp_model_parts: List[Any],
375 fsdp_config: Any,
376 micro_step: int,
377 num_micro_steps: int,
378) -> None:
379 """Reshard model after backward pass."""
380 if (
381 fsdp_config.reshard_after_backward is False
382 and num_micro_steps > 1
383 ):
384 if micro_step == 0:
385 for model_part in hsdp_model_parts:
386 model_part.set_reshard_after_backward(False)
387 elif micro_step == num_micro_steps - 1:
388 for model_part in hsdp_model_parts:
389 model_part.set_reshard_after_backward(True)
392def configure_fsdp_gradient_sync(
393 hsdp_model_parts: List[Any],
394 fsdp_config: Any,
395 dp_replicate_size: int,
396 micro_step: int,
397 num_micro_steps: int,
398) -> None:
399 """Configure FSDP gradient synchronization for one micro step."""
400 if (
401 fsdp_config.dp_shard_size > 1
402 and num_micro_steps > 1
403 ):
404 is_last_micro_batch = micro_step == num_micro_steps - 1
405 requires_gradient_sync = (
406 fsdp_config.requires_grad_sync
407 or is_last_micro_batch
408 )
409 is_hsdp = dp_replicate_size > 1
410 for model_part in hsdp_model_parts:
411 model_part.set_requires_gradient_sync(requires_gradient_sync)
412 model_part.set_is_last_backward(is_last_micro_batch)
413 if is_hsdp:
414 model_part.set_requires_all_reduce(is_last_micro_batch)
417__all__ = [
418 "FSDPRuntimeDiagnostics",
419 "build_fsdp_runtime_diagnostics",
420 "configure_fsdp_gradient_sync",
421 "model_reshard",
422]