Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/compile/pass

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"""

16Parallel Configuration - Graph-mode parallel configuration.

18Pure configuration dataclass: no environment probes. Whether FSDP actually

19runs is decided by ``fsdp_enabled`` here (the user's intent) AND the runtime

20distributed guard inside ``FSDPPass`` (which still early-returns on

21``world_size == 1``). The previous ``fsdp_enabled`` property returned

22``dist.is_initialized() and world_size > 1`` — that turned on FSDP whenever

23distributed was initialized, leaving no way to run pure-TP / pure-PP graph

24mode. The explicit field fixes that. Pipeline-parallel (``pp_enabled``)

25follows the same contract: intent here, runtime guard in ``PpPass``.

27Note:

28 TP / SP / LP are deliberately NOT configurable here: their collectives

29 are baked into module boundary forwards by the automodel flow (shared

30 with eager mode), and no graph pass consumes them. ``PassConfig`` only

31 drives what the graph passes actually execute.

32"""

34__all__ = [

35 "PassConfig",

36 "GradientReduction",

37 "PP_SCHEDULES",

38 "INDUCTOR_MODES",

39 "DP_MODES",

40 "AC_MODES",

41 "AC_REGION_SOURCES",

42 "ACTIVATION_SWAP_MODES",

43 "normalize_dp_mode",

44 "build_pass_config_from_trainer_config",

45]

47import logging

48import math

49from dataclasses import dataclass

50from typing import Any, Dict, List, Literal, Optional

52_LOG = logging.getLogger(__name__)

55# Pipeline schedule names accepted by ``pp_schedule``. Kept torch-free here so

56# the config dataclass stays importable anywhere; ``pp_schedule``'s registry

57# (which does import torch) must expose exactly these keys.

58PP_SCHEDULES = ("gpipe", "1f1b", "interleaved_1f1b", "vpp")

60# Backend compilation modes for the terminal ``InductorPass``:

61# "off" - no Inductor compilation (graph runs as plain FX nodes)

62# "regional" - compile tagged compute regions; collectives and

63# side-effectful nodes stay eager (see InductorPass)

64# "full" - compile the whole joint graph as one Inductor region

65INDUCTOR_MODES = ("off", "regional", "full")

67# Schedule names that can drive multiple virtual stages per physical rank.

68_INTERLEAVED_SCHEDULES = ("interleaved_1f1b", "vpp")

70# Data-parallel modes, mirroring simplefsdp's ``data_parallel`` modes:

71# "fsdp" -> fully_shard (Shard(0) on the fsdp axis)

72# "ddp" -> replicate (Replicate; grads all-reduced)

73# "hsdp" -> hybrid_shard (Replicate x Shard(0); rs + all-reduce)

74DP_MODES = ("fsdp", "ddp", "hsdp")

76# simplefsdp / fully_shard spellings accepted as aliases.

77_DP_MODE_ALIASES = {

78 "fully_shard": "fsdp",

79 "replicate": "ddp",

80 "hybrid_shard": "hsdp",

81}

83# Static-graph activation checkpointing modes for ``ActivationCheckpointPass``,

84# mirroring the dynamic ``ActivationCheckpointConfig.mode`` vocabulary:

85# "off" - save every forward activation (no recompute pass)

86# "full" - recompute every checkpoint-scoped module op in backward

87# "selective" - save compute/communication intensive ops, recompute the rest

88AC_MODES = ("off", "full", "selective")

90# Where the recompute-region shape comes from, mirroring

91# ``ActivationCheckpointSelection.source``. ``default`` selects whole blocks

92# (and lets ``full`` auto-upgrade to a model-owned RecomputePolicy);

93# ``model_adapter_safe_regions`` requires an adapter-declared region.

94AC_REGION_SOURCES = ("default", "model_adapter_safe_regions")

96# Dynamic-graph activation-swap modes reused by graph mode. ``attention``

97# offloads saved activations inside ``self_attn`` modules to host memory

98# (the graph pass equivalent of the eager attention swap wrapper).

99ACTIVATION_SWAP_MODES = ("none", "attention")

100

101

102def normalize_dp_mode(mode: str) -> str:

103 """Return the canonical DP mode for ``mode`` (accepts simplefsdp aliases).

104

105 Args:

106 mode: Requested DP mode spelling.

107

108 Returns:

109 The canonical ``DP_MODES`` entry.

110

111 Raises:

112 ValueError: When ``mode`` is not a string or is unknown.

113 """

114 if not isinstance(mode, str):

115 raise ValueError(f"dp_mode must be a string, got {type(mode).__name__}")

116 key = mode.strip().lower()

117 return _DP_MODE_ALIASES.get(key, key)

118

119

120@dataclass(frozen=True)

121class GradientReduction:

122 """Reduction operation and a divisor independent of the communication mesh.

123

124 The divisor is applied before communication, matching eager FSDP's

125 reduction-precision scaling and its singleton-group behavior.

126 """

127

128 reduce_op: Literal["avg", "sum"] = "avg"

129 divisor: int = 1

130

131 def __post_init__(self) -> None:

132 if self.reduce_op not in ("avg", "sum"):

133 raise ValueError("Gradient reduction must be 'avg' or 'sum'")

134 if isinstance(self.divisor, bool) or not isinstance(self.divisor, int) or self.divisor < 1:

135 raise ValueError("Gradient divisor must be a positive integer")

136

137 @property

138 def scale(self) -> float | None:

139 """Return the pre-reduction multiplier, or None for an identity scale."""

140 return 1.0 / self.divisor if self.divisor > 1 else None

141

142

143

144@dataclass

145class PassConfig:

146 """Parallel configuration for graph-mode FSDP (+ optional TP / SP / PP) training.

147

148 Attributes:

149 enable_overlap: Drive ``AutoOverlapPass`` to move ``wait_tensor`` for

150 communication/compute overlap.

151 fsdp_bucketing: Fuse compatible FSDP all-gathers and reduce-scatters before

152 overlap. Python-only, disabled by default; configure before compilation.

153 fsdp_bucket_modules: Optional module FQN groups, produced by

154 ``build_module_bucket_plans`` or supplied as a list of paths/groups.

155 None groups by immediate parameter parents; an empty list selects nothing.

156 All ranks in each FSDP group must use the same configuration.

157 fsdp_allgather_prefetch: Prefetch forward/backward FSDP all-gathers

158 and delay their wait/view nodes independently of enable_overlap.

159 fsdp_overlap_grad_reduce: Advance backward reduce-scatter launches and

160 delay their waits independently of parameter prefetch and enable_overlap.

161 max_in_flight_gb: GiB of early FSDP all-gather replicas before first use,

162 shared across anchors in each phase. Uses PyTorch's overlap scheduler

163 naming and binary units. Zero disables early launches, not wait sinking.

164 fsdp_enabled: Drive ``FSDPPass`` (data-parallel partitioning: parameter

165 all_gather + gradient reduction + live-model sharding, selected by

166 ``dp_mode``). ``False`` skips data parallelism entirely — set this

167 for pure-TP / pure-PP graph-mode runs. ``FSDPPass`` itself still

168 early-returns when distributed is not initialized or

169 ``world_size == 1``, so single-card runs are a no-op regardless.

170 fsdp_reshard_after_forward: Drive the FSDP reshard optimization.

171 When ``True`` (default, matching FSDP's ``reshard_after_forward``)

172 ``FSDPPass`` sinks each param's all_gather to its first forward

173 use, frees the replicated full parameter once its last forward

174 reader is done, and — when the backward still needs it — inserts a

175 fresh all_gather plus a rematerialization of the saved forward

176 views just before the first backward consumer. Peak memory then

177 tracks the forward working set instead of the sum of all

178 replicated parameters. Set ``False`` to keep every gathered

179 parameter resident for the whole joint graph (previous behavior).

180 fsdp_degree: Size of the FSDP group. ``None`` (default) means

181 "resolve at runtime": the trainer back-fills it from the

182 automodel ``MeshContext`` (TP+FSDP hybrid, where the FSDP group

183 is a proper sub-group of the world), and ``FSDPPass`` falls back

184 to ``world_size`` for the FSDP-only path. Mutating this after

185 construction is supported but discouraged — prefer passing the

186 resolved degree at construction time (see ``GraphTrainer``).

187 For ``dp_mode="hsdp"`` this is the *shard* axis degree.

188 dp_mode: Data-parallel mode, one of ``DP_MODES``: ``"fsdp"``

189 (fully_shard, the default), ``"ddp"`` (replicate; no parameter

190 sharding, gradients all-reduced), or ``"hsdp"`` (hybrid_shard;

191 parameters replicated on the ``dp_replicate`` axis and sharded on

192 the ``fsdp`` axis, gradients reduce-scattered then all-reduced).

193 simplefsdp spellings ``"fully_shard"`` / ``"replicate"`` /

194 ``"hybrid_shard"`` are accepted as aliases and normalized.

195 dp_replicate_degree: Size of the replicate axis for ``"ddp"`` /

196 ``"hsdp"``. ``None`` means "resolve at runtime" (from the mesh,

197 or ``world_size // fsdp_degree`` for HSDP / ``world_size`` for

198 DDP). Kept separate from ``fsdp_degree`` (the shard axis) so an

199 HSDP mesh is unambiguous.

200 pp_enabled: Drive ``PpPass`` (pipeline-parallel graph split). When

201 ``True`` the joint fwd+bwd graph is sliced to this rank's stage

202 along module-FQN boundaries declared in ``GraphParallelPlan``, cross-stage

203 activations are exchanged via P2P ``isend``/``irecv``, and a

204 self-contained pipeline schedule is installed as a ``call_module``

205 stub inside the rewritten graph, so the trainer needs no PP

206 wiring. ``PpPass`` early-returns when distributed is not

207 initialized or ``world_size == 1``.

208 pp_degree: Number of pipeline-parallel ranks (physical stages).

209 ``None`` means "resolve at runtime": pure-PP resolves it from

210 ``world_size``; the PP+DP hybrid resolves it from the 2-D

211 ``(pp, DP)`` mesh built by ``GraphCompiler._init_pp_dp_mesh``

212 (which also back-fills ``fsdp_degree`` /

213 ``dp_replicate_degree``), subject to

214 ``pp_degree * dp_degree == world_size``.

215 pp_chunks: Number of virtual pipeline stages per physical rank.

216 ``1`` (default) is the classic single-stage-per-rank layout used

217 by ``"gpipe"`` / ``"1f1b"``. Values ``> 1`` enable VPP

218 (interleaved 1F1B): the model is split into

219 ``pp_degree * pp_chunks`` virtual stages, rank ``r`` owns virtual

220 stages ``r, r + pp_degree, ...``, and the ``"interleaved_1f1b"``

221 schedule pipelines the virtual stages. Requires

222 ``pp_schedule`` to be an interleaved schedule.

223 pp_microbatch_size: Number of samples per microbatch for the

224 pipeline schedule. The per-step batch size (the leading dim of

225 the trainer's input tensor) must be divisible by this. ``1`` means

226 one sample per microbatch (maximum pipeline concurrency, most

227 P2P traffic); larger values trade bubble size for fewer P2P

228 round-trips.

229 pp_schedule: Pipeline schedule name, one of ``PP_SCHEDULES``.

230 ``"gpipe"`` (default) runs all forwards before all backwards

231 (simplest ordering, largest bubble). ``"1f1b"`` interleaves one

232 backward per forward in steady state, shrinking the pipeline

233 bubble and freeing each microbatch's activations as soon as its

234 backward runs (GPipe retains all of them until the sweep ends);

235 it requires ``num_microbatches >= pp_degree``.

236 ``"interleaved_1f1b"`` (alias ``"vpp"``) is the VPP schedule: it

237 supports ``pp_chunks > 1`` virtual stages per rank and uses the

238 classic interleaved-1F1B warmup / steady / cooldown order.

239 ``PpPass`` resolves the name to a schedule class via

240 ``pp_schedule.get_schedule_class``.

241 inductor_mode: Terminal backend compilation for the joint graph, one

242 of ``INDUCTOR_MODES``. ``"off"`` (default) runs the graph as

243 plain FX nodes; ``"regional"`` compiles tagged compute regions

244 with TorchInductor while keeping communication collectives and

245 side-effectful nodes (e.g. the FSDP reshard storage free) eager;

246 ``"full"`` compiles the whole joint graph as one region. Both

247 compiling modes force ``eager_numerics.division_rounding`` on

248 torch builds that provide it (2.11+) so compiled code stays

249 bitwise-equal to eager, and restore the positional calling

250 convention afterwards. Requires a torch with Inductor. With

251 pipeline parallel enabled, ``PpPass`` runs first and this pass

252 compiles each stage's forward/backward subgraph

253 independently (both modes); cross-stage fusion is inherently

254 impossible because the schedule's P2P sits between the stages.

255 inductor_configs: Optional Inductor config overrides passed to the

256 compiled regions (keys validated by ``torch._inductor.config``),

257 e.g. ``{"reorder_for_peak_memory": True}``. ``"full"`` mode

258 always sets ``reorder_for_peak_memory=True`` unless overridden.

259 ac_mode: Static-graph activation checkpointing mode for

260 ``ActivationCheckpointPass``, one of ``AC_MODES``. Mirrors the

261 dynamic ``ActivationCheckpointConfig.mode`` vocabulary and is

262 projected from a trainer config's ``activation_checkpoint``

263 section by ``build_pass_config_from_trainer_config``. ``"off"``

264 (default) saves every forward activation; ``"full"`` recomputes

265 every checkpoint-scoped module op in backward; ``"selective"``

266 saves the shared compute/communication operator set and

267 recomputes the rest.

268 ac_layer_count: Optional layer selector: checkpoint only the first

269 ``ac_layer_count`` discovered transformer blocks (the same

270 HuggingFace-marker discovery the eager path uses). Mutually

271 exclusive with ``ac_layer_indices``.

272 ac_layer_indices: Optional explicit zero-based indices into the

273 discovered transformer blocks; only nodes inside the selected

274 blocks recompute. Mutually exclusive with ``ac_layer_count``.

275 ac_region_source: Where the recompute-region shape comes from, one of

276 ``AC_REGION_SOURCES``. ``"default"`` selects whole blocks (and

277 lets ``full`` auto-upgrade to a model-owned ``RecomputePolicy``);

278 ``"model_adapter_safe_regions"`` requires an adapter-declared

279 region and its ``exclude_patterns``. Projected from the trainer

280 selection's ``source``.

281 ac_save_forward_all_gather: Save forward (TP/SP/CP) all-gathers instead

282 of replaying them during backward recompute. Graph-mode only;

283 defaults to ``True``, so static-graph recompute does not duplicate

284 a forward collective. FSDP parameter all-gathers are inserted after

285 this pass and are unaffected. Set ``False`` to replay the

286 all-gather, trading communication for less retained memory.

287 activation_swap: Reuse the dynamic-graph ``activation_swap`` setting.

288 ``"attention"`` (the only enabled mode) runs ``CpuOffloadPass``:

289 saved activations produced inside ``self_attn`` modules are copied

290 to host memory asynchronously during forward and reloaded before

291 their first backward consumer. ``"none"`` disables offload.

292 offload_budget_gb: Host-memory budget, in GiB per rank, for activations

293 selected by ``CpuOffloadPass``. Tensors are selected largest-first

294 until the budget is exhausted; ``0`` means no budget. Graph-mode

295 only (the eager attention swap has no budget knob). The budget caps

296 offloaded activation bytes, not the pinned host buffer footprint, so

297 size it against the host's free memory to avoid host OOM.

298 offload_prefetch_n_layers: Move each backward reload this many attention

299 regions earlier to overlap H2D copies with preceding backward

300 compute. ``0`` keeps the reload immediately before its consumer.

301 Graph-mode only.

302

303 Note:

304 ``fsdp_enabled`` / ``pp_enabled`` no longer probe

305 ``torch.distributed``. The distributed-initialized check moved into

306 the respective passes' ``run`` (their original location) so this

307 dataclass stays torch-free and importable anywhere.

308 """

309

310 enable_overlap: bool = True

311 fsdp_enabled: bool = True

312 fsdp_reshard_after_forward: bool = True

313 fsdp_degree: Optional[int] = None

314 dp_mode: str = "fsdp"

315 dp_replicate_degree: Optional[int] = None

316 pp_enabled: bool = False

317 pp_degree: Optional[int] = None

318 pp_chunks: int = 1

319 pp_microbatch_size: int = 1

320 pp_schedule: str = "gpipe"

321 fsdp_bucketing: bool = False

322 fsdp_bucket_modules: Optional[

323 list[str | list[str]] | tuple[tuple[str, ...], ...]

324 ] = None

325 fsdp_allgather_prefetch: bool = True

326 fsdp_overlap_grad_reduce: bool = True

327 max_in_flight_gb: float = 0.25

328 inductor_mode: str = "off"

329 inductor_configs: Optional[Dict[str, Any]] = None

330 ac_mode: str = "off"

331 ac_layer_count: Optional[int] = None

332 ac_layer_indices: Optional[List[int]] = None

333 ac_region_source: str = "default"

334 ac_save_forward_all_gather: bool = True

335 activation_swap: str = "none"

336 offload_budget_gb: float = 64.0

337 offload_prefetch_n_layers: int = 1

338

339 def __post_init__(self) -> None:

340 self.dp_mode = normalize_dp_mode(self.dp_mode)

341 self.validate()

342

343 def validate(self) -> None:

344 """Sanity-check invariants; also re-run after manual field mutation.

345

346 Raises:

347 ValueError: On a non-positive explicit ``fsdp_degree`` /

348 ``pp_degree``, a non-positive ``pp_microbatch_size``, an

349 unknown ``pp_schedule`` / ``dp_mode``, or a non-positive

350 ``dp_replicate_degree``, or ``pp_chunks > 1`` with a

351 non-interleaved schedule, or an invalid activation

352 checkpointing mode / layer selector.

353 """

354 self._validate_flags()

355 self._validate_dp()

356 self._validate_pp()

357 self._validate_ac()

358 self._validate_offload()

359

360 def _validate_ac(self) -> None:

361 """Validate the activation checkpointing mode and layer selectors."""

362 if not isinstance(self.ac_mode, str) or self.ac_mode not in AC_MODES:

363 raise ValueError(f"ac_mode must be one of {AC_MODES}, got {self.ac_mode!r}")

364 if (

365 not isinstance(self.ac_region_source, str)

366 or self.ac_region_source not in AC_REGION_SOURCES

367 ):

368 raise ValueError(

369 f"ac_region_source must be one of {AC_REGION_SOURCES}, "

370 f"got {self.ac_region_source!r}"

371 )

372 if not isinstance(self.ac_save_forward_all_gather, bool):

373 raise ValueError("ac_save_forward_all_gather must be a bool")

374 self._validate_ac_layer_count()

375 self._validate_ac_layer_indices()

376 if self.ac_layer_count is not None and self.ac_layer_indices is not None:

377 raise ValueError(

378 "ac_layer_count and ac_layer_indices are mutually exclusive; "

379 "set exactly one"

380 )

381 if (

382 self.ac_mode != "off"

383 and self.ac_region_source == "model_adapter_safe_regions"

384 and (self.ac_layer_count is None and self.ac_layer_indices is None)

385 ):

386 raise ValueError(

387 "ac_region_source=model_adapter_safe_regions requires exactly "

388 "one of ac_layer_count or ac_layer_indices"

389 )

390

391 def _validate_ac_layer_count(self) -> None:

392 """Validate the optional prefix layer selector."""

393 if self.ac_layer_count is not None and (

394 isinstance(self.ac_layer_count, bool)

395 or not isinstance(self.ac_layer_count, int)

396 or self.ac_layer_count < 0

397 ):

398 raise ValueError(

399 "ac_layer_count must be None or a non-negative integer, "

400 f"got {self.ac_layer_count!r}"

401 )

402

403 def _validate_ac_layer_indices(self) -> None:

404 """Validate the optional explicit layer index selector."""

405 indices = self.ac_layer_indices

406 if indices is None:

407 return

408 invalid = (

409 not isinstance(indices, list)

410 or not indices

411 or any(

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

413 for index in indices

414 )

415 )

416 if invalid:

417 raise ValueError(

418 "ac_layer_indices must be None or a non-empty list of "

419 f"non-negative integers, got {indices!r}"

420 )

421 if len(set(indices)) != len(indices):

422 raise ValueError(

423 f"ac_layer_indices must not contain duplicates, got {indices!r}"

424 )

425

426 def _validate_offload(self) -> None:

427 """Validate graph-mode activation offload (dynamic attention-swap reuse)."""

428 if (

429 not isinstance(self.activation_swap, str)

430 or self.activation_swap not in ACTIVATION_SWAP_MODES

431 ):

432 raise ValueError(

433 "activation_swap must be one of "

434 f"{ACTIVATION_SWAP_MODES}, got {self.activation_swap!r}"

435 )

436 if self.activation_swap == "none":

437 return

438 if self.ac_mode != "off":

439 raise ValueError(

440 "activation_swap='attention' is incompatible with graph-mode "

441 "activation checkpointing; the dynamic path enforces the same "

442 "combination rule"

443 )

444 if self.pp_enabled:

445 raise ValueError(

446 "activation_swap='attention' does not support pipeline "

447 "parallelism; the dynamic path enforces the same rule"

448 )

449 value = self.offload_budget_gb

450 if (

451 isinstance(value, bool)

452 or not isinstance(value, (int, float))

453 or not math.isfinite(value)

454 or value < 0

455 ):

456 raise ValueError("offload_budget_gb must be a finite non-negative number")

457 prefetch = self.offload_prefetch_n_layers

458 if isinstance(prefetch, bool) or not isinstance(prefetch, int) or prefetch < 0:

459 raise ValueError("offload_prefetch_n_layers must be a non-negative integer")

460

461 def _validate_flags(self) -> None:

462 """Validate boolean flags and the overlap memory budget."""

463 for name in (

464 "fsdp_bucketing",

465 "fsdp_allgather_prefetch",

466 "fsdp_overlap_grad_reduce",

467 ):

468 value = getattr(self, name)

469 if not isinstance(value, bool):

470 raise ValueError(f"{name} must be a bool")

471 if self.fsdp_bucket_modules is not None and not isinstance(

472 self.fsdp_bucket_modules, (list, tuple)

473 ):

474 raise ValueError(

475 "fsdp_bucket_modules must be None or a list/tuple of module groups"

476 )

477 value = self.max_in_flight_gb

478 if (

479 isinstance(value, bool)

480 or not isinstance(value, (int, float))

481 or not math.isfinite(value)

482 or value < 0

483 ):

484 raise ValueError("max_in_flight_gb must be a finite non-negative number")

485 self._validate_inductor()

486

487 def _validate_inductor(self) -> None:

488 """Validate the Inductor compilation mode and config overrides."""

489 if not isinstance(self.inductor_mode, str) or (

490 self.inductor_mode.lower() not in INDUCTOR_MODES

491 ):

492 raise ValueError(

493 f"inductor_mode must be one of {INDUCTOR_MODES}, "

494 f"got {self.inductor_mode!r}"

495 )

496 # Canonicalize to one spelling so the pass's mode dispatch cannot be

497 # defeated by an accepted-but-unnormalized value like "FULL".

498 self.inductor_mode = self.inductor_mode.lower()

499 if self.inductor_configs is not None and (

500 not isinstance(self.inductor_configs, dict)

501 or not all(isinstance(key, str) for key in self.inductor_configs)

502 ):

503 raise ValueError(

504 "inductor_configs must be None or a dict of "

505 "torch._inductor.config option names (str) to values"

506 )

507

508 def _validate_dp(self) -> None:

509 """Validate the tensor-parallel and data-parallel degrees."""

510 if self.fsdp_degree is not None and self.fsdp_degree < 1:

511 raise ValueError(

512 f"fsdp_degree must be None or a positive int, got {self.fsdp_degree}"

513 )

514 if self.dp_mode not in DP_MODES:

515 raise ValueError(

516 f"dp_mode must be one of {DP_MODES} "

517 f"(aliases: {tuple(_DP_MODE_ALIASES)}), got {self.dp_mode!r}"

518 )

519 if self.dp_replicate_degree is not None and self.dp_replicate_degree < 1:

520 raise ValueError(

521 f"dp_replicate_degree must be None or a positive int, "

522 f"got {self.dp_replicate_degree}"

523 )

524

525 def _validate_pp(self) -> None:

526 """Validate the pipeline-parallel degree, microbatch size, and schedule."""

527 if self.pp_degree is not None and self.pp_degree < 1:

528 raise ValueError(

529 f"pp_degree must be None or a positive int, got {self.pp_degree}"

530 )

531 if self.pp_chunks < 1:

532 raise ValueError(f"pp_chunks must be >= 1, got {self.pp_chunks}")

533 if self.pp_microbatch_size < 1:

534 raise ValueError(

535 f"pp_microbatch_size must be >= 1, got {self.pp_microbatch_size}"

536 )

537 if not isinstance(self.pp_schedule, str) or (

538 self.pp_schedule.lower() not in PP_SCHEDULES

539 ):

540 raise ValueError(

541 f"pp_schedule must be one of {PP_SCHEDULES}, got {self.pp_schedule!r}"

542 )

543 if (

544 self.pp_chunks > 1

545 and self.pp_schedule.lower() not in _INTERLEAVED_SCHEDULES

546 ):

547 raise ValueError(

548 f"pp_chunks={self.pp_chunks} requires an interleaved schedule "

549 f"(one of {_INTERLEAVED_SCHEDULES}), got {self.pp_schedule!r}"

550 )

551

552

553def build_pass_config_from_trainer_config(

554 config: Any,

555 *,

556 fsdp_enabled: Optional[bool] = None,

557 enable_overlap: bool = False,

558) -> PassConfig:

559 """Project trainer topology intent onto graph-mode pass config.

560

561 This helper stays in the pass-config module because it computes the

562 graph-pass surface from the trainer topology without altering runtime

563 objects. Accepting ``Any`` keeps this module importable without the

564 trainer package. Only fields the graph passes actually consume are

565 projected: TP / SP / LP topology (``accelerator.tp_size`` /

566 ``sequence_parallel`` / ``loss_parallel``) is owned by the automodel

567 flow and deliberately not copied.

568

569 The dynamic ``activation_checkpoint`` section is projected onto the

570 static-graph AC fields: ``mode`` becomes ``ac_mode`` and a

571 ``model_adapter_safe_regions`` selection contributes its layer

572 selectors. ``swap_inputs`` and ``allow_cache_entry_mutation`` have no

573 static-graph counterpart (no eager checkpoint cache) and are ignored

574 with a warning when set, mirroring the dynamic compile path. The

575 top-level dynamic ``activation_swap`` field is projected onto

576 ``activation_swap`` so graph mode reuses the eager attention-swap

577 configuration (the graph pass owns offload instead of the eager wrapper).

578 The graph-only ``compile.offload_budget_gb`` and

579 ``compile.offload_prefetch_n_layers`` knobs are projected too, so Trainer

580 YAML can tune the offload budget instead of being pinned to the 64 GiB

581 default.

582

583 Args:

584 config: Trainer topology (e.g. a ``TrainerConfig``); only the

585 ``fsdp_config``, ``compile`` and ``activation_checkpoint``

586 surfaces are read.

587 fsdp_enabled: Overrides the sharded-DP intent derived from

588 ``config.fsdp_config`` when provided.

589 enable_overlap: Value for the legacy ``AutoOverlapPass`` switch.

590

591 Returns:

592 The graph-mode ``PassConfig`` for this trainer topology.

593 """

594 compile_config = getattr(config, "compile", None)

595 if fsdp_enabled is None:

596 fsdp_enabled = (

597 config.fsdp_config.dp_shard_size > 1

598 or config.fsdp_config.edp_shard_size > 1

599 )

600 fsdp_degree = config.fsdp_config.dp_shard_size if fsdp_enabled else None

601

602 ac_mode = "off"

603 ac_layer_count = None

604 ac_layer_indices = None

605 ac_region_source = "default"

606 ac_config = getattr(config, "activation_checkpoint", None)

607 if ac_config is not None and getattr(ac_config, "mode", None) not in (None, "off"):

608 ac_mode = ac_config.mode

609 selection = getattr(ac_config, "selection", None)

610 ac_region_source = getattr(selection, "source", "default")

611 if ac_region_source == "model_adapter_safe_regions":

612 ac_layer_count = getattr(selection, "layer_count", None)

613 ac_layer_indices = getattr(selection, "layer_indices", None)

614 if getattr(ac_config, "swap_inputs", False):

615 _LOG.warning(

616 "activation_checkpoint.swap_inputs is not supported in graph "

617 "mode; input swapping will be disabled."

618 )

619 if getattr(ac_config, "allow_cache_entry_mutation", False):

620 _LOG.info(

621 "activation_checkpoint.allow_cache_entry_mutation has no "

622 "effect in graph mode; ignoring it."

623 )

624

625 return PassConfig(

626 enable_overlap=enable_overlap,

627 fsdp_enabled=fsdp_enabled,

628 fsdp_degree=fsdp_degree,

629 fsdp_reshard_after_forward=getattr(

630 config.fsdp_config, "reshard_after_forward", True

631 ),

632 fsdp_allgather_prefetch=getattr(

633 compile_config, "fsdp_allgather_prefetch", True

634 ),

635 fsdp_overlap_grad_reduce=getattr(

636 compile_config, "fsdp_overlap_grad_reduce", True

637 ),

638 max_in_flight_gb=getattr(compile_config, "max_in_flight_gb", 0.25),

639 ac_mode=ac_mode,

640 ac_layer_count=ac_layer_count,

641 ac_layer_indices=ac_layer_indices,

642 ac_region_source=ac_region_source,

643 activation_swap=getattr(config, "activation_swap", "none"),

644 offload_budget_gb=getattr(compile_config, "offload_budget_gb", 64.0),

645 offload_prefetch_n_layers=getattr(

646 compile_config, "offload_prefetch_n_layers", 1

647 ),

648 )

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / pass_config.py: 96%

165 statements