Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / fsdp_adapter.py: 27%
311 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"""fsdp_adapter: FSDP2 construction-time adapter (05 §15.2.4).
17Wrap policy, mesh/policy construction, ``fully_shard`` application and
18prefetch configuration all happen at model-construction time (not in the
19Trainer step runtime); ``FSDP2Manager`` / ``_instantiate_fsdp2`` and the
20``parallelize`` signature are unchanged. The source-layout discovery logic
21(FQN/DTensor source-shard resolution) merged into
22``distributed/_builder/source_shard.py`` (05 row 406).
23"""
24# pylint: disable=forbidden-backend-import
26from __future__ import annotations
28import logging
29from collections import defaultdict
30from collections.abc import Mapping, Sequence
31from dataclasses import dataclass
32from typing import TYPE_CHECKING, Any
34import torch
35from torch.nn import Module as ModuleClass
36from torch.nn import Parameter as ParameterClass
38import hyper_parallel.core.fully_shard.utils as fully_shard_utils
39from hyper_parallel import DeviceMesh, HSDPModule, fully_shard
40from hyper_parallel.distributed._builder.source_shard import (
41 SourceShardInfoByFQN, # pylint: disable=unused-import
42 SourceShardInfoByParam, # pylint: disable=unused-import
43 _build_managed_source_shard_info,
44 _build_source_shard_info_by_param,
45 _source_infos_for_fully_shard,
46)
47from hyper_parallel.models.build_options import FSDP2Config
48from hyper_parallel.models.registry import get_model_adapter
50if TYPE_CHECKING:
51 from hyper_parallel.distributed.mesh import (
52 MeshContext,
53 )
55logger = logging.getLogger(__name__)
58@dataclass(frozen=True)
59class _WrapModuleInfo:
60 """One FSDP child unit selected by automatic or adapter policy."""
62 fqn: str
63 module: ModuleClass
66def _is_expert_source_mesh(mesh: DeviceMesh | None) -> bool:
67 """Whether a source mesh carries the expert (EP) axis."""
68 return mesh is not None and "ep" in (getattr(mesh, "mesh_dim_names", None) or ())
71def _get_checkpoint_wrapped_module(module: ModuleClass) -> ModuleClass | None:
72 """Return the module directly wrapped by a supported checkpoint wrapper."""
73 for attr_name in ("_wrapped_module", "_checkpoint_wrapped_module"):
74 wrapped_module = getattr(module, attr_name, None)
75 if isinstance(wrapped_module, ModuleClass):
76 return wrapped_module
77 return None
80class FSDP2Manager:
81 """Configure and apply nested FSDP2 wrapping for a dual-mode model."""
83 def __init__(
84 self,
85 config: FSDP2Config,
86 mesh_context: MeshContext,
87 fp32_main_params: bool = False,
88 ) -> None:
89 """Initialize the manager from resolved YAML config and runtime mesh.
91 Args:
92 config: Resolved top-level ``fsdp_config``.
93 mesh_context: Runtime topology containing the shared world mesh.
94 fp32_main_params: Whether FSDP writes reduced gradients to
95 ``main_grad`` for the fp32 main-parameter optimizer wrapper.
96 """
97 self.config = config
98 self.mesh_context = mesh_context
99 self.fp32_main_params = fp32_main_params
102 def _resolve_replicate_params(
103 self,
104 model: ModuleClass,
105 ) -> set[ParameterClass] | None:
106 """Resolve configured replicate-parameter FQNs on the final model."""
107 if not self.config.replicate_params:
108 return None
110 parameters_by_fqn = dict(model.named_parameters(remove_duplicate=False))
111 replicate_params = set()
112 for parameter_fqn in self.config.replicate_params:
113 parameter = parameters_by_fqn.get(parameter_fqn)
114 if parameter is None:
115 raise ValueError(
116 "replicate_params contains unknown parameter FQN: "
117 f"{parameter_fqn}"
118 )
119 replicate_params.add(parameter)
120 return replicate_params
122 def _build_expert_fsdp_mesh(self) -> DeviceMesh:
123 """Select the expert FSDP child mesh."""
124 expert_mesh = self.mesh_context.fsdp_moe_mesh
125 if expert_mesh is None:
126 raise ValueError("Expert FSDP units require MeshContext.fsdp_moe_mesh")
127 if "edp_replicate" in (expert_mesh.mesh_dim_names or ()):
128 return expert_mesh[("edp_replicate", "edp_shard")]
129 return expert_mesh["edp_shard"]
131 def _build_dense_fsdp_mesh(self) -> DeviceMesh | None:
132 """Select the prebuilt dense FSDP child mesh when available."""
133 dense_mesh = self.mesh_context.fsdp_non_moe_mesh
134 if dense_mesh is None:
135 return None
136 if self.mesh_context.dp_replicate_size > 1:
137 return dense_mesh[("fsdp_replicate", "fsdp_shard")]
138 return dense_mesh["fsdp_shard"]
140 def _build_compatibility_fsdp_mesh(self) -> DeviceMesh:
141 """Build an FSDP mesh from the combined world mesh."""
142 world_mesh = self.mesh_context.device_mesh
143 if world_mesh is None:
144 raise ValueError("FSDP2 requires MeshContext.fsdp_non_moe_mesh")
145 mesh_dim_names = world_mesh.mesh_dim_names
146 if mesh_dim_names is None:
147 raise ValueError("FSDP2 requires named DeviceMesh dimensions")
148 shard_dim_names = tuple(
149 dim_name for dim_name in ("dp_shard", "cp") if dim_name in mesh_dim_names
150 )
151 if not shard_dim_names:
152 raise ValueError("FSDP2 requires a dp_shard or cp dimension in the world mesh")
153 shard_mesh = (
154 world_mesh[shard_dim_names[0]]
155 if len(shard_dim_names) == 1
156 else world_mesh[shard_dim_names].flatten("fsdp_shard")
157 )
158 if "dp_replicate" not in mesh_dim_names:
159 return shard_mesh
160 return DeviceMesh.concatenate([world_mesh["dp_replicate"], shard_mesh])
162 def _build_fsdp_actual_mesh(self, expert: bool = False) -> DeviceMesh:
163 """Select the prebuilt dense or expert FSDP child mesh."""
164 if expert:
165 return self._build_expert_fsdp_mesh()
166 dense_mesh = self._build_dense_fsdp_mesh()
167 if dense_mesh is not None:
168 return dense_mesh
169 return self._build_compatibility_fsdp_mesh()
171 def _build_mixed_precision_policy(self) -> fully_shard_utils.MixedPrecisionPolicy:
172 """Resolve the configured dtype strings to torch dtypes.
174 Dtype strings stay YAML-friendly; ``float32`` resolves to the
175 framework's ``float32`` dtype object. Fully-sharded params without an
176 explicit param dtype fall back to the framework default.
177 """
178 mix_precision = self.config.mix_precision
179 dtype_by_name = {
180 None: None,
181 "bfloat16": torch.bfloat16,
182 "float16": torch.float16,
183 "float32": torch.float32,
184 }
185 return fully_shard_utils.MixedPrecisionPolicy(
186 param_dtype=dtype_by_name[mix_precision.param_dtype],
187 reduce_dtype=dtype_by_name[mix_precision.reduce_dtype],
188 output_dtype=dtype_by_name[mix_precision.output_dtype],
189 cast_forward_inputs=mix_precision.cast_forward_inputs,
190 apply_grad_on_fp32_main_grad=self.fp32_main_params,
191 )
193 def _build_offload_policy(self) -> fully_shard_utils.OffloadPolicy:
194 """Return CPU offload when enabled, otherwise the default no-offload policy."""
195 if not self.config.enable_offload:
196 return fully_shard_utils.OffloadPolicy()
197 return fully_shard_utils.CPUOffloadPolicy()
199 def _build_fully_shard_kwargs(
200 self,
201 fsdp_mesh: DeviceMesh,
202 ) -> tuple[dict[str, Any], dict[str, Any]]:
203 """Build shared sublayer kwargs and the root-specialized copy."""
204 mp_policy = self._build_mixed_precision_policy()
205 offload_policy = self._build_offload_policy()
207 fsdp_sublayer_kwargs = {
208 "mesh": fsdp_mesh,
209 "reshard_after_forward": self.config.reshard_after_forward,
210 "mp_policy": mp_policy,
211 "offload_policy": offload_policy,
212 "comm_fusion": self.config.comm_fusion,
213 "comm_fusion_zero_copy": self.config.comm_fusion_zero_copy,
214 }
215 fsdp_root_kwargs = fsdp_sublayer_kwargs.copy()
216 fsdp_root_kwargs["reshard_after_forward"] = False
217 return fsdp_sublayer_kwargs, fsdp_root_kwargs
219 @staticmethod
220 def _find_transformer_block_modules(
221 model: ModuleClass,
222 excluded_subtree_module_fqns: set[str] | None = None,
223 ) -> tuple[list[_WrapModuleInfo], set[int]]:
224 """Find transformer blocks under HF gradient-checkpointing containers.
226 Args:
227 model: Root model to inspect.
228 excluded_subtree_module_fqns: Roots declared by a model adapter as
229 non-decoder branches or self-contained FSDP units. Automatic
230 decoder discovery skips these roots and their descendants.
231 """
232 wrap_modules = []
233 wrapped_module_ids = set()
234 excluded_subtree_module_fqns = excluded_subtree_module_fqns or set()
236 def is_inside_declared_subtree(module_fqn: str) -> bool:
237 """Whether one module is inside an automatic-discovery exclusion.
239 Args:
240 module_fqn: Fully qualified name of the module.
241 """
242 return any(
243 module_fqn == declared_fqn
244 or module_fqn.startswith(f"{declared_fqn}.")
245 for declared_fqn in excluded_subtree_module_fqns
246 )
248 for container_fqn, container in model.named_modules():
249 if container_fqn and is_inside_declared_subtree(container_fqn):
250 continue
251 if id(container) in wrapped_module_ids:
252 continue
253 if not hasattr(container, "gradient_checkpointing"):
254 continue
255 for child_name, child in container.named_children():
256 child_fqn = f"{container_fqn}.{child_name}" if container_fqn else child_name
257 if is_inside_declared_subtree(child_fqn):
258 continue
259 blocks = list(child.named_children())
260 if not blocks:
261 continue
262 for block_name, block in blocks:
263 if id(block) in wrapped_module_ids:
264 continue
265 wrapped_module_ids.add(id(block))
266 wrapped_module = _get_checkpoint_wrapped_module(block)
267 if wrapped_module is not None:
268 # The wrapper and its direct child represent one logical
269 # transformer block during module-tree traversal.
270 wrapped_module_ids.add(id(wrapped_module))
271 wrap_modules.append(
272 _WrapModuleInfo(f"{child_fqn}.{block_name}", block)
273 )
274 return wrap_modules, wrapped_module_ids
276 def _find_expert_wrap_modules(
277 self,
278 model: ModuleClass,
279 metadata_by_parameter: SourceShardInfoByParam | None,
280 wrapped_module_ids: set[int],
281 ) -> list[_WrapModuleInfo]:
282 """Find nested expert containers whose parameters use the expert mesh."""
283 if metadata_by_parameter is None or self.mesh_context.fsdp_moe_mesh is None:
284 return []
285 expert_parameters = {
286 parameter
287 for parameter, source_shard_info in metadata_by_parameter.items()
288 if _is_expert_source_mesh(source_shard_info.mesh)
289 }
290 expert_module_fqns = {
291 parameter_fqn.split(".experts.", 1)[0] + ".experts"
292 for parameter_fqn, parameter in model.named_parameters(remove_duplicate=False)
293 if parameter in expert_parameters and ".experts." in parameter_fqn
294 }
295 module_by_fqn = dict(model.named_modules())
296 expert_wrap_modules = []
297 for expert_module_fqn in sorted(expert_module_fqns):
298 expert_module = module_by_fqn.get(expert_module_fqn)
299 if expert_module is None or id(expert_module) in wrapped_module_ids:
300 continue
301 wrapped_module_ids.add(id(expert_module))
302 expert_wrap_modules.append(
303 _WrapModuleInfo(expert_module_fqn, expert_module)
304 )
305 return expert_wrap_modules
307 @staticmethod
308 def _get_model_adapter_spec(model: ModuleClass) -> Any | None:
309 """Resolve a model adapter from the final model configuration."""
310 config = getattr(model, "config", None)
311 if config is None:
312 return None
313 identities = (
314 getattr(config, "model_type", ""),
315 *((getattr(config, "architectures", None) or [])[:1]),
316 )
317 for identity in identities:
318 if identity:
319 adapter_spec = get_model_adapter(identity)
320 if adapter_spec is not None:
321 return adapter_spec
322 return None
324 @classmethod
325 def _find_adapter_declared_wrap_modules(
326 cls,
327 model: ModuleClass,
328 ) -> list[_WrapModuleInfo]:
329 """Resolve adapter-declared FSDP units outside HF decoder containers.
331 ``ModelAdapterSpec.fsdp_wrap_modules`` keeps architecture-specific
332 execution boundaries on the family adapter side. It receives the
333 finalized model and returns exact root-relative module FQNs.
335 Args:
336 model: Root model being wrapped.
338 Returns:
339 Additional FSDP child units selected by model declarations.
341 Raises:
342 ValueError: If a declaration has an invalid type, targets no
343 module in the model tree, or aliases the same module twice.
344 """
345 adapter_spec = cls._get_model_adapter_spec(model)
346 if adapter_spec is None or adapter_spec.fsdp_wrap_modules is None:
347 return []
349 declared_fqns = adapter_spec.fsdp_wrap_modules(model)
350 if isinstance(declared_fqns, str) or not isinstance(declared_fqns, Sequence):
351 raise ValueError(
352 "ModelAdapterSpec.fsdp_wrap_modules must return a sequence of exact module FQNs"
353 )
354 module_by_fqn = dict(model.named_modules())
355 declared_wrap_modules = []
356 wrapped_module_ids = set()
357 for module_fqn in declared_fqns:
358 if not isinstance(module_fqn, str) or not module_fqn:
359 raise ValueError(
360 "ModelAdapterSpec.fsdp_wrap_modules entries must be non-empty strings"
361 )
362 wrap_module = module_by_fqn.get(module_fqn)
363 if wrap_module is None:
364 raise ValueError(
365 "ModelAdapterSpec.fsdp_wrap_modules declared an unknown module: "
366 f"{module_fqn}"
367 )
368 if id(wrap_module) in wrapped_module_ids:
369 raise ValueError(
370 "ModelAdapterSpec.fsdp_wrap_modules aliases one module more than once: "
371 f"{module_fqn}"
372 )
373 wrapped_module_ids.add(id(wrap_module))
374 declared_wrap_modules.append(_WrapModuleInfo(module_fqn, wrap_module))
375 return declared_wrap_modules
377 @classmethod
378 def _find_adapter_excluded_subtrees(cls, model: ModuleClass) -> set[str]:
379 """Resolve model branches excluded from automatic decoder discovery."""
380 adapter_spec = cls._get_model_adapter_spec(model)
381 excluded_subtrees_provider = getattr(adapter_spec, "fsdp_excluded_subtrees", None)
382 if excluded_subtrees_provider is None:
383 return set()
384 if not callable(excluded_subtrees_provider):
385 raise ValueError("ModelAdapterSpec.fsdp_excluded_subtrees must be callable")
386 excluded_fqns = excluded_subtrees_provider(model) # pylint: disable=not-callable
387 if isinstance(excluded_fqns, str) or not isinstance(excluded_fqns, Sequence):
388 raise ValueError(
389 "ModelAdapterSpec.fsdp_excluded_subtrees must return a sequence of exact module FQNs"
390 )
391 module_by_fqn = dict(model.named_modules())
392 for module_fqn in excluded_fqns:
393 if not isinstance(module_fqn, str) or not module_fqn:
394 raise ValueError(
395 "ModelAdapterSpec.fsdp_excluded_subtrees entries must be non-empty strings"
396 )
397 if module_fqn not in module_by_fqn:
398 raise ValueError(
399 "ModelAdapterSpec.fsdp_excluded_subtrees declared an unknown module: "
400 f"{module_fqn}"
401 )
402 return set(excluded_fqns)
404 def _find_wrap_modules(
405 self,
406 model: ModuleClass,
407 metadata_by_parameter: SourceShardInfoByParam | None = None,
408 ) -> list[_WrapModuleInfo]:
409 """Find automatic decoder/expert units and adapter-declared units."""
410 declared_wrap_modules = self._find_adapter_declared_wrap_modules(model)
411 if declared_wrap_modules:
412 logger.info(
413 "Adapter-declared FSDP child units: %s",
414 [wrap_module.fqn for wrap_module in declared_wrap_modules],
415 )
416 excluded_subtree_module_fqns = {
417 wrap_module.fqn for wrap_module in declared_wrap_modules
418 }
419 excluded_subtree_module_fqns.update(
420 self._find_adapter_excluded_subtrees(model)
421 )
422 wrap_modules, wrapped_module_ids = self._find_transformer_block_modules(
423 model,
424 excluded_subtree_module_fqns=excluded_subtree_module_fqns,
425 )
426 wrapped_module_ids.update(
427 id(wrap_module.module) for wrap_module in declared_wrap_modules
428 )
429 wrap_modules.extend(declared_wrap_modules)
430 wrap_modules.extend(
431 self._find_expert_wrap_modules(
432 model,
433 metadata_by_parameter,
434 wrapped_module_ids,
435 )
436 )
437 if not wrap_modules:
438 raise ValueError(
439 "wrap_policy='transformer_block' found no decoder, expert, or "
440 "adapter-declared FSDP child unit"
441 )
442 return wrap_modules
444 @staticmethod
445 def _resolve_parameter_owners(
446 model: ModuleClass,
447 wrap_modules: list[_WrapModuleInfo],
448 ) -> dict[ParameterClass, ModuleClass]:
449 """Assign each parameter to its deepest common FSDP wrap module."""
450 aliases_by_parameter: dict[ParameterClass, list[str]] = defaultdict(list)
451 for parameter_fqn, parameter in model.named_parameters(remove_duplicate=False):
452 aliases_by_parameter[parameter].append(parameter_fqn)
454 owner_by_parameter = {}
455 for parameter, parameter_fqns in aliases_by_parameter.items():
456 common_wrap_modules = [
457 wrap_module
458 for wrap_module in wrap_modules
459 if all(
460 parameter_fqn.startswith(f"{wrap_module.fqn}.")
461 for parameter_fqn in parameter_fqns
462 )
463 ]
464 if common_wrap_modules:
465 owner_by_parameter[parameter] = max(
466 common_wrap_modules,
467 key=lambda wrap_module: wrap_module.fqn.count("."),
468 ).module
469 else:
470 owner_by_parameter[parameter] = model
471 return owner_by_parameter
474 def _uses_expert_mesh(self, metadata_by_parameter: SourceShardInfoByParam | None) -> bool:
475 """Return whether one FSDP unit owns routed-expert source metadata."""
476 if metadata_by_parameter is None or self.mesh_context.fsdp_moe_mesh is None:
477 return False
478 return any(
479 _is_expert_source_mesh(source_shard_info.mesh)
480 for source_shard_info in metadata_by_parameter.values()
481 )
484 @staticmethod
485 def _build_managed_replicate_params(
486 owner: ModuleClass,
487 owner_by_parameter: Mapping[ParameterClass, ModuleClass],
488 replicate_params: set[ParameterClass] | None,
489 ignored_params: set[ParameterClass] | None = None,
490 ) -> set[ParameterClass] | None:
491 """Select configured replicate parameters owned by one FSDP module."""
492 if replicate_params is None:
493 return None
494 return {
495 parameter
496 for parameter, parameter_owner in owner_by_parameter.items()
497 if parameter_owner is owner and parameter in replicate_params
498 and parameter not in (ignored_params or ())
499 }
501 def _configure_source_layout_gradient_scaling(
502 self,
503 hsdp_module: HSDPModule,
504 managed_source_shard_info: SourceShardInfoByParam | None,
505 ) -> bool:
506 """Restore global-mean gradients for one source-layout FSDP unit.
508 Master defaults every ``fully_shard`` unit to AVG reduction, so units
509 with source-layout metadata explicitly select SUM. The token-weighted
510 backward loss is scaled by the FSDP data domain (DP x CP), so a
511 source-layout SUM-reduced unit must apply its inverse exactly once.
512 With loss parallelism, TP ranks hold shards or partial contributions
513 to one logical loss and TP is not another averaging axis. Without loss
514 parallelism, logits and the loss are replicated over TP, so the SUM
515 path also accumulates ``tp_size`` identical loss replicas and must
516 divide by TP. EP is excluded because expert ranks own different
517 parameters.
519 Returns:
520 Whether a gradient scaling factor was configured for the unit.
521 """
522 if not managed_source_shard_info:
523 return False
524 hsdp_module.set_reduce_op_type("sum", recurse=False)
526 tp_loss_replica_size = (
527 1 if self.mesh_context.loss_parallel else self.mesh_context.tp_size
528 )
529 source_gradient_domain_size = (
530 self.mesh_context.dp_size
531 * self.mesh_context.cp_size
532 * tp_loss_replica_size
533 )
534 if source_gradient_domain_size <= 1:
535 return False
536 hsdp_module.set_gradient_scaling_factor(1.0 / source_gradient_domain_size)
537 return True
539 def _configure_prefetch(self, hsdp_modules: list[HSDPModule]) -> None:
540 """Configure forward and backward prefetch targets by list depth."""
541 for module_index, hsdp_module in enumerate(hsdp_modules):
542 forward_end = module_index + self.config.forward_prefetch_depth + 1
543 hsdp_module.set_modules_to_forward_prefetch(
544 hsdp_modules[module_index + 1:forward_end]
545 )
547 backward_start = max(
548 0, module_index - self.config.backward_prefetch_depth
549 )
550 hsdp_module.set_modules_to_backward_prefetch(
551 list(reversed(hsdp_modules[backward_start:module_index]))
552 )
554 def _order_wrap_modules(
555 self,
556 model: ModuleClass,
557 wrap_modules: list[_WrapModuleInfo],
558 ) -> list[_WrapModuleInfo]:
559 """Order FSDP units by adapter execution order or module traversal."""
560 module_order = {
561 id(module): module_index
562 for module_index, (_, module) in enumerate(model.named_modules())
563 }
564 default_order = sorted(
565 wrap_modules,
566 key=lambda wrap_module: module_order.get(id(wrap_module.module), len(module_order)),
567 )
568 adapter_spec = self._get_model_adapter_spec(model)
569 execution_order_provider = getattr(adapter_spec, "fsdp_execution_order", None)
570 if execution_order_provider is None:
571 return default_order
572 if not callable(execution_order_provider):
573 raise ValueError("ModelAdapterSpec.fsdp_execution_order must be callable")
575 wrap_module_by_fqn = {
576 wrap_module.fqn: wrap_module for wrap_module in wrap_modules
577 }
578 execution_fqns = execution_order_provider( # pylint: disable=not-callable
579 model,
580 tuple(wrap_module.fqn for wrap_module in default_order),
581 )
582 if isinstance(execution_fqns, str) or not isinstance(execution_fqns, Sequence):
583 raise ValueError(
584 "ModelAdapterSpec.fsdp_execution_order must return a sequence of module FQNs"
585 )
586 if any(not isinstance(module_fqn, str) for module_fqn in execution_fqns):
587 raise ValueError(
588 "ModelAdapterSpec.fsdp_execution_order entries must be module FQN strings"
589 )
590 if (
591 len(execution_fqns) != len(wrap_module_by_fqn)
592 or len(set(execution_fqns)) != len(execution_fqns)
593 or set(execution_fqns) != set(wrap_module_by_fqn)
594 ):
595 raise ValueError(
596 "ModelAdapterSpec.fsdp_execution_order must return every selected FSDP "
597 "child-unit FQN exactly once"
598 )
599 logger.info("Adapter-declared FSDP execution order: %s", list(execution_fqns))
600 return [wrap_module_by_fqn[module_fqn] for module_fqn in execution_fqns]
602 def _parallelize_child_units(
603 self,
604 wrap_modules: list[_WrapModuleInfo],
605 owner_by_parameter,
606 metadata_by_parameter,
607 replicate_params,
608 dense_fsdp_kwargs: dict[str, Any],
609 ignored_params: set[ParameterClass] | None = None,
610 ) -> int:
611 """Apply FSDP to child units and return the gradient-scaled count."""
612 gradient_scaled_units = 0
613 wrapping_order = sorted(
614 wrap_modules, key=lambda wrap_module: wrap_module.fqn.count("."), reverse=True
615 )
616 for wrap_module in wrapping_order:
617 managed_source_info = _build_managed_source_shard_info(
618 self,
619 wrap_module.module,
620 metadata_by_parameter,
621 ignored_params,
622 )
623 fsdp_sublayer_kwargs = dense_fsdp_kwargs
624 if self._uses_expert_mesh(managed_source_info):
625 if self.mesh_context.fsdp_moe_mesh is None:
626 raise ValueError("Expert TP metadata requires MeshContext.fsdp_moe_mesh")
627 fsdp_sublayer_kwargs, _ = self._build_fully_shard_kwargs(
628 self._build_fsdp_actual_mesh(expert=True)
629 )
630 fully_shard( # pylint: disable=unexpected-keyword-arg
631 wrap_module.module,
632 source_shard_infos=_source_infos_for_fully_shard(managed_source_info),
633 replicate_params=self._build_managed_replicate_params(
634 wrap_module.module, owner_by_parameter, replicate_params, ignored_params
635 ),
636 ignored_params=ignored_params,
637 **fsdp_sublayer_kwargs,
638 )
639 if self._configure_source_layout_gradient_scaling(wrap_module.module, managed_source_info):
640 gradient_scaled_units += 1
641 return gradient_scaled_units
643 def _parallelize_root(
644 self,
645 model: ModuleClass,
646 owner_by_parameter,
647 metadata_by_parameter,
648 replicate_params,
649 dense_root_kwargs: dict[str, Any],
650 ignored_params: set[ParameterClass] | None = None,
651 ) -> int:
652 """Apply root FSDP and report whether gradient scaling was configured."""
653 root_source_info = _build_managed_source_shard_info(
654 self,
655 model,
656 metadata_by_parameter,
657 ignored_params,
658 )
659 if self._uses_expert_mesh(root_source_info):
660 raise ValueError("Routed expert parameters must belong to a nested experts FSDP unit")
661 fully_shard( # pylint: disable=unexpected-keyword-arg
662 model,
663 source_shard_infos=_source_infos_for_fully_shard(root_source_info),
664 replicate_params=self._build_managed_replicate_params(
665 model, owner_by_parameter, replicate_params, ignored_params
666 ),
667 ignored_params=ignored_params,
668 **dense_root_kwargs,
669 )
670 return int(self._configure_source_layout_gradient_scaling(model, root_source_info))
672 def _log_gradient_scaling(self, gradient_scaled_units: int) -> None:
673 """Log the source-layout gradient scaling applied to FSDP units."""
674 if not gradient_scaled_units:
675 return
676 domain_size = (
677 self.mesh_context.dp_size
678 * self.mesh_context.cp_size
679 * (1 if self.mesh_context.loss_parallel else self.mesh_context.tp_size)
680 )
681 logger.info(
682 "Applied source-layout SUM gradient scaling: dp_size=%d, cp_size=%d, tp_size=%d, "
683 "source_gradient_domain_size=%d, factor=%s, units=%d",
684 self.mesh_context.dp_size,
685 self.mesh_context.cp_size,
686 self.mesh_context.tp_size,
687 domain_size,
688 1.0 / domain_size,
689 gradient_scaled_units,
690 )
692 def parallelize(
693 self,
694 model: ModuleClass,
695 source_shard_info: SourceShardInfoByFQN | None = None,
696 ignored_params: set[ParameterClass] | None = None,
697 ) -> ModuleClass:
698 """Apply configured child FSDP units, root FSDP, and prefetch links.
700 Args:
701 model: Model to wrap in place.
702 source_shard_info: FQN-keyed source-layout metadata returned by
703 ``apply_sharding_plan``. The manager resolves parameter
704 identities before applying nested FSDP units.
706 Returns:
707 The model enhanced with the HyperParallel HSDPModule interface.
709 Raises:
710 ValueError: If TP is enabled without FQN metadata or configured
711 parameter FQNs cannot be resolved.
712 """
713 metadata_by_parameter = _build_source_shard_info_by_param(
714 self,
715 model,
716 source_shard_info,
717 )
718 replicate_params = self._resolve_replicate_params(model)
720 wrap_modules = self._find_wrap_modules(model, metadata_by_parameter)
721 owner_by_parameter = self._resolve_parameter_owners(model, wrap_modules)
722 dense_fsdp_kwargs, dense_root_kwargs = self._build_fully_shard_kwargs(
723 self._build_fsdp_actual_mesh()
724 )
726 gradient_scaled_units = self._parallelize_child_units(
727 wrap_modules,
728 owner_by_parameter,
729 metadata_by_parameter,
730 replicate_params,
731 dense_fsdp_kwargs,
732 ignored_params,
733 )
734 gradient_scaled_units += self._parallelize_root(
735 model,
736 owner_by_parameter,
737 metadata_by_parameter,
738 replicate_params,
739 dense_root_kwargs,
740 ignored_params,
741 )
742 self._log_gradient_scaling(gradient_scaled_units)
743 ordered_wrap_modules = self._order_wrap_modules(model, wrap_modules)
744 self._configure_prefetch(
745 [wrap_module.module for wrap_module in ordered_wrap_modules]
746 )
747 logger.info(
748 "Applied FSDP2 to %d child units and the root module",
749 len(wrap_modules),
750 )
751 return model
754def _instantiate_fsdp2(
755 *,
756 config: Any,
757 mesh_context: MeshContext,
758 fp32_main_params: bool = False,
759) -> FSDP2Manager | None:
760 """Create an FSDP2 manager for the resolved strategy config."""
761 if config is None:
762 return None
763 if isinstance(config, FSDP2Config):
764 return FSDP2Manager(
765 config=config,
766 mesh_context=mesh_context,
767 fp32_main_params=fp32_main_params,
768 )
769 raise ValueError(f"FSDP2Manager requires FSDP2Config, got {type(config).__name__}")
772def _apply_fsdp2(
773 model: "nn.Module",
774 fsdp2_manager: Any,
775 source_shard_info: Any,
776 ignored_params: set[ParameterClass] | None = None,
777) -> "nn.Module":
778 """Apply FSDP2; Torch scheduler hooks remain outside Dynamo by design.
780 Moved from ``_transformers/infrastructure.py`` (05 §15.2.1) so the FSDP
781 application step lives with ``FSDP2Manager.parallelize``.
782 """
783 if fsdp2_manager is None:
784 return model
785 if not isinstance(fsdp2_manager, FSDP2Manager):
786 logger.warning("fsdp2_manager is not an FSDP2Manager instance")
787 return model
788 model = fsdp2_manager.parallelize(
789 model,
790 source_shard_info=source_shard_info,
791 ignored_params=ignored_params,
792 )
793 logger.info("FSDP2 wrap applied")
794 return model