Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / parallelism.py: 91%

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"""Parallel topology, activation-checkpoint and plan-override configuration. 

16 

17Split from ``auto_models/trainer/config.py`` in stage 7 (05 §15.2.5). 

18``PlanOverride`` stays the YAML DTO that desugars to 

19``auto_models.distributed.recipe_spec.ModuleShardingSpec``. 

20""" 

21# pylint: disable=unused-import 

22 

23import importlib 

24import logging 

25from collections.abc import Mapping 

26from dataclasses import dataclass, field 

27from types import MappingProxyType 

28from typing import Any, List, Literal, Optional, Union 

29 

30from torch import nn # pylint: disable=forbidden-backend-import 

31 

32from hyper_parallel.components.quantization.config import ( 

33 LowPrecisionConfig, 

34 LowPrecisionDtypeScheme, 

35) 

36from hyper_parallel.distributed.encoder_parallel.config import EncoderParallelConfig 

37from hyper_parallel.models.replacement import ( 

38 MethodReplacementSpec, 

39 ModuleReplacementFactory, 

40 ModuleReplacementSpec, 

41 ModuleReplacementRequest, 

42 method_replacement, 

43 module_replacement, 

44) 

45 

46from hyper_parallel.components.modules.family import get_module_family 

47from hyper_parallel.trainer.config.replacement import ReplacementTarget 

48from hyper_parallel.trainer.config.target import Target, _serialize_config_value 

49 

50logger = logging.getLogger(__name__) 

51 

52 

53@dataclass 

54class AcceleratorConfig: 

55 """Parallel topology and target-selected strategy configurations.""" 

56 

57 tp_size: int = 1 

58 cp_size: int = 1 

59 ep_size: int = 1 

60 pp_size: int = 1 

61 sequence_parallel: bool = False 

62 loss_parallel: bool = False 

63 encoder_parallel: EncoderParallelConfig = field(default_factory=EncoderParallelConfig) 

64 

65 

66@dataclass 

67class ActivationCheckpointSelection: 

68 """Select adapter-declared recompute regions and exact layer coverage. 

69 

70 ``layer_count`` selects the contiguous prefix ``[0, layer_count)`` from 

71 the flattened layer-container order. ``layer_indices`` selects explicit 

72 zero-based entries from that same order. 

73 

74 ``model_adapter_safe_regions`` wraps the outermost module matching the 

75 adapter's ``region_patterns`` inside every selected layer and keeps every 

76 ``exclude_patterns`` descendant outside the recompute region. Which shape that 

77 produces is the model's declaration, not the Trainer's: leaf patterns give 

78 one region per submodule, a layer pattern gives one region per layer. 

79 """ 

80 

81 source: Literal["default", "model_adapter_safe_regions"] = "default" 

82 layer_count: Optional[int] = None 

83 layer_indices: Optional[List[int]] = None 

84 

85 def __post_init__(self) -> None: 

86 """Validate the mutually exclusive layer-count and layer-index selectors.""" 

87 self._validate_selector_choice() 

88 self._validate_selector_values() 

89 

90 def _validate_selector_choice(self) -> None: 

91 """Require one supported source and its matching selector form.""" 

92 if self.source not in ("default", "model_adapter_safe_regions"): 

93 raise ValueError( 

94 "activation_checkpoint.selection.source must be default or " 

95 "model_adapter_safe_regions" 

96 ) 

97 selectors = (self.layer_count is not None, self.layer_indices is not None) 

98 if self.source == "default" and any(selectors): 

99 raise ValueError( 

100 "activation_checkpoint.selection.layer_count and layer_indices " 

101 "are only valid when " 

102 "selection.source=model_adapter_safe_regions" 

103 ) 

104 if self.source == "model_adapter_safe_regions" and sum(selectors) != 1: 

105 raise ValueError( 

106 "activation_checkpoint.selection requires exactly one of " 

107 "layer_count or layer_indices when " 

108 "selection.source=model_adapter_safe_regions" 

109 ) 

110 

111 def _validate_selector_values(self) -> None: 

112 """Require non-negative, distinct layer numbers.""" 

113 if self.layer_count is not None and ( 

114 isinstance(self.layer_count, bool) 

115 or not isinstance(self.layer_count, int) 

116 or self.layer_count < 0 

117 ): 

118 raise ValueError( 

119 "activation_checkpoint.selection.layer_count must be a " 

120 "non-negative integer" 

121 ) 

122 if self.layer_indices is None: 

123 return 

124 if not isinstance(self.layer_indices, list) or not self.layer_indices: 

125 raise ValueError( 

126 "activation_checkpoint.selection.layer_indices must be a " 

127 "non-empty list of non-negative integers" 

128 ) 

129 if any( 

130 isinstance(index, bool) or not isinstance(index, int) or index < 0 

131 for index in self.layer_indices 

132 ): 

133 raise ValueError( 

134 "activation_checkpoint.selection.layer_indices must contain " 

135 "only non-negative integers" 

136 ) 

137 if len(set(self.layer_indices)) != len(self.layer_indices): 

138 raise ValueError( 

139 "activation_checkpoint.selection.layer_indices must not " 

140 "contain duplicates" 

141 ) 

142 

143 

144@dataclass 

145class ActivationCheckpointConfig: 

146 """Activation-checkpoint options exposed by the initial YAML schema. 

147 

148 ``swap_inputs`` is consumed only when ``mode`` is ``"full"`` or 

149 ``"selective"``. 

150 """ 

151 

152 mode: Optional[Literal["off", "full", "selective"]] = "off" 

153 swap_inputs: bool = False 

154 selection: ActivationCheckpointSelection = field( 

155 default_factory=ActivationCheckpointSelection 

156 ) 

157 # Escape hatch for SAC's cached-tensor mutation check: safe when 

158 # every in-place rewrite of a cached operator output is 

159 # deterministic, so the recompute replays the same values (the 

160 # CSA indexer's masked_fill_ family). 

161 allow_cache_entry_mutation: bool = False 

162 

163 def __post_init__(self) -> None: 

164 """Reject ambiguous values for activation input swapping.""" 

165 if not isinstance(self.swap_inputs, bool): 

166 raise TypeError("activation_checkpoint.swap_inputs must be a bool") 

167 if not isinstance(self.allow_cache_entry_mutation, bool): 

168 raise TypeError( 

169 "activation_checkpoint.allow_cache_entry_mutation must be a bool") 

170 

171 

172def _resolve_replace_implementation(value: Union[str, Target[Any]]) -> Target[Any]: 

173 """Resolve a replacement implementation without invoking its callable.""" 

174 if isinstance(value, Target): 

175 return value 

176 if not isinstance(value, str) or not value.strip(): 

177 raise TypeError("replace.implementation must be a non-empty import path or Target") 

178 parts = value.split(".") 

179 for split_at in range(len(parts), 0, -1): 

180 module_name = ".".join(parts[:split_at]) 

181 try: 

182 symbol = importlib.import_module(module_name) 

183 except ModuleNotFoundError as exc: 

184 if exc.name == module_name or module_name.startswith(f"{exc.name}."): 

185 continue 

186 raise ValueError(f"replace.implementation {value!r} failed while importing {exc.name!r}") from exc 

187 except ImportError as exc: 

188 raise ValueError(f"replace.implementation {value!r} failed while importing: {exc}") from exc 

189 for attribute in parts[split_at:]: 

190 if not hasattr(symbol, attribute): 

191 raise ValueError(f"replace.implementation {value!r} has no attribute {attribute!r}") 

192 symbol = getattr(symbol, attribute) 

193 if not callable(symbol): 

194 raise TypeError(f"replace.implementation {value!r} must name a callable") 

195 return Target(symbol, target_path=value) 

196 raise ValueError(f"replace.implementation {value!r} could not be imported") 

197 

198 

199@dataclass(frozen=True) 

200class ReplaceAction: 

201 """One public module or method replacement declaration.""" 

202 

203 kind: Literal["module", "method"] 

204 implementation: Union[str, Target[Any]] 

205 source_type: Optional[str] = None 

206 exact_type: bool = False 

207 

208 def __post_init__(self) -> None: 

209 """Validate kind-specific fields and normalize the delayed target.""" 

210 if self.kind not in ("module", "method"): 

211 raise ValueError("replace.kind must be 'module' or 'method'") 

212 if not isinstance(self.exact_type, bool): 

213 raise TypeError("replace.exact_type must be a bool") 

214 if self.kind == "module": 

215 if self.source_type is not None and ( 

216 not isinstance(self.source_type, str) or not self.source_type.strip()): 

217 raise ValueError("replace.source_type must be a non-empty import path when provided") 

218 if self.exact_type and self.source_type is None: 

219 raise ValueError("replace.exact_type requires replace.source_type") 

220 elif self.source_type is not None or self.exact_type: 

221 raise ValueError("replace.source_type and replace.exact_type are only valid for kind='module'") 

222 object.__setattr__( 

223 self, "implementation", _resolve_replace_implementation(self.implementation) 

224 ) 

225 

226 

227@dataclass 

228class PlanOverride: 

229 """One plan_overrides entry (YAML ``plan_overrides`` list item). 

230 

231 This is the **YAML transport form** of ``ShardingPlanner(plan_overrides= 

232 {...})`` — the planner's single override interface. Placement objects are 

233 not YAML-serializable, so contract fields use the string DSL 

234 (``"replicate"`` / ``"partial"`` / ``"shard(N)"``) plus the sentinels 

235 ``"auto"`` (explicit inherit) / ``"none"`` (explicit clear); they are 

236 desugared to real objects in ``to_override()`` — the planner never sees 

237 strings. The desugar happens trainer-side in 

238 ``instantiate_infrastructure`` BEFORE planner construction. 

239 

240 Fields: 

241 replace: public replacement action. ``kind='module'`` requires a 

242 decorated module factory at ``implementation``; optional 

243 ``source_type`` checks the expected source module type; 

244 ``kind='method'`` selects a literal method suffix in ``match``; 

245 ``.forward`` uses the existing boundary executor and defaults to 

246 ``region_dispatch=False``, while other methods use a decorated 

247 ``@method_replacement`` factory before sharding. 

248 match: module FQN/glob, or module glob plus a literal method suffix for 

249 ``kind='method'``. Boundary globs first merge matching planner boundaries; 

250 a glob with a concrete params/I/O contract also expands against 

251 the final model's ``named_modules`` and creates missing boundaries. 

252 A merge-only glob fails if it matches a real module that has no 

253 planner boundary, with a suggested contract skeleton. 

254 ``*`` spans dots, so ``"*.self_attn"`` hits 

255 ``model.layers.0.self_attn``. 

256 when: optional activation condition. Sharding actions accept ``"cp"`` 

257 (active when cp_size>1), ``"ep"`` (ep_size>1), 

258 ``"sequence_parallel"``, or ``"low_precision"`` (active when 

259 online low precision is enabled). Module replacement accepts only 

260 ``"low_precision"`` and cannot depend on the parallel topology. 

261 low_precision_dtype_scheme: a named ``training.low_precision. 

262 dtype_schemes`` entry bound to this match. Only valid on 

263 ``when="low_precision"`` replacement entries; the resolved scheme 

264 is injected into the replacement factory context so the strategy 

265 factory never resolves a catalog at replacement time. 

266 Forward replacements use ``@forward_replacement`` or ``@local_compute`` 

267 factories through ``replace.implementation``. 

268 region_dispatch: validate execution-mode declaration; forward method 

269 replacement defaults to ``false``. ``false`` = the region 

270 computation (the module's own forward or the injected function) 

271 cannot be dispatched: it contains communication primitives / 

272 custom kernels (e.g. an in-house EP-aware MoE all-to-all inside 

273 forward, or the in-repo CP/EP reference implementations), so it 

274 runs as a local-region skeleton black box; ``true`` = the 

275 injected code is pure standard operators (fused kernels / 

276 scripted-style optimizations), so under validate the dispatch 

277 passes through and out_src is genuinely checked. This field is 

278 not needed without an injection (ordinary boundaries pass through 

279 naturally). 

280 params / in_src / in_dst / out_src / out_dst: contract fields in the 

281 YAML form ``{name: {axis: placement_str}}`` (out_* also accept 

282 the scalar shorthand ``{axis: placement_str}``), or the sentinels 

283 ``"auto"`` / ``"none"``. Merge mode (match hits a derived 

284 boundary): an omitted field inherits and an explicit ``{}`` 

285 clears. Insert mode (misses every boundary): nothing can inherit; 

286 at least one concrete parameter/I/O field must declare the new 

287 boundary, and all fields required by that module's contract must 

288 be written explicitly. 

289 tp_divide_attrs: optional module-instance integer attributes divided 

290 exactly by the active TP size when the module forward runs on 

291 local tensors. Omit for no user adjustment; an explicit empty 

292 list clears an inherited glob declaration. 

293 """ 

294 

295 match: Union[str, List[str]] 

296 replace: Optional[ReplaceAction] = None 

297 when: Optional[Literal[ 

298 "cp", "ep", "sequence_parallel", "low_precision" 

299 ]] = None 

300 low_precision_dtype_scheme: Optional[str] = None 

301 region_dispatch: Optional[bool] = None 

302 params: Optional[Any] = None 

303 in_src: Optional[Any] = None 

304 in_dst: Optional[Any] = None 

305 out_src: Optional[Any] = None 

306 out_dst: Optional[Any] = None 

307 tp_divide_attrs: Optional[List[str]] = None 

308 

309 # Contract fields (string sentinels are resolved at planner merge time; 

310 # insert mode rejects them) 

311 _CONTRACT_FIELDS = ("params", "in_src", "in_dst", "out_src", "out_dst") 

312 _BOUNDARY_FIELDS = _CONTRACT_FIELDS + ("tp_divide_attrs",) 

313 

314 def __post_init__(self) -> None: 

315 """Validate the public replacement action and its boundary contract.""" 

316 if self.replace is None: 

317 return 

318 if not isinstance(self.replace, ReplaceAction): 

319 raise TypeError("plan_overrides.replace must be a ReplaceAction") 

320 if self.replace.kind == "module": 

321 self._validate_module_replacement() 

322 return 

323 self._validate_method_replacement() 

324 

325 def _validate_module_replacement(self) -> None: 

326 """Validate fields exclusive to a module replacement rule.""" 

327 if self.when not in (None, "low_precision"): 

328 raise ValueError("kind='module' only supports when='low_precision'") 

329 if self.region_dispatch is not None: 

330 raise ValueError("region_dispatch cannot appear with kind='module'") 

331 if any(getattr(self, name) is not None for name in self._BOUNDARY_FIELDS): 

332 raise ValueError("kind='module' must use its own plan_overrides rule") 

333 

334 def _validate_method_replacement(self) -> None: 

335 """Validate a literal method match and dispatch to its protocol checks.""" 

336 if isinstance(self.match, list): 

337 raise ValueError(f"kind={self.replace.kind!r} requires one string match per rule") 

338 if self.low_precision_dtype_scheme is not None: 

339 raise ValueError("low_precision_dtype_scheme is only valid for kind='module'") 

340 if not isinstance(self.match, str): 

341 raise TypeError("kind='method' match must be a string") 

342 owner_pattern, separator, method_name = self.match.rpartition(".") 

343 if (not separator or not owner_pattern or not method_name.isidentifier() 

344 or (method_name.startswith("__") and method_name.endswith("__"))): 

345 raise ValueError("kind='method' match must end in a literal method name") 

346 if method_name != "forward": 

347 self._validate_non_forward_method() 

348 return 

349 self._validate_forward_method() 

350 

351 def _validate_non_forward_method(self) -> None: 

352 """Validate a decorated method factory without a planner contract.""" 

353 if self.region_dispatch is not None or any( 

354 getattr(self, name) is not None for name in self._BOUNDARY_FIELDS): 

355 raise ValueError("kind='method' on a non-forward method cannot declare a boundary contract") 

356 if not getattr(self.replace.implementation.callable, "_hp_method_replacement", False): 

357 raise TypeError("kind='method' implementation must use @method_replacement") 

358 

359 def _validate_forward_method(self) -> None: 

360 """Validate a forward factory and its boundary dispatch declaration.""" 

361 self.region_dispatch = False if self.region_dispatch is None else self.region_dispatch 

362 if not isinstance(self.region_dispatch, bool): 

363 raise TypeError("region_dispatch must be a bool") 

364 implementation = self.replace.implementation 

365 meta = getattr(implementation.callable, "_injection_meta", None) 

366 if meta is None: 

367 raise TypeError( 

368 "kind='method' on forward implementation must use @forward_replacement or @local_compute" 

369 ) 

370 if meta.kind not in ("local_compute", "inner_wrapper"): 

371 raise TypeError(f"unsupported forward target metadata kind {meta.kind!r}") 

372 

373 def _serialized_replacement(self): 

374 """Serialize only the replacement action, independently of layout fields.""" 

375 implementation_config = _serialize_config_value(self.replace.implementation) 

376 implementation_value = ( 

377 implementation_config["_target_"] 

378 if isinstance(implementation_config, dict) and set(implementation_config) == {"_target_"} 

379 else implementation_config 

380 ) 

381 action = {"kind": self.replace.kind, "implementation": implementation_value} 

382 if self.replace.kind == "module": 

383 if self.replace.source_type is not None: 

384 action["source_type"] = self.replace.source_type 

385 if self.replace.exact_type: 

386 action["exact_type"] = True 

387 return action 

388 

389 def to_dict(self) -> dict[str, Any]: 

390 """Serialize the public declaration.""" 

391 result: dict[str, Any] = {"match": self.match} 

392 if self.when is not None: 

393 result["when"] = self.when 

394 if self.low_precision_dtype_scheme is not None: 

395 result["low_precision_dtype_scheme"] = self.low_precision_dtype_scheme 

396 if self.replace is not None: 

397 result["replace"] = self._serialized_replacement() 

398 for name in ("region_dispatch", "params", "in_src", "in_dst", "out_src", "out_dst", "tp_divide_attrs"): 

399 value = getattr(self, name) 

400 default_forward_dispatch = ( 

401 name == "region_dispatch" and value is False 

402 and self.replace is not None and self.replace.kind == "method" 

403 and self.match.rpartition(".")[2] == "forward" 

404 ) 

405 if value is not None and not default_forward_dispatch: 

406 result[name] = _serialize_config_value(value) 

407 return result 

408 

409 def _validate_override_match(self): 

410 """Reject method-only and empty declarations before creating a boundary.""" 

411 if self.replace is not None and self.replace.kind == "method" and ( 

412 self.match.rpartition(".")[2] != "forward"): 

413 raise ValueError("non-forward method replacements do not create planner boundaries") 

414 

415 if not self.match: 

416 raise ValueError( 

417 "plan_overrides entry is missing 'match' (an fqn or " 

418 "fqn glob such as '*.self_attn' or '*.mlp')") 

419 

420 def to_override(self) -> "tuple[str, Any]": 

421 """Desugar to a ``(match, ModuleShardingSpec)`` plan_overrides entry. 

422 

423 Placement DSL strings become real Placement objects here; the 

424 ``"auto"``/``"none"`` sentinels pass through AS STRINGS into the spec 

425 (the planner's merge resolves them; insert rejects them). Imports are 

426 lazy: trainer.config must not become an import-time dependency of 

427 components.distributed (the zero-dependency boundary is guarded by 

428 test_s5_zero_dep_lint). 

429 """ 

430 # Keep trainer.config outside the distributed component's import-time dependency graph. 

431 # pylint: disable-next=import-outside-toplevel 

432 from hyper_parallel.distributed.recipe_spec import ( 

433 ModuleShardingSpec, 

434 ) 

435 

436 self._validate_override_match() 

437 implementation = self.replace.implementation if self.replace is not None else None 

438 meta = getattr(implementation.callable, "_injection_meta", None) if implementation is not None else None 

439 spec = ModuleShardingSpec( 

440 local_compute_fn=implementation if meta is not None and meta.kind == "local_compute" else None, 

441 inner_target="self" if meta is not None and meta.kind == "inner_wrapper" else None, 

442 inner_wrapper=implementation if meta is not None and meta.kind == "inner_wrapper" else None, 

443 region_dispatch=self.region_dispatch, 

444 tp_divide_attrs=self._validate_tp_divide_attrs(), 

445 ) 

446 for attr in self._CONTRACT_FIELDS: 

447 raw = getattr(self, attr) 

448 if raw is None: 

449 continue 

450 setattr(spec, attr, self._parse_contract_field(attr, raw)) 

451 match = ( 

452 self.match.rpartition(".")[0] 

453 if self.replace is not None and self.replace.kind == "method" 

454 else self.match 

455 ) 

456 return match, spec 

457 

458 def _validate_tp_divide_attrs(self) -> Optional[List[str]]: 

459 """Validate the YAML transport shape for TP-local attributes.""" 

460 attrs = self.tp_divide_attrs 

461 if attrs is None: 

462 return None 

463 if not isinstance(attrs, list): 

464 raise ValueError( 

465 f"plan_overrides match={self.match!r}: tp_divide_attrs " 

466 f"must be a list of attribute names, got {type(attrs).__name__}") 

467 seen = set() 

468 for attr in attrs: 

469 if not isinstance(attr, str) or not attr or not attr.isidentifier(): 

470 raise ValueError( 

471 f"plan_overrides match={self.match!r}: " 

472 f"tp_divide_attrs may only contain valid attribute names, got {attr!r}") 

473 if attr in seen: 

474 raise ValueError( 

475 f"plan_overrides match={self.match!r}: " 

476 f"tp_divide_attrs contains duplicate attribute {attr!r}") 

477 seen.add(attr) 

478 return list(attrs) 

479 

480 def _parse_contract_field(self, attr, raw): 

481 """YAML form → spec field value (DSL parse / sentinel pass-through).""" 

482 # Keep trainer.config outside the distributed component's import-time dependency graph. 

483 # pylint: disable-next=import-outside-toplevel 

484 from hyper_parallel.distributed.recipe_spec import ( 

485 parse_named_placement, 

486 ) 

487 

488 if isinstance(raw, str): 

489 if raw in ("auto", "none"): 

490 return raw 

491 raise ValueError( 

492 f"plan_overrides match={self.match!r}: string value of contract " 

493 f"field {attr} only accepts the sentinels 'auto' (explicit " 

494 f"inherit) / 'none' (explicit clear), got {raw!r}") 

495 if not isinstance(raw, dict): 

496 raise ValueError( 

497 f"plan_overrides match={self.match!r}: contract field {attr} " 

498 f"expects a mapping or a sentinel string, got {raw!r}") 

499 if not raw: 

500 return {} # explicit empty ("as written": clear / no sharding), unlike omitting the field 

501 if attr in ("out_src", "out_dst") and all( 

502 isinstance(v, str) for v in raw.values()): 

503 # Scalar shorthand {axis: placement} — already a NamedPlacement 

504 # (the output name is filled in uniformly by the planner's 

505 # _normalize_out_fields) 

506 return parse_named_placement( 

507 raw, path=f"plan_overrides[{self.match!r}].{attr}") 

508 return { 

509 name: parse_named_placement( 

510 named, path=f"plan_overrides[{self.match!r}].{attr}.{name}") 

511 for name, named in raw.items() 

512 } 

513 

514 

515_WHEN_CONDITIONS = ("cp", "ep", "sequence_parallel", "low_precision") 

516 

517 

518def _import_module_type(path: str) -> type[nn.Module]: 

519 """Import the source module type named by a replacement YAML entry.""" 

520 

521 if not isinstance(path, str) or not path.strip(): 

522 raise ValueError("plan_overrides.replace.source_type must be a non-empty symbol path") 

523 parts = path.split(".") 

524 for split_at in range(len(parts), 0, -1): 

525 module_name = ".".join(parts[:split_at]) 

526 try: 

527 symbol = importlib.import_module(module_name) 

528 except ModuleNotFoundError as exc: 

529 if exc.name == module_name or module_name.startswith(f"{exc.name}."): 

530 continue 

531 raise ValueError( 

532 f"plan_overrides.replace.source_type {path!r} failed while importing {exc.name!r}" 

533 ) from exc 

534 except ImportError as exc: 

535 raise ValueError( 

536 f"plan_overrides.replace.source_type {path!r} failed while importing: {exc}" 

537 ) from exc 

538 for attribute in parts[split_at:]: 

539 if not hasattr(symbol, attribute): 

540 raise ValueError( 

541 f"plan_overrides.replace.source_type {path!r} has no attribute {attribute!r}" 

542 ) 

543 symbol = getattr(symbol, attribute) 

544 if not isinstance(symbol, type) or not issubclass(symbol, nn.Module): 

545 raise TypeError(f"plan_overrides.replace.source_type {path!r} must name an nn.Module type") 

546 return symbol 

547 raise ValueError(f"plan_overrides.replace.source_type {path!r} could not be imported") 

548 

549 

550def _target_replacement_factory( 

551 target: Target[Any], 

552 low_precision_policy: Optional[LowPrecisionDtypeScheme] = None, 

553) -> ModuleReplacementFactory: 

554 """Bind a YAML Target's static args to the replacement factory protocol.""" 

555 

556 if not isinstance(target, ReplacementTarget) and get_module_family(target.callable) is not None: 

557 # Preserve configured Python objects; only Family baselines need selection. 

558 # Target has no public accessors for its raw constructor values. 

559 target = ReplacementTarget( 

560 target.callable, target_path=target._target_path, **target._kwargs, 

561 ) 

562 

563 if isinstance(target, ReplacementTarget) and target.low_precision_dtype_scheme is not None: 

564 target = target.bind_dtype_scheme(low_precision_policy) 

565 

566 if not getattr(target.callable, "_hp_module_replacement", False): 

567 raise TypeError( 

568 "plan_overrides replace.implementation for kind='module' must use @module_replacement" 

569 ) 

570 

571 @module_replacement 

572 def factory( 

573 *, 

574 module: nn.Module, 

575 module_fqn: str, 

576 context: Mapping[str, Any], 

577 ) -> nn.Module: 

578 """Build one replacement module by delegating to the bound YAML Target.""" 

579 if low_precision_policy is not None: 

580 factory_context = dict(context) 

581 factory_context["low_precision"] = low_precision_policy 

582 context = MappingProxyType(factory_context) 

583 return target.build(module=module, module_fqn=module_fqn, context=context) 

584 

585 return factory 

586 

587 

588def _target_method_factory(target: Target[Any]): 

589 """Bind YAML static arguments to the method replacement factory protocol.""" 

590 if not getattr(target.callable, "_hp_method_replacement", False): 

591 raise TypeError("plan_overrides replace.implementation for kind='method' must use @method_replacement") 

592 

593 @method_replacement 

594 def factory( 

595 *, module: nn.Module, method_fqn: str, original_method: Any, context: Mapping[str, Any], 

596 ) -> Any: 

597 """Build one replacement callable from the configured target.""" 

598 return target.build( 

599 module=module, method_fqn=method_fqn, original_method=original_method, context=context, 

600 ) 

601 

602 return factory 

603 

604 

605def entries_to_method_replacements( 

606 entries: List[PlanOverride], *, cp_size: int = 1, ep_size: int = 1, 

607 sequence_parallel: bool = False, low_precision_enabled: bool = False, 

608) -> tuple[MethodReplacementSpec, ...]: 

609 """Resolve active non-forward method actions into executable rules. 

610 

611 Args: 

612 entries: Trainer plan override declarations. 

613 cp_size: Active context parallel degree. 

614 ep_size: Active expert parallel degree. 

615 sequence_parallel: Whether sequence parallelism is active. 

616 low_precision_enabled: Whether online low precision is enabled. 

617 

618 Returns: 

619 Rules applied after module replacement and before sharding. 

620 """ 

621 rules = [] 

622 for entry in entries: 

623 if entry.replace is None or entry.replace.kind != "method": 

624 continue 

625 if entry.match.rpartition(".")[2] == "forward": 

626 continue 

627 _validate_when(entry) 

628 active = { 

629 None: True, 

630 "cp": cp_size > 1, 

631 "ep": ep_size > 1, 

632 "sequence_parallel": sequence_parallel, 

633 "low_precision": low_precision_enabled, 

634 }[entry.when] 

635 if active: 

636 rules.append(MethodReplacementSpec( 

637 match=entry.match, 

638 factory=_target_method_factory(entry.replace.implementation), 

639 )) 

640 else: 

641 logger.info("plan_overrides: match=%r skipped because condition %r is inactive", 

642 entry.match, entry.when) 

643 return tuple(rules) 

644 

645 

646def entries_to_module_replacement_requests( 

647 entries: List[PlanOverride], *, low_precision_enabled: bool = False, 

648 low_precision_config: Optional[LowPrecisionConfig] = None, 

649) -> tuple[ModuleReplacementRequest, ...]: 

650 """Preserve public family identity and explicit parameters through desugaring. 

651 

652 Args: 

653 entries: Value passed to this operation. 

654 """ 

655 effective = low_precision_config if low_precision_enabled else None 

656 result = [] 

657 for entry_index, entry in enumerate(entries): 

658 if entry.replace is None or entry.replace.kind != "module": 

659 continue 

660 if entry.when == "low_precision" and not low_precision_enabled: 

661 continue 

662 _validate_replacement_entry(entry, effective) 

663 target = entry.replace.implementation 

664 baseline = target.requested_target if isinstance(target, ReplacementTarget) else target.callable 

665 arguments = {name: value for name, value in target._kwargs.items() # pylint: disable=protected-access 

666 if name not in {"module", "module_fqn", "context"}} 

667 if isinstance(target, ReplacementTarget): 

668 arguments.update(target.selectors) 

669 if entry.when == "low_precision": 

670 # Legacy policy-bearing closures remain on the explicit compatibility path. 

671 continue 

672 result.append(ModuleReplacementRequest( 

673 match=(entry.match,) if isinstance(entry.match, str) else tuple(entry.match), 

674 target=baseline, kwargs=arguments, origin="user", 

675 declaration_id=f"plan_override:{entry_index}:replace", 

676 module_type=_import_module_type(entry.replace.source_type) if entry.replace.source_type else None, 

677 exact_type=entry.replace.exact_type, 

678 )) 

679 return tuple(result) 

680 

681 

682def entries_to_module_replacements( 

683 entries: List[PlanOverride], 

684 *, 

685 low_precision_enabled: bool = False, 

686 low_precision_config: Optional[LowPrecisionConfig] = None, 

687) -> tuple[ModuleReplacementSpec, ...]: 

688 """Desugar replacement actions and bind their named dtype schemes. 

689 

690 Args: 

691 entries: Raw planner override entries. 

692 low_precision_enabled: Whether low-precision replacement is active. 

693 low_precision_config: Named dtype schemes for active replacements. 

694 """ 

695 

696 effective_config = low_precision_config if low_precision_enabled else None 

697 rules = [] 

698 for entry_index, entry in enumerate(entries): 

699 if entry.replace is None or entry.replace.kind != "module": 

700 continue 

701 if entry.when == "low_precision" and not low_precision_enabled: 

702 continue 

703 _validate_replacement_entry(entry, effective_config) 

704 patterns = (entry.match,) if isinstance(entry.match, str) else tuple(entry.match) 

705 policy = None 

706 scheme_name = getattr(entry.replace.implementation, "low_precision_dtype_scheme", None) 

707 if scheme_name is not None: 

708 if entry.low_precision_dtype_scheme is not None: 

709 raise ValueError( 

710 f"plan_overrides match={entry.match!r}: specify dtype scheme only inside replace.implementation") 

711 if low_precision_config is None: 

712 raise ValueError(f"plan_overrides match={entry.match!r}: dtype scheme requires training.low_precision") 

713 policy = low_precision_config.resolve_dtype_scheme(scheme_name) 

714 elif entry.when == "low_precision" and effective_config is not None: 

715 policy = effective_config.resolve_dtype_scheme(entry.low_precision_dtype_scheme) 

716 serialized_target = entry.replace.implementation.to_dict() 

717 rules.append( 

718 ModuleReplacementSpec( 

719 match=patterns, 

720 factory=_target_replacement_factory(entry.replace.implementation, policy), 

721 module_type=( 

722 _import_module_type(entry.replace.source_type) 

723 if entry.replace.source_type is not None else None 

724 ), 

725 exact_type=entry.replace.exact_type, 

726 declaration_id=f"plan_override:{entry_index}:replace", 

727 replacement_target=serialized_target["_target_"], 

728 static_arguments={ 

729 key: value 

730 for key, value in serialized_target.items() 

731 if key != "_target_" 

732 }, 

733 ) 

734 ) 

735 return tuple(rules) 

736 

737 

738def _validate_replacement_entry(entry: PlanOverride, effective_config: Optional[LowPrecisionConfig]) -> None: 

739 """Validate one replacement before constructing its planner rule.""" 

740 if entry.when not in (None, "low_precision"): 

741 raise ValueError( 

742 f"plan_overrides match={entry.match!r} uses kind='module' and " 

743 "does not support 'when' values other than 'low_precision' " 

744 "(module replacements are not conditioned on parallel topology)" 

745 ) 

746 if entry.low_precision_dtype_scheme is not None and ( 

747 not isinstance(entry.low_precision_dtype_scheme, str) 

748 or not entry.low_precision_dtype_scheme.strip() 

749 ): 

750 raise ValueError( 

751 f"plan_overrides match={entry.match!r}: " 

752 "low_precision_dtype_scheme must be a non-empty string" 

753 ) 

754 if entry.low_precision_dtype_scheme is not None and entry.when != "low_precision": 

755 raise ValueError( 

756 f"plan_overrides match={entry.match!r} sets " 

757 "low_precision_dtype_scheme but is not a low-precision entry" 

758 ) 

759 if entry.when == "low_precision" and effective_config is None: 

760 raise ValueError( 

761 f"plan_overrides match={entry.match!r} enables a low-precision " 

762 "replacement but training.low_precision is missing" 

763 ) 

764 if not isinstance(entry.replace.implementation, Target): 

765 raise TypeError( 

766 f"plan_overrides match={entry.match!r} replace.implementation must be a YAML Target" 

767 ) 

768 

769 

770def _has_sharding_action(entry: PlanOverride) -> bool: 

771 action_names = ( 

772 "region_dispatch", "params", "in_src", "in_dst", "out_src", "out_dst", 

773 ) 

774 return any(getattr(entry, name) is not None for name in action_names) 

775 

776 

777def _validate_when(entry: PlanOverride) -> None: 

778 """Validate programmatic when values that bypass the YAML Literal check.""" 

779 

780 if entry.when is not None and entry.when not in _WHEN_CONDITIONS: 

781 raise ValueError( 

782 f"plan_overrides match={entry.match!r} has invalid when={entry.when!r}; " 

783 f"expected one of {list(_WHEN_CONDITIONS)}" 

784 ) 

785 

786 

787def _merge_plan_override(overrides: dict[str, Any], match: str, spec: Any) -> None: 

788 """Merge one desugared spec into an existing match, with later fields winning.""" 

789 if match not in overrides: 

790 overrides[match] = spec 

791 return 

792 previous = overrides[match] 

793 for name in ("local_compute_fn", "inner_target", "inner_wrapper", 

794 "inner_out_src", "params", "in_src", "in_dst", 

795 "out_src", "out_dst", "tp_divide_attrs"): 

796 value = getattr(spec, name) 

797 if value is not None: 

798 setattr(previous, name, value) 

799 if spec.region_dispatch is not None: 

800 previous.region_dispatch = spec.region_dispatch 

801 

802 

803def _validate_override_transport(entry): 

804 """Validate YAML module-selection transport separately from placement lowering.""" 

805 if isinstance(entry.match, list) and ( 

806 entry.replace is None or entry.replace.kind != "module" or _has_sharding_action(entry) 

807 ): 

808 raise ValueError( 

809 "plan_overrides match lists are supported only for replace.kind='module' " 

810 "entries; use separate string matches for sharding actions" 

811 ) 

812 if entry.replace is not None and entry.replace.kind == "module" and entry.when not in ( 

813 None, 

814 "low_precision", 

815 ): 

816 raise ValueError( 

817 f"plan_overrides match={entry.match!r} uses replace.kind='module' and " 

818 "does not support 'when' values other than 'low_precision' " 

819 "(module replacements are not conditioned on parallel topology)" 

820 ) 

821 

822 

823def entries_to_plan_overrides( 

824 entries: "List[PlanOverride]", *, cp_size: int = 1, ep_size: int = 1,