Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / fsdp_adapter.py: 82%

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). 

16 

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 

25 

26from __future__ import annotations 

27 

28import fnmatch 

29import logging 

30from collections import defaultdict 

31from collections.abc import Mapping, Sequence 

32from dataclasses import dataclass 

33from typing import TYPE_CHECKING, Any 

34 

35import torch 

36from torch.nn import Module as ModuleClass 

37from torch.nn import Parameter as ParameterClass 

38 

39from hyper_parallel.components.modules.module_family import collect_family_modules 

40import hyper_parallel.core.fully_shard.utils as fully_shard_utils 

41from hyper_parallel import DeviceMesh, HSDPModule, fully_shard 

42from hyper_parallel.core.fully_shard.hsdp_utils import get_managed_modules_parameters 

43from hyper_parallel.distributed._builder.source_shard import ( 

44 SourceShardInfoByFQN, # pylint: disable=unused-import 

45 SourceShardInfoByParam, # pylint: disable=unused-import 

46 _build_managed_source_shard_info, 

47 _build_source_shard_info_by_param, 

48 _source_infos_for_fully_shard, 

49) 

50from hyper_parallel.models.build_options import FSDP2Config 

51from hyper_parallel.models.registry import get_model_adapter 

52 

53if TYPE_CHECKING: 

54 from hyper_parallel.distributed.mesh import ( 

55 MeshContext, 

56 ) 

57 

58logger = logging.getLogger(__name__) 

59 

60 

61@dataclass(frozen=True) 

62class _WrapModuleInfo: 

63 """One FSDP child unit selected by automatic or adapter policy.""" 

64 

65 fqn: str 

66 module: ModuleClass 

67 

68 

69def _is_expert_source_mesh(mesh: DeviceMesh | None) -> bool: 

70 """Whether a source mesh carries the expert (EP) axis.""" 

71 return mesh is not None and "ep" in (getattr(mesh, "mesh_dim_names", None) or ()) 

72 

73 

74def _get_checkpoint_wrapped_module(module: ModuleClass) -> ModuleClass | None: 

75 """Return the module directly wrapped by a supported checkpoint wrapper.""" 

76 for attr_name in ("_wrapped_module", "_checkpoint_wrapped_module"): 

77 wrapped_module = getattr(module, attr_name, None) 

78 if isinstance(wrapped_module, ModuleClass): 

79 return wrapped_module 

80 return None 

81 

82 

83def _declared_expert_unit_fqns(model, expert_parameters): 

84 """Find family/native execution units by identity rather than an experts path suffix.""" 

85 expert_module_fqns = set() 

86 for boundary in collect_family_modules(model): 

87 parameters = set(boundary.module.parameters()) 

88 if parameters and parameters <= expert_parameters: 

89 expert_module_fqns.add(boundary.fqns[0]) 

90 for fqn, module in model.named_modules(): 

91 binding = getattr(module, "_hp_parallel_binding", None) 

92 if binding is not None and binding.ep_context is not None: 

93 expert_module_fqns.add(fqn) 

94 expert_module_fqns.discard("") 

95 return expert_module_fqns 

96 

97 

98class FSDP2Manager: 

99 """Configure and apply nested FSDP2 wrapping for a dual-mode model.""" 

100 

101 def __init__( 

102 self, 

103 config: FSDP2Config, 

104 mesh_context: MeshContext, 

105 fp32_main_params: bool = False, 

106 ) -> None: 

107 """Initialize the manager from resolved YAML config and runtime mesh. 

108 

109 Args: 

110 config: Resolved top-level ``fsdp_config``. 

111 mesh_context: Runtime topology containing the shared world mesh. 

112 fp32_main_params: Whether FSDP writes reduced gradients to 

113 ``main_grad`` for the fp32 main-parameter optimizer wrapper. 

114 """ 

115 self.config = config 

116 self.mesh_context = mesh_context 

117 self.fp32_main_params = fp32_main_params 

118 self._mp_policy: fully_shard_utils.MixedPrecisionPolicy | None = None 

119 

120 

121 def _resolve_replicate_params( 

122 self, 

123 model: ModuleClass, 

124 ) -> set[ParameterClass] | None: 

125 """Resolve configured replicate-parameter FQNs on the final model.""" 

126 if not self.config.replicate_params: 

127 return None 

128 

129 parameters_by_fqn = dict(model.named_parameters(remove_duplicate=False)) 

130 replicate_params = set() 

131 for parameter_fqn in self.config.replicate_params: 

132 parameter = parameters_by_fqn.get(parameter_fqn) 

133 if parameter is None: 

134 raise ValueError( 

135 "replicate_params contains unknown parameter FQN: " 

136 f"{parameter_fqn}" 

137 ) 

138 replicate_params.add(parameter) 

139 return replicate_params 

140 

141 def _build_expert_fsdp_mesh(self) -> DeviceMesh: 

142 """Select the expert FSDP child mesh.""" 

143 expert_mesh = self.mesh_context.fsdp_moe_mesh 

144 if expert_mesh is None: 

145 raise ValueError("Expert FSDP units require MeshContext.fsdp_moe_mesh") 

146 if "edp_replicate" in (expert_mesh.mesh_dim_names or ()): 

147 return expert_mesh[("edp_replicate", "edp_shard")] 

148 return expert_mesh["edp_shard"] 

149 

150 def _build_dense_fsdp_mesh(self) -> DeviceMesh | None: 

151 """Select the prebuilt dense FSDP child mesh when available.""" 

152 dense_mesh = self.mesh_context.fsdp_non_moe_mesh 

153 if dense_mesh is None: 

154 return None 

155 if self.mesh_context.dp_replicate_size > 1: 

156 return dense_mesh[("fsdp_replicate", "fsdp_shard")] 

157 return dense_mesh["fsdp_shard"] 

158 

159 def _build_compatibility_fsdp_mesh(self) -> DeviceMesh: 

160 """Build an FSDP mesh from the combined world mesh.""" 

161 world_mesh = self.mesh_context.device_mesh 

162 if world_mesh is None: 

163 raise ValueError("FSDP2 requires MeshContext.fsdp_non_moe_mesh") 

164 mesh_dim_names = world_mesh.mesh_dim_names 

165 if mesh_dim_names is None: 

166 raise ValueError("FSDP2 requires named DeviceMesh dimensions") 

167 shard_dim_names = tuple( 

168 dim_name for dim_name in ("dp_shard", "cp") if dim_name in mesh_dim_names 

169 ) 

170 if not shard_dim_names: 

171 raise ValueError("FSDP2 requires a dp_shard or cp dimension in the world mesh") 

172 shard_mesh = ( 

173 world_mesh[shard_dim_names[0]] 

174 if len(shard_dim_names) == 1 

175 else world_mesh[shard_dim_names].flatten("fsdp_shard") 

176 ) 

177 if "dp_replicate" not in mesh_dim_names: 

178 return shard_mesh 

179 return DeviceMesh.concatenate([world_mesh["dp_replicate"], shard_mesh]) 

180 

181 def _build_fsdp_actual_mesh(self, expert: bool = False) -> DeviceMesh: 

182 """Select the prebuilt dense or expert FSDP child mesh.""" 

183 if expert: 

184 return self._build_expert_fsdp_mesh() 

185 dense_mesh = self._build_dense_fsdp_mesh() 

186 if dense_mesh is not None: 

187 return dense_mesh 

188 return self._build_compatibility_fsdp_mesh() 

189 

190 def parameter_mesh_bindings( 

191 self, model: ModuleClass, source_shard_info: SourceShardInfoByFQN | None, 

192 ignored_params: set[ParameterClass] | None = None, 

193 ) -> dict[ParameterClass, DeviceMesh]: 

194 """Resolve unit-owned parameter meshes without installing FSDP hooks. 

195 

196 Args: 

197 model: Final model before graph capture or eager wrapping. 

198 source_shard_info: The planner's non-FSDP parameter placements. 

199 ignored_params: Parameters managed outside FSDP. 

200 

201 The same managed-parameter selector runs inside-out in both paths. 

202 Claimed identities model the initialized inner units, including tied 

203 aliases; nearest-parent ownership alone would not match fully_shard. 

204 """ 

205 metadata = _build_source_shard_info_by_param(self, model, source_shard_info) 

206 units = self._find_wrap_modules(model, metadata) 

207 units = [unit.module for unit in sorted(units, key=lambda unit: unit.fqn.count("."), reverse=True)] 

208 claimed = set(ignored_params or ()) 

209 bindings = {} 

210 for owner in [*units, model]: 

211 parameters = get_managed_modules_parameters((owner,), claimed) 

212 source = _build_managed_source_shard_info(self, owner, metadata, claimed) 

213 expert = self._uses_expert_mesh(source) 

214 if owner is model and expert: 

215 raise ValueError("Routed expert parameters must belong to a nested experts FSDP unit") 

216 mesh = self._build_fsdp_actual_mesh(expert) 

217 bindings.update(dict.fromkeys(parameters, mesh)) 

218 claimed.update(parameters) 

219 return bindings 

220 

221 def _resolve_custom_params( 

222 self, parameters_by_fqn: Mapping[str, ParameterClass] 

223 ) -> dict[ParameterClass, dict[str, Any]]: 

224 """Resolve configured names and patterns against the supplied parameter mapping.""" 

225 custom_params = {} 

226 for parameter_fqn, custom_dtypes in self.config.mix_precision.custom_params.items(): 

227 matched_fqns = [parameter_fqn] if parameter_fqn in parameters_by_fqn else [ 

228 fqn for fqn in parameters_by_fqn if fnmatch.fnmatchcase(fqn, parameter_fqn) 

229 ] 

230 if not matched_fqns: 

231 raise ValueError(f"mix_precision.custom_params matches no parameter: {parameter_fqn}") 

232 for matched_fqn in matched_fqns: 

233 parameter = parameters_by_fqn[matched_fqn] 

234 if parameter in custom_params and custom_params[parameter] != custom_dtypes: 

235 raise ValueError(f"mix_precision.custom_params has conflicting overrides for {matched_fqn}") 

236 custom_params[parameter] = dict(custom_dtypes) 

237 return custom_params 

238 

239 def _build_mixed_precision_policy( 

240 self, custom_params: Mapping[ParameterClass, Mapping[str, Any]] | None = None 

241 ) -> fully_shard_utils.MixedPrecisionPolicy: 

242 """Build or reuse the shared precision policy for the current model.""" 

243 if custom_params is None and self._mp_policy is not None: 

244 return self._mp_policy 

245 mix_precision = self.config.mix_precision 

246 dtype_by_name = { 

247 None: None, 

248 "bfloat16": torch.bfloat16, 

249 "float16": torch.float16, 

250 "float32": torch.float32, 

251 } 

252 custom_params = { 

253 parameter: {name: dtype_by_name[dtype] for name, dtype in custom_dtypes.items()} 

254 for parameter, custom_dtypes in (custom_params or {}).items() 

255 } 

256 self._mp_policy = fully_shard_utils.MixedPrecisionPolicy( 

257 param_dtype=dtype_by_name[mix_precision.param_dtype], 

258 reduce_dtype=dtype_by_name[mix_precision.reduce_dtype], 

259 output_dtype=dtype_by_name[mix_precision.output_dtype], 

260 cast_forward_inputs=mix_precision.cast_forward_inputs, 

261 apply_grad_on_fp32_main_grad=self.fp32_main_params, 

262 custom_params=custom_params if custom_params else None, 

263 ) 

264 return self._mp_policy 

265 

266 def _build_offload_policy(self) -> fully_shard_utils.OffloadPolicy: 

267 """Return CPU offload when enabled, otherwise the default no-offload policy.""" 

268 if not self.config.enable_offload: 

269 return fully_shard_utils.OffloadPolicy() 

270 return fully_shard_utils.CPUOffloadPolicy() 

271 

272 def _build_fully_shard_kwargs( 

273 self, 

274 fsdp_mesh: DeviceMesh, 

275 ) -> tuple[dict[str, Any], dict[str, Any]]: 

276 """Build child and root kwargs with the same configured reshard policy.""" 

277 mp_policy = self._build_mixed_precision_policy() 

278 offload_policy = self._build_offload_policy() 

279 

280 fsdp_sublayer_kwargs = { 

281 "mesh": fsdp_mesh, 

282 "reshard_after_forward": self.config.reshard_after_forward, 

283 "mp_policy": mp_policy, 

284 "offload_policy": offload_policy, 

285 "comm_fusion": self.config.comm_fusion, 

286 "comm_fusion_zero_copy": self.config.comm_fusion_zero_copy, 

287 } 

288 fsdp_root_kwargs = fsdp_sublayer_kwargs.copy() 

289 return fsdp_sublayer_kwargs, fsdp_root_kwargs 

290 

291 @staticmethod 

292 def _find_transformer_block_modules( 

293 model: ModuleClass, 

294 excluded_subtree_module_fqns: set[str] | None = None, 

295 ) -> tuple[list[_WrapModuleInfo], set[int]]: 

296 """Find transformer blocks under HF gradient-checkpointing containers. 

297 

298 Args: 

299 model: Root model to inspect. 

300 excluded_subtree_module_fqns: Roots declared by a model adapter as 

301 non-decoder branches or self-contained FSDP units. Automatic 

302 decoder discovery skips these roots and their descendants. 

303 """ 

304 wrap_modules = [] 

305 wrapped_module_ids = set() 

306 excluded_subtree_module_fqns = excluded_subtree_module_fqns or set() 

307 

308 def is_inside_declared_subtree(module_fqn: str) -> bool: 

309 """Whether one module is inside an automatic-discovery exclusion. 

310 

311 Args: 

312 module_fqn: Fully qualified name of the module. 

313 """ 

314 return any( 

315 module_fqn == declared_fqn 

316 or module_fqn.startswith(f"{declared_fqn}.") 

317 for declared_fqn in excluded_subtree_module_fqns 

318 ) 

319 

320 for container_fqn, container in model.named_modules(): 

321 if container_fqn and is_inside_declared_subtree(container_fqn): 

322 continue 

323 if id(container) in wrapped_module_ids: 

324 continue 

325 if not hasattr(container, "gradient_checkpointing"): 

326 continue 

327 for child_name, child in container.named_children(): 

328 child_fqn = f"{container_fqn}.{child_name}" if container_fqn else child_name 

329 if is_inside_declared_subtree(child_fqn): 

330 continue 

331 blocks = list(child.named_children()) 

332 if not blocks: 

333 continue 

334 for block_name, block in blocks: 

335 if id(block) in wrapped_module_ids: 

336 continue 

337 wrapped_module_ids.add(id(block)) 

338 wrapped_module = _get_checkpoint_wrapped_module(block) 

339 if wrapped_module is not None: 

340 # The wrapper and its direct child represent one logical 

341 # transformer block during module-tree traversal. 

342 wrapped_module_ids.add(id(wrapped_module)) 

343 wrap_modules.append( 

344 _WrapModuleInfo(f"{child_fqn}.{block_name}", block) 

345 ) 

346 return wrap_modules, wrapped_module_ids 

347 

348 def _find_expert_wrap_modules( 

349 self, 

350 model: ModuleClass, 

351 metadata_by_parameter: SourceShardInfoByParam | None, 

352 wrapped_module_ids: set[int], 

353 ) -> list[_WrapModuleInfo]: 

354 """Find nested expert containers whose parameters use the expert mesh.""" 

355 if metadata_by_parameter is None or self.mesh_context.fsdp_moe_mesh is None: 

356 return [] 

357 expert_parameters = { 

358 parameter 

359 for parameter, source_shard_info in metadata_by_parameter.items() 

360 if _is_expert_source_mesh(source_shard_info.mesh) 

361 } 

362 expert_module_fqns = { 

363 parameter_fqn.split(".experts.", 1)[0] + ".experts" 

364 for parameter_fqn, parameter in model.named_parameters(remove_duplicate=False) 

365 if parameter in expert_parameters and ".experts." in parameter_fqn 

366 } 

367 expert_module_fqns.update(_declared_expert_unit_fqns(model, expert_parameters)) 

368 expert_module_fqns.discard("") 

369 module_by_fqn = dict(model.named_modules()) 

370 expert_wrap_modules = [] 

371 for expert_module_fqn in sorted(expert_module_fqns): 

372 expert_module = module_by_fqn.get(expert_module_fqn) 

373 if expert_module is None or id(expert_module) in wrapped_module_ids: 

374 continue 

375 wrapped_module_ids.add(id(expert_module)) 

376 expert_wrap_modules.append( 

377 _WrapModuleInfo(expert_module_fqn, expert_module) 

378 ) 

379 return expert_wrap_modules 

380 

381 @staticmethod 

382 def _get_model_adapter_spec(model: ModuleClass) -> Any | None: 

383 """Resolve a model adapter from the final model configuration.""" 

384 if hasattr(model, "_hp_model_adapter_spec"): 

385 return model._hp_model_adapter_spec # pylint: disable=protected-access 

386 config = getattr(model, "config", None) 

387 if config is None: 

388 return None 

389 identities = ( 

390 getattr(config, "model_type", ""), 

391 *((getattr(config, "architectures", None) or [])[:1]), 

392 ) 

393 for identity in identities: 

394 if identity: 

395 adapter_spec = get_model_adapter(identity) 

396 if adapter_spec is not None: 

397 return adapter_spec 

398 return None 

399 

400 @classmethod 

401 def _find_adapter_declared_wrap_modules( 

402 cls, 

403 model: ModuleClass, 

404 ) -> list[_WrapModuleInfo]: 

405 """Resolve adapter-declared FSDP units outside HF decoder containers. 

406 

407 ``ModelAdapterSpec.fsdp_wrap_modules`` keeps architecture-specific 

408 execution boundaries on the family adapter side. It receives the 

409 finalized model and returns exact root-relative module FQNs. 

410 

411 Args: 

412 model: Root model being wrapped. 

413 

414 Returns: 

415 Additional FSDP child units selected by model declarations. 

416 

417 Raises: 

418 ValueError: If a declaration has an invalid type, targets no 

419 module in the model tree, or aliases the same module twice. 

420 """ 

421 adapter_spec = cls._get_model_adapter_spec(model) 

422 if adapter_spec is None or adapter_spec.fsdp_wrap_modules is None: 

423 return [] 

424 

425 declared_fqns = adapter_spec.fsdp_wrap_modules(model) 

426 if isinstance(declared_fqns, str) or not isinstance(declared_fqns, Sequence): 

427 raise ValueError( 

428 "ModelAdapterSpec.fsdp_wrap_modules must return a sequence of exact module FQNs" 

429 ) 

430 module_by_fqn = dict(model.named_modules()) 

431 declared_wrap_modules = [] 

432 wrapped_module_ids = set() 

433 for module_fqn in declared_fqns: 

434 if not isinstance(module_fqn, str) or not module_fqn: 

435 raise ValueError( 

436 "ModelAdapterSpec.fsdp_wrap_modules entries must be non-empty strings" 

437 ) 

438 wrap_module = module_by_fqn.get(module_fqn) 

439 if wrap_module is None: 

440 raise ValueError( 

441 "ModelAdapterSpec.fsdp_wrap_modules declared an unknown module: " 

442 f"{module_fqn}" 

443 ) 

444 if id(wrap_module) in wrapped_module_ids: 

445 raise ValueError( 

446 "ModelAdapterSpec.fsdp_wrap_modules aliases one module more than once: " 

447 f"{module_fqn}" 

448 ) 

449 wrapped_module_ids.add(id(wrap_module)) 

450 declared_wrap_modules.append(_WrapModuleInfo(module_fqn, wrap_module)) 

451 return declared_wrap_modules 

452 

453 @classmethod 

454 def _find_adapter_excluded_subtrees(cls, model: ModuleClass) -> set[str]: 

455 """Resolve model branches excluded from automatic decoder discovery.""" 

456 adapter_spec = cls._get_model_adapter_spec(model) 

457 excluded_subtrees_provider = getattr(adapter_spec, "fsdp_excluded_subtrees", None) 

458 if excluded_subtrees_provider is None: 

459 return set() 

460 if not callable(excluded_subtrees_provider): 

461 raise ValueError("ModelAdapterSpec.fsdp_excluded_subtrees must be callable") 

462 excluded_fqns = excluded_subtrees_provider(model) # pylint: disable=not-callable 

463 if isinstance(excluded_fqns, str) or not isinstance(excluded_fqns, Sequence): 

464 raise ValueError( 

465 "ModelAdapterSpec.fsdp_excluded_subtrees must return a sequence of exact module FQNs" 

466 ) 

467 module_by_fqn = dict(model.named_modules()) 

468 for module_fqn in excluded_fqns: 

469 if not isinstance(module_fqn, str) or not module_fqn: 

470 raise ValueError( 

471 "ModelAdapterSpec.fsdp_excluded_subtrees entries must be non-empty strings" 

472 ) 

473 if module_fqn not in module_by_fqn: 

474 raise ValueError( 

475 "ModelAdapterSpec.fsdp_excluded_subtrees declared an unknown module: " 

476 f"{module_fqn}" 

477 ) 

478 return set(excluded_fqns) 

479 

480 def _find_class_wrap_modules(self, model: ModuleClass) -> list[_WrapModuleInfo]: 

481 """Select explicit classes, retaining checkpoint wrappers as the units.""" 

482 requested = set(self.config.wrapper_class) 

483 matched = set() 

484 selected = [] 

485 seen = set() 

486 aliases = defaultdict(list) 

487 for name, parameter in model.named_parameters(remove_duplicate=False): 

488 aliases[parameter].append(name) 

489 for fqn, module in model.named_modules(): 

490 if not fqn or id(module) in seen: 

491 continue 

492 underlying = _get_checkpoint_wrapped_module(module) 

493 candidate = underlying if underlying is not None else module 

494 class_name = type(candidate).__name__ 

495 if class_name not in requested: 

496 continue 

497 matched.add(class_name) 

498 parameters = list(module.parameters()) 

499 if not parameters: 

500 raise ValueError(f"wrapper_class selected a parameterless module: {fqn}") 

501 if any(not name.startswith(f"{fqn}.") for param in parameters for name in aliases[param]): 

502 raise ValueError(f"wrapper_class module shares parameters outside its subtree: {fqn}") 

503 selected.append(_WrapModuleInfo(fqn, module)) 

504 seen.add(id(module)) 

505 seen.add(id(candidate)) 

506 if requested - matched: 

507 raise ValueError(f"wrapper_class matched no child module for: {sorted(requested - matched)}") 

508 return selected 

509 

510 def _find_wrap_modules( 

511 self, 

512 model: ModuleClass, 

513 metadata_by_parameter: SourceShardInfoByParam | None = None, 

514 ) -> list[_WrapModuleInfo]: 

515 """Find automatic decoder/expert units and adapter-declared units.""" 

516 if self.config.wrapper_class: 

517 wrap_modules = self._find_class_wrap_modules(model) 

518 wrapped_module_ids = {id(unit.module) for unit in wrap_modules} 

519 else: 

520 declared_wrap_modules = self._find_adapter_declared_wrap_modules(model) 

521 if declared_wrap_modules: 

522 logger.info( 

523 "Adapter-declared FSDP child units: %s", 

524 [wrap_module.fqn for wrap_module in declared_wrap_modules], 

525 ) 

526 excluded_subtree_module_fqns = { 

527 wrap_module.fqn for wrap_module in declared_wrap_modules 

528 } 

529 excluded_subtree_module_fqns.update( 

530 self._find_adapter_excluded_subtrees(model) 

531 ) 

532 wrap_modules, wrapped_module_ids = self._find_transformer_block_modules( 

533 model, 

534 excluded_subtree_module_fqns=excluded_subtree_module_fqns, 

535 ) 

536 wrapped_module_ids.update( 

537 id(wrap_module.module) for wrap_module in declared_wrap_modules 

538 ) 

539 wrap_modules.extend(declared_wrap_modules) 

540 wrap_modules.extend( 

541 self._find_expert_wrap_modules( 

542 model, 

543 metadata_by_parameter, 

544 wrapped_module_ids, 

545 ) 

546 ) 

547 if not wrap_modules: 

548 raise ValueError( 

549 "wrap_policy='transformer_block' found no decoder, expert, or " 

550 "adapter-declared FSDP child unit" 

551 ) 

552 return wrap_modules 

553 

554 @staticmethod 

555 def _resolve_parameter_owners( 

556 model: ModuleClass, 

557 wrap_modules: list[_WrapModuleInfo], 

558 ) -> dict[ParameterClass, ModuleClass]: 

559 """Assign each parameter to its deepest common FSDP wrap module.""" 

560 aliases_by_parameter: dict[ParameterClass, list[str]] = defaultdict(list) 

561 for parameter_fqn, parameter in model.named_parameters(remove_duplicate=False): 

562 aliases_by_parameter[parameter].append(parameter_fqn) 

563 

564 owner_by_parameter = {} 

565 for parameter, parameter_fqns in aliases_by_parameter.items(): 

566 common_wrap_modules = [ 

567 wrap_module 

568 for wrap_module in wrap_modules 

569 if all( 

570 parameter_fqn.startswith(f"{wrap_module.fqn}.") 

571 for parameter_fqn in parameter_fqns 

572 ) 

573 ] 

574 if common_wrap_modules: 

575 owner_by_parameter[parameter] = max( 

576 common_wrap_modules, 

577 key=lambda wrap_module: wrap_module.fqn.count("."), 

578 ).module 

579 else: 

580 owner_by_parameter[parameter] = model 

581 return owner_by_parameter 

582 

583 

584 def _uses_expert_mesh(self, metadata_by_parameter: SourceShardInfoByParam | None) -> bool: 

585 """Return whether one FSDP unit owns routed-expert source metadata.""" 

586 if metadata_by_parameter is None or self.mesh_context.fsdp_moe_mesh is None: 

587 return False 

588 return any( 

589 _is_expert_source_mesh(source_shard_info.mesh) 

590 for source_shard_info in metadata_by_parameter.values() 

591 ) 

592 

593 

594 @staticmethod 

595 def _build_managed_replicate_params( 

596 owner: ModuleClass, 

597 owner_by_parameter: Mapping[ParameterClass, ModuleClass], 

598 replicate_params: set[ParameterClass] | None, 

599 ignored_params: set[ParameterClass] | None = None, 

600 ) -> set[ParameterClass] | None: 

601 """Select configured replicate parameters owned by one FSDP module.""" 

602 if replicate_params is None: 

603 return None 

604 return { 

605 parameter 

606 for parameter, parameter_owner in owner_by_parameter.items() 

607 if parameter_owner is owner and parameter in replicate_params 

608 and parameter not in (ignored_params or ()) 

609 } 

610 

611 def _configure_source_layout_gradient_scaling( 

612 self, 

613 hsdp_module: HSDPModule, 

614 managed_source_shard_info: SourceShardInfoByParam | None, 

615 ) -> bool: 

616 """Restore global-mean gradients for one source-layout FSDP unit. 

617 

618 Master defaults every ``fully_shard`` unit to AVG reduction, so units 

619 with source-layout metadata explicitly select SUM. The token-weighted 

620 backward loss is scaled by the FSDP data domain (DP x CP), so a 

621 source-layout SUM-reduced unit must apply its inverse exactly once. 

622 With loss parallelism, TP ranks hold shards or partial contributions 

623 to one logical loss and TP is not another averaging axis. Without loss 

624 parallelism, logits and the loss are replicated over TP, so the SUM 

625 path also accumulates ``tp_size`` identical loss replicas and must 

626 divide by TP. EP is excluded because expert ranks own different 

627 parameters. 

628 

629 Returns: 

630 Whether a gradient scaling factor was configured for the unit. 

631 """ 

632 if not managed_source_shard_info: 

633 return False 

634 hsdp_module.set_reduce_op_type("sum", recurse=False) 

635 

636 tp_loss_replica_size = ( 

637 1 if self.mesh_context.loss_parallel else self.mesh_context.tp_size 

638 ) 

639 source_gradient_domain_size = ( 

640 self.mesh_context.dp_size 

641 * self.mesh_context.cp_size 

642 * tp_loss_replica_size 

643 ) 

644 if source_gradient_domain_size <= 1: 

645 return False 

646 hsdp_module.set_gradient_scaling_factor(1.0 / source_gradient_domain_size) 

647 return True 

648 

649 def _configure_prefetch(self, hsdp_modules: list[HSDPModule]) -> None: 

650 """Configure forward and backward prefetch targets by list depth.""" 

651 for module_index, hsdp_module in enumerate(hsdp_modules): 

652 forward_end = module_index + self.config.forward_prefetch_depth + 1 

653 hsdp_module.set_modules_to_forward_prefetch( 

654 hsdp_modules[module_index + 1:forward_end] 

655 ) 

656 

657 backward_start = max( 

658 0, module_index - self.config.backward_prefetch_depth 

659 ) 

660 hsdp_module.set_modules_to_backward_prefetch( 

661 list(reversed(hsdp_modules[backward_start:module_index])) 

662 ) 

663 

664 def _order_wrap_modules( 

665 self, 

666 model: ModuleClass, 

667 wrap_modules: list[_WrapModuleInfo], 

668 ) -> list[_WrapModuleInfo]: 

669 """Order FSDP units by adapter execution order or module traversal.""" 

670 module_order = { 

671 id(module): module_index 

672 for module_index, (_, module) in enumerate(model.named_modules()) 

673 } 

674 default_order = sorted( 

675 wrap_modules, 

676 key=lambda wrap_module: module_order.get(id(wrap_module.module), len(module_order)), 

677 ) 

678 adapter_spec = self._get_model_adapter_spec(model) 

679 execution_order_provider = getattr(adapter_spec, "fsdp_execution_order", None) 

680 if execution_order_provider is None: 

681 return default_order 

682 if not callable(execution_order_provider): 

683 raise ValueError("ModelAdapterSpec.fsdp_execution_order must be callable") 

684 

685 wrap_module_by_fqn = { 

686 wrap_module.fqn: wrap_module for wrap_module in wrap_modules 

687 } 

688 execution_fqns = execution_order_provider( # pylint: disable=not-callable 

689 model, 

690 tuple(wrap_module.fqn for wrap_module in default_order), 

691 ) 

692 if isinstance(execution_fqns, str) or not isinstance(execution_fqns, Sequence): 

693 raise ValueError( 

694 "ModelAdapterSpec.fsdp_execution_order must return a sequence of module FQNs" 

695 ) 

696 if any(not isinstance(module_fqn, str) for module_fqn in execution_fqns): 

697 raise ValueError( 

698 "ModelAdapterSpec.fsdp_execution_order entries must be module FQN strings" 

699 ) 

700 if ( 

701 len(execution_fqns) != len(wrap_module_by_fqn) 

702 or len(set(execution_fqns)) != len(execution_fqns) 

703 or set(execution_fqns) != set(wrap_module_by_fqn) 

704 ): 

705 raise ValueError( 

706 "ModelAdapterSpec.fsdp_execution_order must return every selected FSDP " 

707 "child-unit FQN exactly once" 

708 ) 

709 logger.info("Adapter-declared FSDP execution order: %s", list(execution_fqns)) 

710 return [wrap_module_by_fqn[module_fqn] for module_fqn in execution_fqns] 

711 

712 def _parallelize_child_units( 

713 self, 

714 wrap_modules: list[_WrapModuleInfo], 

715 owner_by_parameter, 

716 metadata_by_parameter, 

717 replicate_params, 

718 dense_fsdp_kwargs: dict[str, Any], 

719 ignored_params: set[ParameterClass] | None = None, 

720 ) -> int: 

721 """Apply FSDP to child units and return the gradient-scaled count.""" 

722 gradient_scaled_units = 0 

723 wrapping_order = sorted( 

724 wrap_modules, key=lambda wrap_module: wrap_module.fqn.count("."), reverse=True 

725 ) 

726 for wrap_module in wrapping_order: 

727 managed_source_info = _build_managed_source_shard_info( 

728 self, 

729 wrap_module.module, 

730 metadata_by_parameter, 

731 ignored_params, 

732 ) 

733 fsdp_sublayer_kwargs = dense_fsdp_kwargs 

734 if self._uses_expert_mesh(managed_source_info): 

735 if self.mesh_context.fsdp_moe_mesh is None: 

736 raise ValueError("Expert TP metadata requires MeshContext.fsdp_moe_mesh") 

737 fsdp_sublayer_kwargs, _ = self._build_fully_shard_kwargs( 

738 self._build_fsdp_actual_mesh(expert=True) 

739 ) 

740 fully_shard( # pylint: disable=unexpected-keyword-arg 

741 wrap_module.module, 

742 source_shard_infos=_source_infos_for_fully_shard(managed_source_info), 

743 replicate_params=self._build_managed_replicate_params( 

744 wrap_module.module, owner_by_parameter, replicate_params, ignored_params 

745 ), 

746 ignored_params=ignored_params, 

747 **fsdp_sublayer_kwargs, 

748 ) 

749 if self._configure_source_layout_gradient_scaling(wrap_module.module, managed_source_info): 

750 gradient_scaled_units += 1 

751 return gradient_scaled_units 

752 

753 def _parallelize_root( 

754 self, 

755 model: ModuleClass, 

756 owner_by_parameter, 

757 metadata_by_parameter, 

758 replicate_params, 

759 dense_root_kwargs: dict[str, Any], 

760 ignored_params: set[ParameterClass] | None = None, 

761 ) -> int: 

762 """Apply root FSDP and report whether gradient scaling was configured.""" 

763 root_source_info = _build_managed_source_shard_info( 

764 self, 

765 model, 

766 metadata_by_parameter, 

767 ignored_params, 

768 ) 

769 if self._uses_expert_mesh(root_source_info): 

770 raise ValueError("Routed expert parameters must belong to a nested experts FSDP unit") 

771 fully_shard( # pylint: disable=unexpected-keyword-arg 

772 model, 

773 source_shard_infos=_source_infos_for_fully_shard(root_source_info), 

774 replicate_params=self._build_managed_replicate_params( 

775 model, owner_by_parameter, replicate_params, ignored_params 

776 ), 

777 ignored_params=ignored_params, 

778 **dense_root_kwargs, 

779 ) 

780 return int(self._configure_source_layout_gradient_scaling(model, root_source_info)) 

781 

782 def _log_gradient_scaling(self, gradient_scaled_units: int) -> None: 

783 """Log the source-layout gradient scaling applied to FSDP units.""" 

784 if not gradient_scaled_units: 

785 return 

786 domain_size = ( 

787 self.mesh_context.dp_size 

788 * self.mesh_context.cp_size 

789 * (1 if self.mesh_context.loss_parallel else self.mesh_context.tp_size) 

790 ) 

791 logger.info( 

792 "Applied source-layout SUM gradient scaling: dp_size=%d, cp_size=%d, tp_size=%d, " 

793 "source_gradient_domain_size=%d, factor=%s, units=%d", 

794 self.mesh_context.dp_size, 

795 self.mesh_context.cp_size, 

796 self.mesh_context.tp_size, 

797 domain_size, 

798 1.0 / domain_size, 

799 gradient_scaled_units, 

800 ) 

801 

802 def parallelize( 

803 self, 

804 model: ModuleClass, 

805 source_shard_info: SourceShardInfoByFQN | None = None, 

806 ignored_params: set[ParameterClass] | None = None, 

807 ) -> ModuleClass: 

808 """Apply configured child FSDP units, root FSDP, and prefetch links. 

809 

810 Args: 

811 model: Model to wrap in place. 

812 source_shard_info: FQN-keyed source-layout metadata returned by 

813 ``apply_sharding_plan``. The manager resolves parameter 

814 identities before applying nested FSDP units. 

815 

816 Returns: 

817 The model enhanced with the HyperParallel HSDPModule interface. 

818 

819 Raises: 

820 ValueError: If TP is enabled without FQN metadata or configured 

821 parameter FQNs cannot be resolved. 

822 """ 

823 self._mp_policy = None 

824 metadata_by_parameter = _build_source_shard_info_by_param( 

825 self, 

826 model, 

827 source_shard_info, 

828 ) 

829 replicate_params = self._resolve_replicate_params(model) 

830 custom_params = None 

831 if self.config.mix_precision.custom_params: 

832 custom_params = self._resolve_custom_params(dict(model.named_parameters(remove_duplicate=False))) 

833 self._build_mixed_precision_policy(custom_params) 

834 

835 wrap_modules = self._find_wrap_modules(model, metadata_by_parameter) 

836 owner_by_parameter = self._resolve_parameter_owners(model, wrap_modules)