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

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"""Activation Checkpoint Pass - static-graph activation checkpointing.

17FX-level activation recomputation for the joint forward+backward graph,

18following torchtitan graph_trainer's ``tag_with_memory_policy_pass`` +

19``selective_activation_remat_pass`` design while reusing the dynamic-graph

20configuration vocabulary (``ActivationCheckpointConfig``: ``mode`` =

21off/full/selective, layer selection, model-adapter safe regions):

231. Tag: forward ``call_function`` nodes inside the checkpoint scope get a

24 ``CheckpointPolicy`` in ``node.meta["recompute"]``. The scope is the

25 discovered transformer blocks (the same HuggingFace-marker discovery the

26 eager path requires), narrowed by the optional layer selectors and the

27 model adapter's safe regions. Nodes outside the blocks (embedding, head,

28 the ``train_fn`` loss chain) and compute/communication intensive

29 operators under ``selective`` stay saved; everything else is a recompute

30 candidate (``full`` forces every scoped node). Structural exceptions are

31 always saved: RNG ops (the graph cannot replay random state), in-place

32 mutations with their consumers and the readers of every tensor sharing

33 the mutated storage (including view aliases), SymInt scalars.

342. Remat: every tagged node that the backward half consumes (transitively

35 through other tagged nodes) is duplicated in front of its earliest

36 backward consumer and the backward consumers are redirected to the

37 duplicate. A forward original with no remaining readers is erased; one

38 still read by forward code stays, and its backward readers now rebuild

39 the value during backward. Either way the activation's forward lifetime

40 ends at its last forward use instead of spanning the whole step. A

41 cross-unit forward consumer forces the producing node to stay saved

42 (recreating it would only waste compute), which keeps each recompute

43 chain inside one block.

45The backward half is identified by ``meta["autograd_backward"]``, which the

46joint tracer guarantees on every supported torch (engine tag on patched

47torch, structural reconstruction otherwise -- see

48``tracer/graph_tracer.py``).

50Ordering contract: the pass must run BEFORE ``FSDPPass``. Recomputed

51duplicates are ordinary forward ops at tag time, so ``FSDPPass`` inserts

52parameter all-gathers and the backward re-gathers around them correctly

53(duplicates are tagged ``autograd_backward`` and count as backward

54consumers). Running AC after FSDP would make recompute read sharded

55placeholders. ``PpPass`` runs later and inherits the rewrite value-level.

57Operator policy sets and transformer-block discovery are shared with the

58eager path (``hyper_parallel.distributed.activation_checkpoint``), so a

59model sees the same save/recompute vocabulary in both modes. RNG ops are

60saved in every mode (a graph cannot replay random state). For collectives,

61the eager selective must-save set (all-to-all/reduce-scatter/in-place

62allreduce) is always saved; other collectives follow the mode policy.

63``ac_save_forward_all_gather`` may opt forward (TP/SP/CP) all-gathers into

64backward replay.

65"""

67from __future__ import annotations

69__all__ = ["ActivationCheckpointPass", "prepare_activation_checkpoint"]

71import functools

72import logging

73import operator

74from types import SimpleNamespace

75from typing import Any, Callable, Dict, List, Optional, Set, Tuple

77import torch

78from torch import fx

79from torch._functorch import compile_utils

80from torch._ops import OpOverload

81from torch.utils.checkpoint import CheckpointPolicy

83# Same-package reuse of the eager AC machinery: operator save/recompute sets,

84# HF transformer-block discovery and model-adapter recompute specs. They are

85# module-level privates because the eager module never exported them, but the

86# coupling is the point: one policy vocabulary for both modes.

87from hyper_parallel.distributed.activation_checkpoint import (

88 _SELECTIVE_AC_COMM_OP_NAMES,

89 _SELECTIVE_AC_FORCE_RECOMPUTE_OPS,

90 _SELECTIVE_AC_MUST_SAVE_OPS,

91 _adapter_recompute_spec,

92 _existing_ops,

93 _find_transformer_layer_container_infos,

94 _resolve_op_attr,

95 _resolve_torch_op,

96 _select_excluded_targets,

97 _select_layer_region_roots,

98 _selected_checkpoint_layer_fqns,

99 _validate_region_exclusions,

100)

101

102from ..pass_config import PassConfig

103from .base import GraphPass

104

105_LOG = logging.getLogger(__name__)

106

107# FSDP all-gathers are inserted after this pass, but forward (TP/SP/CP)

108# all-gathers are already in the traced forward; saving them avoids a second

109# collective during backward recompute. ``PassConfig.ac_save_forward_all_gather``

110# selects that superset. Only the selective communication set is always saved:

111# other collectives follow the mode policy, matching the eager path.

112_GRAPH_AC_COMM_OPS = _existing_ops(

113 *(_resolve_torch_op(name) for name in _SELECTIVE_AC_COMM_OP_NAMES),

114 _resolve_op_attr(torch.ops, "deepep.dispatch.default"),

115 _resolve_op_attr(torch.ops, "deepep.combine.default"),

116 _resolve_op_attr(torch.ops, "hybridep.dispatch.default"),

117 _resolve_op_attr(torch.ops, "hybridep.combine.default"),

118)

119_GRAPH_FORWARD_ALL_GATHER_OPS = _existing_ops(

120 _resolve_torch_op("aten.all_gather_into_tensor"),

121 _resolve_torch_op("_c10d_functional.all_gather_into_tensor"),

122)

123

124

125def _resolve_wait_tensor() -> Optional[Any]:

126 """Return the c10d ``wait_tensor`` overload when this torch ships it."""

127 functional = getattr(torch.ops, "_c10d_functional", None)

128 wait = getattr(functional, "wait_tensor", None)

129 if wait is None:

130 return None

131 return getattr(wait, "default", None)

132

133

134# getitem / wait_tensor extract or await their producer's value: their

135# save/recompute decision is inherited, and remat walks through them.

136_TRANSPARENT_OPS: Tuple[Any, ...] = tuple(

137 target

138 for target in (operator.getitem, _resolve_wait_tensor())

139 if target is not None

140)

141

142_RECOMPUTE_POLICIES = (

143 CheckpointPolicy.PREFER_RECOMPUTE,

144 CheckpointPolicy.MUST_RECOMPUTE,

145)

146

147

148def _must_recompute(node: fx.Node) -> bool:

149 """Whether ``node`` carries a policy the partitioner would recompute."""

150 return node.meta.get("recompute") in _RECOMPUTE_POLICIES

151

152

153def _node_stack_fqns(node: fx.Node) -> List[str]:

154 """Module FQNs on ``node``'s ``nn_module_stack``, innermost first."""

155 stack = node.meta.get("nn_module_stack") or {}

156 return [fqn for fqn, _ in reversed(list(stack.values()))]

157

158

159def _at_or_under(fqn: str, root: str) -> bool:

160 """Whether ``fqn`` equals ``root`` or lives below it."""

161 return fqn == root or fqn.startswith(f"{root}.")

162

163

164def _first_container(fqns: List[str], roots: frozenset) -> Optional[str]:

165 """Return the first root that contains one of ``fqns`` (innermost wins).

166

167 ``roots`` is sorted so a (currently impossible) nested-root match is still

168 deterministic. ``fqns`` is innermost-first, so the first hit is the

169 innermost containing root.

170 """

171 for fqn in fqns:

172 for root in sorted(roots):

173 if _at_or_under(fqn, root):

174 return root

175 return None

176

177

178class _CheckpointScope:

179 """Resolved recompute scope for one compilation.

180

181 ``_resolve_scope`` always fills every field; the ``None`` defaults exist

182 for direct construction in tests and mean "no region" / "no boundary

183 units", which makes every node ineligible / disables the cross-unit rule.

184

185 Attributes:

186 region_fqns: FQNs whose subtree may recompute. Empty (or ``None``)

187 makes every node ineligible (the pass becomes a no-op).

188 exclude_fqns: FQNs whose subtree must stay saved (adapter exclusions).

189 unit_fqns: Boundary units for the cross-unit save rule. Empty (or

190 ``None``) disables the rule.

191 on_applied: Adapter runtime hook that must run before the model's

192 forward is traced, when the adapter's region shape is in effect.

193 """

194

195 def __init__(

196 self,

197 region_fqns: Optional[frozenset] = None,

198 exclude_fqns: frozenset = frozenset(),

199 unit_fqns: Optional[frozenset] = None,

200 on_applied: Optional[Callable[[torch.nn.Module], None]] = None,

201 ) -> None:

202 """Store the resolved scope membership sets."""

203 self.region_fqns = region_fqns

204 self.exclude_fqns = exclude_fqns

205 self.unit_fqns = unit_fqns

206 self.on_applied = on_applied

207

208

209def _adapter_scope(

210 model: torch.nn.Module,

211 selected_layers: frozenset,

212 unit_fqns: frozenset,

213 recompute_spec: Any,

214) -> _CheckpointScope:

215 """Build a scope from the adapter-declared recompute regions."""

216 module_by_fqn = dict(model.named_modules())

217 roots = _select_layer_region_roots(module_by_fqn, selected_layers, recompute_spec)

218 if not roots:

219 raise ValueError(

220 "model adapter recompute patterns matched no region in the selected layers"

221 )

222 excluded = _select_excluded_targets(module_by_fqn, roots, recompute_spec)

223 if recompute_spec.exclude_patterns and not excluded:

224 raise ValueError(

225 "model adapter recompute exclusions matched no module "

226 "inside the selected regions"

227 )

228 _validate_region_exclusions(roots, excluded)

229 return _CheckpointScope(

230 region_fqns=frozenset(fqn for fqn, _ in roots),

231 exclude_fqns=frozenset(fqn for fqn, _ in excluded),

232 unit_fqns=unit_fqns,

233 on_applied=recompute_spec.on_applied,

234 )

235

236

237def _resolve_scope(

238 model: Optional[torch.nn.Module], pass_config: PassConfig

239) -> _CheckpointScope:

240 """Resolve the checkpoint scope from the dynamic-config vocabulary.

241

242 The recompute scope and the boundary units are the discovered

243 transformer blocks — the same HuggingFace-marker discovery the eager

244 path requires (``_resolve_activation_checkpoint_containers``), so a

245 model is AC-eligible in both modes under the same contract. Nodes

246 outside the blocks (embedding, head, the ``train_fn`` loss chain) are

247 never tagged, which also keeps the head/epilogue backward from pulling

248 the whole forward into one recompute chain. With

249 ``ac_layer_count`` / ``ac_layer_indices`` only the selected blocks

250 recompute; an adapter ``RecomputePolicy`` (``model_adapter_safe_regions``

251 in the dynamic config) narrows the selected blocks to its region

252 patterns and pulls ``exclude_patterns`` back out of the recompute.

253

254 ``full`` with the default selection mirrors the eager auto-upgrade: when

255 the model owns a ``RecomputePolicy``, all blocks are selected through its

256 regions instead of recomputing every op in the blocks.

257

258 Raises:

259 ValueError: When the live model is missing, no transformer blocks

260 are discovered, an explicitly requested safe-region selection has

261 no adapter, or the adapter declares no recompute region.

262 """

263 if model is None:

264 raise ValueError(

265 "ActivationCheckpointPass requires the live model via kwargs "

266 "(model=...) so it can discover the transformer blocks to "

267 "checkpoint"

268 )

269 containers = _find_transformer_layer_container_infos(model)

270 if not containers:

271 raise ValueError(

272 f"{type(model).__name__} has no module with a 'gradient_checkpointing' "

273 "attribute and a non-empty repeated block container; activation "

274 "checkpointing cannot resolve its layers"

275 )

276 block_fqns = frozenset(

277 block.fqn for container in containers for block in container.blocks

278 )

279 explicit_regions = pass_config.ac_region_source == "model_adapter_safe_regions"

280 has_count = pass_config.ac_layer_count is not None

281 has_indices = pass_config.ac_layer_indices is not None

282 if not has_count and not has_indices:

283 if pass_config.ac_mode == "full":

284 recompute_spec = _adapter_recompute_spec(model, required=explicit_regions)

285 if recompute_spec is not None:

286 if not recompute_spec.region_patterns:

287 raise ValueError(

288 "model adapter recompute requires "

289 "RecomputePolicy.region_patterns"

290 )

291 return _adapter_scope(model, block_fqns, block_fqns, recompute_spec)

292 return _CheckpointScope(region_fqns=block_fqns, unit_fqns=block_fqns)

293 selection = SimpleNamespace(

294 layer_count=pass_config.ac_layer_count,

295 layer_indices=pass_config.ac_layer_indices,

296 )

297 selected_layers = frozenset(_selected_checkpoint_layer_fqns(containers, selection))

298 recompute_spec = _adapter_recompute_spec(model, required=explicit_regions)

299 if explicit_regions:

300 if recompute_spec is None:

301 raise ValueError(

302 "activation_checkpoint.selection.source=model_adapter_safe_regions "

303 "requires ModelAdapterSpec.recompute"

304 )

305 if not recompute_spec.region_patterns:

306 raise ValueError(

307 "model adapter recompute requires RecomputePolicy.region_patterns"

308 )

309 if not selected_layers:

310 # Zero selected layers is a no-op, not an error.

311 _LOG.info("activation checkpoint layer selection matched zero layers")

312 return _CheckpointScope(region_fqns=frozenset(), unit_fqns=block_fqns)

313 if recompute_spec is None or not recompute_spec.region_patterns:

314 return _CheckpointScope(region_fqns=selected_layers, unit_fqns=block_fqns)

315 return _adapter_scope(model, selected_layers, block_fqns, recompute_spec)

316

317

318def prepare_activation_checkpoint(

319 model: Optional[torch.nn.Module], pass_config: PassConfig

320) -> _CheckpointScope:

321 """Resolve the AC scope and run the adapter hook before tracing.

322

323 The adapter's ``on_applied`` hook configures runtime state the model's

324 forward reads, so it must run before the joint graph is traced.

325 ``GraphCompiler`` calls this once pre-trace and hands the scope to the

326 pass, which avoids a second hook invocation.

327

328 Args:

329 model: The live model whose transformer blocks define the scope.

330 pass_config: Parallel configuration driving ``ac_mode``.

331

332 Returns:

333 The resolved scope, including the adapter runtime hook when one is

334 declared.

335 """

336 scope = _resolve_scope(model, pass_config)

337 if scope.on_applied is not None and model is not None:

338 scope.on_applied(model)

339 return scope

340

341

342def _is_eligible(node: fx.Node, scope: _CheckpointScope) -> bool:

343 """Whether ``node`` may carry a recompute decision at all.

344

345 Nodes outside the checkpoint regions stay saved: activation

346 checkpointing targets transformer-block forwards, mirroring the eager

347 path that wraps the discovered blocks only.

348 """

349 if not scope.region_fqns:

350 return False

351 fqns = _node_stack_fqns(node)

352 if not fqns:

353 return False

354 if _first_container(fqns, scope.exclude_fqns) is not None:

355 return False

356 return _first_container(fqns, scope.region_fqns) is not None

357

358

359def _make_policy_fn(pass_config: PassConfig) -> Callable[[fx.Node], CheckpointPolicy]:

360 """Build the per-node policy for ``ac_mode``.

361

362 ``full`` recomputes every scoped node. ``selective`` saves the shared

363 compute/communication operator set and forces recompute only on the

364 operators the eager policy must re-replay (in-place-unsafe outputs such

365 as ``topk``); everything else prefers recompute.

366 """

367 if pass_config.ac_mode == "full":

368 return lambda node: CheckpointPolicy.MUST_RECOMPUTE

369

370 must_save = _SELECTIVE_AC_MUST_SAVE_OPS

371 force_recompute = _SELECTIVE_AC_FORCE_RECOMPUTE_OPS

372

373 def selective_policy(node: fx.Node) -> CheckpointPolicy:

374 """Apply the shared selective-AC operator policy."""

375 if node.target in force_recompute:

376 return CheckpointPolicy.MUST_RECOMPUTE

377 if node.target in must_save:

378 return CheckpointPolicy.MUST_SAVE

379 return CheckpointPolicy.PREFER_RECOMPUTE

380

381 return selective_policy

382

383

384def _unit_of(node: fx.Node, scope: _CheckpointScope) -> Optional[str]:

385 """Return ``node``'s boundary block unit, if any."""

386 if not scope.unit_fqns:

387 return None

388 return _first_container(_node_stack_fqns(node), scope.unit_fqns)

389

390

391def _tag_policy_pass(

392 graph_module: fx.GraphModule,

393 pass_config: PassConfig,

394 scope: _CheckpointScope,

395) -> None:

396 """Tag forward nodes per policy, applying the structural exceptions.

397

398 RNG ops and SymInt scalars stay saved; getitem/wait_tensor inherit their

399 parent's decision. The selective communication set is saved in every mode

400 (including ``full``), matching the eager selective policy; all-gathers join

401 the saved set when ``ac_save_forward_all_gather`` is enabled. Other

402 collectives (functional all-reduce/broadcast) follow the mode policy.

403 """

404 policy_fn = _make_policy_fn(pass_config)

405 comm_ops = _GRAPH_AC_COMM_OPS

406 if pass_config.ac_save_forward_all_gather:

407 comm_ops = comm_ops | _GRAPH_FORWARD_ALL_GATHER_OPS

408 for node in graph_module.graph.nodes:

409 if node.op != "call_function" or node.meta.get("autograd_backward", False):

410 continue

411 if node.target in _TRANSPARENT_OPS:

412 parent = node.args[0]

413 if isinstance(parent, fx.Node) and "recompute" in parent.meta:

414 node.meta["recompute"] = parent.meta["recompute"]

415 continue

416 if not _is_eligible(node, scope):

417 continue

418 if _is_rng_op(node) or node.target in comm_ops:

419 node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

420 continue

421 value = node.meta.get("val")

422 if isinstance(value, torch.SymInt):

423 node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

424 continue

425 node.meta["recompute"] = policy_fn(node)

426

427

428def _flatten_getitems(graph_module: fx.GraphModule) -> None:

429 """Flatten tuple projections when this torch build provides the helper."""

430 raise_getitems = getattr(compile_utils, "raise_getitems", None)

431 if raise_getitems is not None:

432 raise_getitems(graph_module)

433

434

435def _is_mutable_op(target: Any) -> bool:

436 """Return whether ``target`` is an in-place ``OpOverload``."""

437 if not isinstance(target, OpOverload):

438 return False

439 return bool(getattr(getattr(target, "_schema", None), "is_mutable", False))

440

441

442def _is_rng_op(node: fx.Node) -> bool:

443 """Return True when a node's seeded nondeterminism cannot be ruled out."""

444 return bool(

445 hasattr(node.target, "tags")

446 and torch.Tag.nondeterministic_seeded in node.target.tags

447 )

448

449

450def _mutated_arg_nodes(node: fx.Node) -> List[fx.Node]:

451 """Return the input nodes whose tensors ``node`` overwrites in place."""

452 schema = getattr(node.target, "_schema", None)

453 if schema is None:

454 return []

455 mutated: List[fx.Node] = []

456 for position, schema_arg in enumerate(schema.arguments):

457 alias_info = getattr(schema_arg, "alias_info", None)

458 if alias_info is None or not alias_info.is_write:

459 continue

460 arg = (

461 node.args[position]

462 if position < len(node.args)

463 else node.kwargs.get(schema_arg.name)

464 )

465 if isinstance(arg, fx.Node):

466 mutated.append(arg)

467 return mutated

468

469

470def _aliased_input_nodes(node: fx.Node) -> List[fx.Node]:

471 """Return input nodes whose storage ``node``'s outputs may share.

472

473 Alias relationships come from the operator schema. A return carrying

474 ``alias_info`` aliases the arguments whose alias sets intersect; several

475 multi-output views (``split`` / ``chunk`` / ``unbind``) instead mark the

476 input with an ``after_set`` and leave the returns unannotated. In-place

477 writes always share their argument's storage. ``getitem`` picks one output

478 of a producer, so it inherits the producer's storage even though the

479 builtin itself has no schema.

480 """

481 if node.op != "call_function":

482 return []

483 if node.target is operator.getitem:

484 producer = node.args[0] if node.args else None

485 return [producer] if isinstance(producer, fx.Node) else []

486 schema = getattr(node.target, "_schema", None)

487 if schema is None:

488 return []

489 result_names: Set[str] = set()

490 for result in schema.returns:

491 alias_info = getattr(result, "alias_info", None)

492 if alias_info is not None:

493 result_names.update(alias_info.before_set)

494 result_names.update(alias_info.after_set)

495 wildcard_result = "*" in result_names

496 aliased: List[fx.Node] = []

497 for position, schema_arg in enumerate(schema.arguments):

498 alias_info = getattr(schema_arg, "alias_info", None)

499 if alias_info is None:

500 continue

501 arg_names = set(alias_info.before_set) | set(alias_info.after_set)

502 if not (

503 getattr(alias_info, "is_write", False)

504 or wildcard_result

505 or "*" in arg_names

506 or bool(arg_names & result_names)

507 ):

508 continue

509 arg = (

510 node.args[position]

511 if position < len(node.args)

512 else node.kwargs.get(schema_arg.name)

513 )

514 if isinstance(arg, fx.Node):

515 aliased.append(arg)

516 return aliased

517

518

519def _forward_alias_neighbors(

520 graph_module: fx.GraphModule,

521) -> Dict[fx.Node, Set[fx.Node]]:

522 """Build the undirected shared-storage graph over the forward graph."""

523 neighbors: Dict[fx.Node, Set[fx.Node]] = {}

524 for node in graph_module.graph.nodes:

525 if node.meta.get("autograd_backward", False):

526 continue

527 for arg in _aliased_input_nodes(node):

528 if arg.meta.get("autograd_backward", False):

529 continue

530 neighbors.setdefault(node, set()).add(arg)

531 neighbors.setdefault(arg, set()).add(node)

532 return neighbors

533

534

535def _alias_closure(

536 roots: Set[fx.Node], neighbors: Dict[fx.Node, Set[fx.Node]]

537) -> Set[fx.Node]:

538 """Return every forward node sharing storage with any of ``roots``.

539

540 The walk is undirected so a mutation through one view reaches its base

541 (backward along the alias edge) and then the sibling views of that base

542 (forward along their edges).

543 """

544 closure = set(roots)

545 stack = list(roots)

546 while stack:

547 current = stack.pop()

548 for neighbor in neighbors.get(current, ()):

549 if neighbor not in closure:

550 closure.add(neighbor)

551 stack.append(neighbor)

552 return closure

553

554

555def _tag_inplace_saves(graph_module: fx.GraphModule) -> None:

556 """Force saves around in-place mutations and their aliased-storage readers.

557

558 The mutation must run exactly once in forward, so it stays saved together

559 with its forward consumers. Every forward reader of the mutated storage

560 stays saved as well: a recompute duplicate would re-execute against either

561 the post-mutation original or a pre-mutation duplicate of the producer,

562 and neither matches what forward computed at the reader's own position.

563 Eager autograd stays silent for such a read — its value reaches backward

564 only through the recompute chain, never as a saved tensor. The storage

565 closure follows operator-schema aliases, so readers of a view (or of a

566 sibling view when the write goes through another view) are covered, and

567 mutations outside the checkpoint scope still contribute their storage.

568 """

569 nodes = list(graph_module.graph.nodes)

570 mutated_bases: Set[fx.Node] = set()

571 for node in nodes:

572 if (

573 node.op != "call_function"

574 or node.meta.get("autograd_backward", False)

575 or not _is_mutable_op(node.target)

576 ):

577 continue

578 mutated_bases.update(_mutated_arg_nodes(node))

579 if "recompute" not in node.meta:

580 continue

581 node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

582 for user in node.users:

583 if (

584 not user.meta.get("autograd_backward", False)

585 and "recompute" in user.meta

586 ):

587 user.meta["recompute"] = CheckpointPolicy.MUST_SAVE

588 if not mutated_bases:

589 return

590 aliased_storage = _alias_closure(

591 mutated_bases, _forward_alias_neighbors(graph_module)

592 )

593 for node in nodes:

594 if (

595 node.op != "call_function"

596 or node.meta.get("autograd_backward", False)

597 or node.meta.get("recompute") not in _RECOMPUTE_POLICIES

598 ):

599 continue

600 if not aliased_storage.isdisjoint(node.all_input_nodes):

601 node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

602

603

604def _tag_cross_unit_saves(

605 graph_module: fx.GraphModule, scope: _CheckpointScope

606) -> None:

607 """Save recomputable nodes read forward across a unit boundary.

608

609 Their original activation stays alive for the forward consumer anyway, so

610 recomputing them in backward would be pure waste.

611 """

612 for node in graph_module.graph.nodes:

613 if (

614 node.op != "call_function"

615 or node.meta.get("autograd_backward", False)

616 or node.meta.get("recompute") not in _RECOMPUTE_POLICIES

617 ):

618 continue

619 if _has_cross_unit_forward_consumer(node, scope):

620 node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

621

622

623def _tag_sac_policy(

624 graph_module: fx.GraphModule,

625 pass_config: PassConfig,

626 scope: _CheckpointScope,

627) -> None:

628 """Annotate forward nodes with ``meta["recompute"]`` policies.

629

630 Pass 1 tags per policy with the structural exceptions; pass 2 forces saves

631 on in-place mutations, their consumers and the forward readers of the

632 mutated storage (view aliases included); pass 3 forces saves on

633 recomputable nodes read forward across a unit boundary (their original

634 stays alive anyway, so the backward recompute would be pure waste).

635 """

636 if not scope.region_fqns:

637 return

638 _tag_policy_pass(graph_module, pass_config, scope)

639 _tag_inplace_saves(graph_module)

640 if scope.unit_fqns:

641 _tag_cross_unit_saves(graph_module, scope)

642

643

644def _has_cross_unit_forward_consumer(node: fx.Node, scope: _CheckpointScope) -> bool:

645 """Whether ``node`` is read forward outside its own boundary unit.

646

647 Transparent (inherited) getitem/wait_tensor chains are walked through;

648 any other forward consumer in a different unit (or outside every unit,

649 e.g. the epilogue) means the original must stay alive anyway.

650 """

651 node_unit = _unit_of(node, scope)

652 policy = node.meta["recompute"]

653 seen = {node}

654 stack = [node]

655 while stack:

656 current = stack.pop()

657 for user in current.users:

658 if user in seen:

659 continue

660 seen.add(user)

661 if user.op == "output" or user.meta.get("autograd_backward", False):

662 continue

663 if (

664 user.op == "call_function"

665 and user.target in _TRANSPARENT_OPS

666 and user.meta.get("recompute") == policy

667 ):

668 stack.append(user)

669 continue

670 if _unit_of(user, scope) != node_unit:

671 return True

672 return False

673

674

675def _get_attr_value(graph_module: fx.GraphModule, target: str) -> Any:

676 """Resolve an FX ``get_attr`` target, including dotted qualified names."""

677 try:

678 return graph_module.get_submodule(target)

679 except AttributeError:

680 return getattr(graph_module, target, None)

681

682

683def _duplicate_subgraph_get_attrs(

684 graph_module: fx.GraphModule, fwd_node: fx.Node, bwd_target: fx.Node

685) -> Dict[fx.Node, fx.Node]:

686 """Give a recompute duplicate private copies of HOP subgraph get_attrs.

687

688 Higher-order ops (flex_attention etc.) reference their subgraphs through

689 ``get_attr`` nodes. Sharing the original would let a later regional

690 Inductor pass place the shared subgraph in only one partition; fresh

691 get_attr nodes keep each region self-contained.

692 """

693 duplicates: Dict[fx.Node, fx.Node] = {}

694 for arg in fwd_node.all_input_nodes:

695 if arg.op != "get_attr":

696 continue

697 if not isinstance(_get_attr_value(graph_module, arg.target), fx.GraphModule):

698 continue

699 with graph_module.graph.inserting_before(bwd_target):

700 dup_attr = graph_module.graph.get_attr(arg.target)

701 dup_attr.meta.update(arg.meta)

702 duplicates[arg] = dup_attr

703 return duplicates

704

705

706def _remat(graph_module: fx.GraphModule) -> int:

707 """Duplicate tagged recompute nodes in front of their backward consumers.

708

709 Returns the number of duplicated forward nodes. The graph is mutated in

710 place; originals whose consumers were all redirected are erased, and

711 ``raise_getitems`` flattens output/tuple projections for better memory.

712 """

713 graph = graph_module.graph

714 all_nodes = list(graph.nodes)

715 order = {node: idx for idx, node in enumerate(all_nodes)}

716

717 # Map each tagged forward node to its earliest backward consumer; the

718 # transitive walk through tagged ancestors mirrors the partitioner's

719 # recompute-closure gathering (saved ancestors stay live, so only tagged

720 # ones need duplicating).

721 remat_targets: Dict[fx.Node, fx.Node] = {}

722 for bwd_node in all_nodes:

723 if not bwd_node.meta.get("autograd_backward", False):

724 continue

725 seen: Set[fx.Node] = set()

726 stack = list(bwd_node.all_input_nodes)

727 while stack:

728 current = stack.pop()

729 if current in seen or not _must_recompute(current):

730 continue

731 seen.add(current)

732 remat_targets.setdefault(current, bwd_node)

733 stack.extend(current.all_input_nodes)

734

735 if not remat_targets:

736 return 0

737

738 recomputed: Dict[fx.Node, fx.Node] = {}

739

740 def remat_arg(arg: Any, get_attr_dups: Dict[fx.Node, fx.Node]) -> Any:

741 """Redirect duplicate args to their recompute copies."""

742 if isinstance(arg, fx.Node):

743 if arg in get_attr_dups:

744 return get_attr_dups[arg]

745 if arg in recomputed:

746 return recomputed[arg]

747 return arg

748

749 for fwd_node in sorted(remat_targets, key=order.__getitem__):

750 bwd_target = remat_targets[fwd_node]

751 get_attr_dups = _duplicate_subgraph_get_attrs(

752 graph_module, fwd_node, bwd_target

753 )

754 with graph.inserting_before(bwd_target):

755 duplicate = graph.node_copy(

756 fwd_node,

757 functools.partial(remat_arg, get_attr_dups=get_attr_dups),

758 )

759 # Leave the name to ``node_copy``/``create_node`` so the graph's name

760 # namespace owns uniqueness; a direct ``.name =`` assignment bypasses

761 # it and can collide with an existing node.

762 duplicate.meta["autograd_backward"] = True

763 duplicate.meta["ac_recomputed"] = True

764 duplicate.meta["ac_recomputed_from"] = fwd_node.name

765 # node_copy shares meta, including the policy tag. Drop it so a second

766 # pass run cannot mistake the duplicate for a recompute candidate.

767 duplicate.meta.pop("recompute", None)

768 recomputed[fwd_node] = duplicate

769

770 def redirect(arg: Any) -> Any:

771 """Arg transform for backward consumers of recomputed originals."""

772 if isinstance(arg, fx.Node) and arg in recomputed:

773 return recomputed[arg]

774 return arg

775

776 direct_backward_consumers = {

777 user

778 for fwd_node in recomputed

779 for user in fwd_node.users

780 if user.meta.get("autograd_backward", False)

781 }

782 for bwd_node in direct_backward_consumers:

783 bwd_node.args = fx.map_arg(bwd_node.args, redirect)

784 bwd_node.kwargs = fx.map_arg(bwd_node.kwargs, redirect)

785

786 for original in reversed(list(recomputed)):

787 if not original.users:

788 graph.erase_node(original)

789

790 _flatten_getitems(graph_module)

791 graph.lint()

792 # ``raise_getitems`` recompiles when present; the explicit call also covers

793 # torch builds without it, where the inserted/erased nodes would otherwise

794 # leave the GraphModule's generated code stale.

795 graph_module.recompile()

796 return len(recomputed)

797

798

799class ActivationCheckpointPass(GraphPass):

800 """Static-graph activation checkpointing (tag + recompute-materialize).

801

802 Configure through ``PassConfig.ac_mode`` (``"off"`` / ``"full"`` /

803 ``"selective"``) and the optional ``ac_layer_count`` /

804 ``ac_layer_indices`` layer selectors. The trainer-level

805 ``activation_checkpoint`` section is projected onto these fields by

806 ``build_pass_config_from_trainer_config``.

807 """

808

809 name = "activation_checkpoint"

810

811 def run(

812 self,

813 graph_module: fx.GraphModule,

814 pass_config: PassConfig,

815 **kwargs: Any,

816 ) -> fx.GraphModule:

817 """Tag the joint graph and materialize recompute duplicates.

818

819 Args:

820 graph_module: The joint forward+backward FX graph.

821 pass_config: Parallel configuration driving ``ac_mode``.

822 **kwargs: ``model`` (the live module) is required whenever

823 ``ac_mode != "off"``: block discovery resolves the recompute

824 scope from it. ``ac_scope`` is the pre-resolved scope from

825 ``GraphCompiler`` (it also ran the adapter hook before

826 tracing); direct callers may omit it.

827

828 Returns:

829 The transformed graph module.

830 """

831 if pass_config.ac_mode == "off":

832 return graph_module

833 scope = kwargs.get("ac_scope")

834 if scope is None:

835 scope = prepare_activation_checkpoint(kwargs.get("model"), pass_config)

836 already_rematted = any(

837 node.meta.get("ac_recomputed") for node in graph_module.graph.nodes

838 )

839 _tag_sac_policy(graph_module, pass_config, scope)

840 duplicated = _remat(graph_module)

841 _LOG.info(

842 "Activation checkpointing (mode=%s) recomputed %d node(s)",

843 pass_config.ac_mode,

844 duplicated,

845 )

846 if duplicated == 0 and scope.region_fqns and not already_rematted:

847 _LOG.warning(

848 "activation checkpointing matched scope %s but recomputed no "

849 "node; the model may already be wrapped by eager activation "

850 "checkpointing, so graph-mode recompute is inactive",

851 sorted(scope.region_fqns),

852 )

853 return graph_module

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / activation_checkpoint.py: 92%

349 statements