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

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

7# http://www.apache.org/licenses/LICENSE-2.0

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.

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

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

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

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 ModuleReplacementFactory,

39 ModuleReplacementSpec,

40 module_replacement,

41)

43from hyper_parallel.trainer.config.target import Target

45logger = logging.getLogger(__name__)

48@dataclass

49class AcceleratorConfig:

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

52 tp_size: int = 1

53 cp_size: int = 1

54 ep_size: int = 1

55 pp_size: int = 1

56 sequence_parallel: bool = False

57 loss_parallel: bool = False

58 encoder_parallel: EncoderParallelConfig = field(default_factory=EncoderParallelConfig)

61@dataclass

62class ActivationCheckpointSelection:

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

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

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

67 zero-based entries from that same order.

69 ``model_adapter_safe_regions`` wraps the outermost module matching the

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

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

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

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

74 """

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

77 layer_count: Optional[int] = None

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

80 def __post_init__(self) -> None:

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

82 self._validate_selector_choice()

83 self._validate_selector_values()

85 def _validate_selector_choice(self) -> None:

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

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

88 raise ValueError(

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

90 "model_adapter_safe_regions"

91 )

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

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

94 raise ValueError(

95 "activation_checkpoint.selection.layer_count and layer_indices "

96 "are only valid when "

97 "selection.source=model_adapter_safe_regions"

98 )

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

100 raise ValueError(

101 "activation_checkpoint.selection requires exactly one of "

102 "layer_count or layer_indices when "

103 "selection.source=model_adapter_safe_regions"

104 )

105

106 def _validate_selector_values(self) -> None:

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

108 if self.layer_count is not None and (

109 isinstance(self.layer_count, bool)

110 or not isinstance(self.layer_count, int)

111 or self.layer_count < 0

112 ):

113 raise ValueError(

114 "activation_checkpoint.selection.layer_count must be a "

115 "non-negative integer"

116 )

117 if self.layer_indices is None:

118 return

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

120 raise ValueError(

121 "activation_checkpoint.selection.layer_indices must be a "

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

123 )

124 if any(

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

126 for index in self.layer_indices

127 ):

128 raise ValueError(

129 "activation_checkpoint.selection.layer_indices must contain "

130 "only non-negative integers"

131 )

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

133 raise ValueError(

134 "activation_checkpoint.selection.layer_indices must not "

135 "contain duplicates"

136 )

137

138

139@dataclass

140class ActivationCheckpointConfig:

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

142

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

144 ``"selective"``.

145 """

146

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

148 swap_inputs: bool = False

149 selection: ActivationCheckpointSelection = field(

150 default_factory=ActivationCheckpointSelection

151 )

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

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

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

155 # CSA indexer's masked_fill_ family).

156 allow_cache_entry_mutation: bool = False

157

158 def __post_init__(self) -> None:

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

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

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

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

163 raise TypeError(

164 "activation_checkpoint.allow_cache_entry_mutation must be a bool")

165

166

167@dataclass

168class PlanOverride:

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

170

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

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

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

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

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

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

177 strings. The desugar happens trainer-side in

178 ``instantiate_infrastructure`` BEFORE planner construction.

179

180 Fields:

181 module_type: importable source ``nn.Module`` type required by

182 ``replace_module``. The symbol is resolved but never constructed.

183 exact_type: replacement type check mode. Defaults to subclass matching.

184 replace_module: Target factory decorated with ``@module_replacement``.

185 It receives ``module``, ``module_fqn``, and a read-only context,

186 and must return a structure-preserving replacement. A list

187 ``match`` is supported only for replacement-only entries.

188 match: fqn or fqn glob. Globs first merge matching planner boundaries;

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

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

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

192 planner boundary, with a suggested contract skeleton.

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

194 ``model.layers.0.self_attn``.

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

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

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

198 online low precision is enabled). A replacement action accepts

199 only ``"low_precision"``; module replacement must not depend on

200 the parallel topology.

201 low_precision_dtype_scheme: a named ``training.low_precision.

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

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

204 is injected into the replacement factory context so the strategy

205 factory never resolves a catalog at replacement time.

206 local_compute_fn: a Target whose callable is a **factory** — must be

207 decorated ``@local_compute`` (injection discipline); the mesh

208 family ``mesh``/``tp_mesh``/``cp_mesh``/``ep_mesh`` is mandatory

209 context, all filled by the framework at apply time (``module``/

210 ``spec`` optional); the factory must return the region compute

211 fn ``fn(module, *local_args)`` (params validated against the

212 module's forward at apply time).

213 This is also the **performance-replacement** channel: point it at

214 any factory returning a faster compute fn (see

215 examples/distributed/perf_replacement.py). Shipped EP archetype

216 factories (complete semantics, explicitly chosen — see the

217 archetype table in recipes.py):

218 ``hyper_parallel.distributed.expert_parallel.recipes.qwen2moe_ep_compute_fn``

219 / ``mixtral_ep_compute_fn`` /

220 ``routed_only_ep_compute_fn`` / ``deepseekv3_ep_compute_fn``; the

221 Qwen3-MoE archetype lives at its model-adapter home

222 ``hyper_parallel.models.qwen3_moe.adapter.distributed.expert_parallel.qwen3moe_ep_compute_fn``

223 (M3; the legacy ``recipes.qwen3moe_ep_compute_fn`` delegate was

224 removed in M5). A

225 MoE whose behavior matches no archetype writes its own factory,

226 referencing ``MOE_ROUTER_ADAPTERS`` entries by name (reference:

227 examples/distributed/ep_factories.py).

228 inner_target: attribute name of the inner submodule whose forward is

229 wrapped (``"self"`` = the boundary module itself; default:

230 auto-location).

231 inner_wrapper: a ``INNER_WRAPPER_REGISTRY`` name (``"sdpa_qkv"`` /

232 ``"sdpa_hf"`` / ``"flex_qkv"`` / ``"flex_hf"``) or a Target

233 pointing at an ``@inner_wrapper``-decorated wrapper fn

234 ``fn(target_module, mesh, tp_mesh, cp_mesh, ep_mesh)``

235 — shipped defaults live in

236 ``hyper_parallel.distributed.context_parallel.wrappers``. The

237 mechanism is not CP-gated (declaration == application), but the

238 four shipped names are CP schemes and require an active cp axis

239 (fail-fast otherwise); custom wrappers receive ``cp_mesh=None``

240 when no cp axis exists.

241 inner_out_src: required when inner_target points at a submodule

242 (not ``self``) — the inner output placement declaration:

243 the sentinel ``"first_input"`` (output layout == first DTensor

244 input's layout, for attention-style layout-preserving

245 wrappers), ``{axis: placement_str}`` (single output), or

246 ``{name: {axis: placement_str}}`` (multi-output). The framework

247 never derives inner output layouts — undeclared fails fast at

248 apply time.

249 region_dispatch: validate execution-mode declaration (no default —

250 required when declaring an injection). ``false`` = the region

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

252 cannot be dispatched: it contains communication primitives /

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

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

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

256 injected code is pure standard operators (fused kernels /

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

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

259 not needed without an injection (ordinary boundaries pass through

260 naturally).

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

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

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

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

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

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

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

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

269 be written explicitly.

270 tp_divide_attrs: optional module-instance integer attributes divided

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

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

273 list clears an inherited glob declaration.

274 """

275

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

277 when: Optional[Literal[

278 "cp", "ep", "sequence_parallel", "low_precision"

279 ]] = None

280 low_precision_dtype_scheme: Optional[str] = None

281 module_type: Optional[str] = None

282 exact_type: bool = False

283 replace_module: Optional[Target[Any]] = None

284 local_compute_fn: Optional[Target[Any]] = None

285 inner_target: Optional[str] = None

286 inner_wrapper: Optional[Union[str, Target[Any]]] = None

287 inner_out_src: Optional[Any] = None

288 region_dispatch: Optional[bool] = None

289 params: Optional[Any] = None

290 in_src: Optional[Any] = None

291 in_dst: Optional[Any] = None

292 out_src: Optional[Any] = None

293 out_dst: Optional[Any] = None

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

295

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

297 # insert mode rejects them)

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

299

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

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

302

303 Placement DSL strings become real Placement objects here; the

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

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

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

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

308 test_s5_zero_dep_lint).

309 """

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

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

312 from hyper_parallel.distributed.recipe_spec import (

313 ModuleShardingSpec,

314 )

315

316 if not self.match:

317 raise ValueError(

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

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

320 spec = ModuleShardingSpec(

321 local_compute_fn=self.local_compute_fn,

322 inner_target=self.inner_target,

323 inner_wrapper=self.inner_wrapper,

324 inner_out_src=self._parse_inner_out_src(),

325 region_dispatch=self.region_dispatch,

326 tp_divide_attrs=self._validate_tp_divide_attrs(),

327 )

328 for attr in self._CONTRACT_FIELDS:

329 raw = getattr(self, attr)

330 if raw is None:

331 continue

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

333 return self.match, spec

334

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

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

337 attrs = self.tp_divide_attrs

338 if attrs is None:

339 return None

340 if not isinstance(attrs, list):

341 raise ValueError(

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

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

344 seen = set()

345 for attr in attrs:

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

347 raise ValueError(

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

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

350 if attr in seen:

351 raise ValueError(

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

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

354 seen.add(attr)

355 return list(attrs)

356

357 def _parse_inner_out_src(self):

358 """Desugar the YAML form of inner_out_src: sentinel / single-output DSL / multi-output DSL."""

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

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

361 from hyper_parallel.distributed.recipe_spec import (

362 parse_named_placement,

363 )

364

365 raw = self.inner_out_src

366 if raw is None:

367 return None

368 path = f"plan_overrides[{self.match!r}].inner_out_src"

369 if isinstance(raw, str):

370 if raw == "first_input":

371 return raw

372 raise ValueError(

373 f"{path}: string value only accepts the sentinel 'first_input' "

374 f"(output layout == layout of the first DTensor input), got {raw!r}")

375 if not isinstance(raw, dict) or not raw:

376 raise ValueError(

377 f"{path}: expected 'first_input', {{axis: placement}} or "

378 f"{{name: {{axis: placement}}}}, got {raw!r}")

379 if all(isinstance(v, str) for v in raw.values()):

380 return parse_named_placement(raw, path=path) # single output

381 return {name: parse_named_placement(named, path=f"{path}.{name}")

382 for name, named in raw.items()} # multi-output

383

384 def _parse_contract_field(self, attr, raw):

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

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

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

388 from hyper_parallel.distributed.recipe_spec import (

389 parse_named_placement,

390 )

391

392 if isinstance(raw, str):

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

394 return raw

395 raise ValueError(

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

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

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

399 if not isinstance(raw, dict):

400 raise ValueError(

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

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

403 if not raw:

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

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

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

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

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

409 # _normalize_out_fields)

410 return parse_named_placement(

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

412 return {

413 name: parse_named_placement(

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

415 for name, named in raw.items()

416 }

417

418

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

420

421

422def _import_module_type(path: str) -> type:

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

424

425 if not isinstance(path, str) or not path:

426 raise ValueError("plan_overrides.module_type must be a non-empty symbol path")

427 parts = path.split(".")

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

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

430 try:

431 symbol = importlib.import_module(module_name)

432 except ModuleNotFoundError as exc:

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

434 continue

435 raise ValueError(

436 f"plan_overrides.module_type {path!r} failed while importing {exc.name!r}"

437 ) from exc

438 except ImportError as exc:

439 raise ValueError(

440 f"plan_overrides.module_type {path!r} failed while importing: {exc}"

441 ) from exc

442 for attribute in parts[split_at:]:

443 if not hasattr(symbol, attribute):

444 raise ValueError(

445 f"plan_overrides.module_type {path!r} has no attribute {attribute!r}"

446 )

447 symbol = getattr(symbol, attribute)

448 if not isinstance(symbol, type):

449 raise TypeError(f"plan_overrides.module_type {path!r} must name a type")

450 return symbol

451 raise ValueError(f"plan_overrides.module_type {path!r} could not be imported")

452

453

454def _target_replacement_factory(

455 target: Target[Any],

456 low_precision_policy: Optional[LowPrecisionDtypeScheme] = None,

457) -> ModuleReplacementFactory:

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

459

460 if not getattr(target._target_, "_hp_module_replacement", False): # pylint: disable=protected-access

461 raise TypeError(

462 "plan_overrides replace_module target must be decorated with "

463 "@module_replacement"

464 )

465

466 @module_replacement

467 def factory(

468 *,

469 module: nn.Module,

470 module_fqn: str,

471 context: Mapping[str, Any],

472 ) -> nn.Module:

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

474 if low_precision_policy is not None:

475 factory_context = dict(context)

476 factory_context["low_precision"] = low_precision_policy

477 context = MappingProxyType(factory_context)

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

479

480 return factory

481

482

483def entries_to_module_replacements(

484 entries: List[PlanOverride],

485 *,

486 low_precision_enabled: bool = False,

487 low_precision_config: Optional[LowPrecisionConfig] = None,

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

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

490

491 Args:

492 entries: Raw planner override entries.

493 low_precision_enabled: Whether low-precision replacement is active.

494 low_precision_config: Named dtype schemes for active replacements.

495 """

496

497 effective_config = low_precision_config if low_precision_enabled else None

498 rules = []

499 for entry in entries:

500 if entry.replace_module is None:

501 continue

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

503 continue

504 _validate_replacement_entry(entry, effective_config)

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

506 policy = None

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

508 policy = effective_config.resolve_dtype_scheme(entry.low_precision_dtype_scheme)

509 rules.append(

510 ModuleReplacementSpec(

511 match=patterns,

512 factory=_target_replacement_factory(entry.replace_module, policy),

513 module_type=_import_module_type(entry.module_type),

514 exact_type=entry.exact_type,

515 )

516 )

517 return tuple(rules)

518

519

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

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

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

523 raise ValueError(

524 f"plan_overrides match={entry.match!r} uses replace_module and "

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

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

527 )

528 if entry.low_precision_dtype_scheme is not None and (

529 not isinstance(entry.low_precision_dtype_scheme, str)

530 or not entry.low_precision_dtype_scheme.strip()

531 ):

532 raise ValueError(

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

534 "low_precision_dtype_scheme must be a non-empty string"

535 )

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

537 raise ValueError(

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

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

540 )

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

542 raise ValueError(

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

544 "replacement but training.low_precision is missing"

545 )

546 if entry.module_type is None:

547 raise ValueError(

548 f"plan_overrides match={entry.match!r} uses replace_module but omits module_type"

549 )

550 if not isinstance(entry.replace_module, Target):

551 raise TypeError(

552 f"plan_overrides match={entry.match!r} replace_module must be a YAML Target"

553 )

554

555

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

557 action_names = (

558 "local_compute_fn", "inner_target", "inner_wrapper", "inner_out_src",

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

560 )

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

562

563

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

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

566

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

568 raise ValueError(

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

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

571 )

572

573

574def entries_to_plan_overrides(

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

576 sequence_parallel: bool = False,

577 low_precision_enabled: bool = False,

578) -> "dict[str, Any]":

579 """Desugar PlanOverride entries into a ``plan_overrides`` dict.

580

581 Args:

582 entries: Raw YAML planner override entries.

583 cp_size: Context-parallel degree.

584 ep_size: Expert-parallel degree.

585 sequence_parallel: Whether decoder sequence parallelism is active.

586 low_precision_enabled: Whether low-precision rules are active.

587

588 - ``when`` filter: an entry whose declared condition is inactive

589 (``when="cp"`` with cp_size==1 etc.) is SKIPPED with an INFO log —

590 declared gating, never silent application. An unknown ``when`` value

591 fails fast listing the valid conditions (programmatic construction

592 bypasses the resolver's Literal check);

593 - two entries sharing the same ``match`` are merged field-wise

594 (non-None fields, later entry wins) — the planner's interface is

595 exactly this dict.

596 """

597 overrides: dict[str, Any] = {}

598 for entry in entries:

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

600 entry.replace_module is None or _has_sharding_action(entry)

601 ):

602 raise ValueError(

603 "plan_overrides match lists are supported only for replace_module "

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

605 )

606 if entry.replace_module is not None and entry.when not in (

607 None,

608 "low_precision",

609 ):

610 raise ValueError(

611 f"plan_overrides match={entry.match!r} uses replace_module and "

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

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

614 )

615 _validate_when(entry)

616 if entry.when is not None:

617 active = {

618 "cp": cp_size > 1,

619 "ep": ep_size > 1,

620 "sequence_parallel": sequence_parallel,

621 "low_precision": low_precision_enabled,

622 }[entry.when]

623 if not active:

624 logger.info(

625 "plan_overrides: match=%r skipped because condition %r "

626 "is inactive",

627 entry.match, entry.when)

628 continue

629 if entry.replace_module is not None and not _has_sharding_action(entry):

630 continue

631 match, spec = entry.to_override()

632 if match in overrides:

633 prev = overrides[match]

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

635 "inner_out_src", "params", "in_src", "in_dst",

636 "out_src", "out_dst", "tp_divide_attrs"):

637 value = getattr(spec, name)

638 if value is not None:

639 setattr(prev, name, value)

640 if spec.region_dispatch is not None:

641 prev.region_dispatch = spec.region_dispatch

642 else:

643 overrides[match] = spec

644 return overrides

645

646

647def normalize_distributed_setup_overrides(

648 distributed_setup: Any,

649 config: "TrainerConfig",

650) -> None:

651 """Fill the normalized plan fields of a ``DistributedSetup`` (05 §15.2.6).

652

653 Args:

654 distributed_setup: Mutable setup receiving normalized planner rules.

655 config: Trainer configuration containing raw override entries.

656

657 This is the single Trainer-side desugar point: raw YAML ``PlanOverride``

658 entries (``TrainerConfig.plan_overrides``) are converted HERE —

659 when-filtered against the built mesh and compiled into replacement

660 rules — so the AutoModels build pipeline consumes only normalized

661 objects and never imports trainer config.

662 """

663 entries = list(getattr(config, "plan_overrides", None) or [])

664 mesh_context = getattr(distributed_setup, "mesh_context", None)

665 low_precision_enabled = bool(

666 getattr(

667 getattr(distributed_setup, "low_precision_config", None),

668 "enabled",

669 False,

670 )

671 )

672 if low_precision_enabled and not any(

673 entry.replace_module is not None and entry.when == "low_precision"

674 for entry in entries

675 ):

676 raise ValueError(

677 "low_precision is enabled but no low-precision module replacement "

678 "is configured"

679 )

680 distributed_setup.module_replacements = entries_to_module_replacements(

681 entries,

682 low_precision_enabled=low_precision_enabled,

683 low_precision_config=getattr(distributed_setup, "low_precision_config", None),

684 )

685 distributed_setup.plan_overrides = (

686 entries_to_plan_overrides(

687 entries,

688 cp_size=getattr(mesh_context, "cp_size", 1),

689 ep_size=getattr(mesh_context, "ep_size", 1),

690 sequence_parallel=getattr(mesh_context, "sequence_parallel", False),

691 low_precision_enabled=low_precision_enabled,

692 )

693 if entries

694 else None

695 )

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

234 statements