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

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

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

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Static-graph CPU offload - host-side activation swap. 

16 

17Graph-mode counterpart of the dynamic ``activation_swap="attention"`` setting 

18(``hyper_parallel.distributed.attention_swap``). It reuses the same public 

19configuration surface - the Trainer's top-level ``activation_swap`` field - 

20and is driven by ``PassConfig.activation_swap``. The implementation follows 

21torchtitan graph_trainer's ``cpu_offload.py``: 

22 

231. Select forward activations produced inside ``self_attn`` modules that the 

24 backward half consumes (directly or through a chain of aliases), are large 

25 enough to amortize a transfer, are contiguous and independently owned, are 

26 never written in place, and fit the host-memory budget (largest-first). 

27 Aliases are grouped by real storage identity, so replay ops without schema 

28 ``alias_info`` (for example ``aten::_unsafe_view``) share the owner's single 

29 offload and lifetime. In-place writes are detected from the schema (including 

30 ops like ``aten::_foreach_add_`` that return no tensor), so a storage that 

31 changes after an earlier offload is never moved to host. 

322. Insert the asynchronous transfer after each producer: 

33 

34 ``cpu = ao.offload(gpu)`` is launched right after the producer, and 

35 ``cpu = ao.wait_tensor(cpu, keepalive=gpu, last_use_of_storage=...)`` is 

36 placed after the last forward read of the whole storage group. The D2H copy 

37 therefore overlaps the remaining forward compute, and the GPU storage is 

38 released as soon as the wait runs. 

393. Insert ``gpu = ao.reload(cpu, device)`` plus a ``wait_tensor`` before the 

40 first backward consumer, redirecting every backward consumer (including 

41 consumers reached through replayed aliases) to the reloaded tensor. Reloads 

42 may be prefetched a few attention regions 

43 earlier to overlap H2D with preceding backward compute. 

44 

45The pass runs before ``FSDPPass``: it only rewrites saved forward values, so 

46FSDP sees the resulting backward consumers when it inserts parameter 

47re-gathers. Attention swap is incompatible with activation checkpointing (the 

48dynamic rule is enforced in ``PassConfig``), so no recompute interacts with it. 

49 

50Only ``activation_swap="attention"`` is enabled because that is the only mode 

51the dynamic attention-swap path exposes. Models without ``self_attn`` modules 

52raise the same error as the eager path. Symbolic (dynamic-shape) activations are 

53supported: selection heuristics (the 1 MiB threshold and the host budget) read 

54the trace-time size hints without specializing the graph, and the transfer ops 

55carry only the device so runtime shapes/strides come from the offloaded tensor. 

56Unbacked (data-dependent) sizes have no hint and are skipped. 

57""" 

58 

59from __future__ import annotations 

60 

61__all__ = ["CpuOffloadPass"] 

62 

63import contextlib 

64import logging 

65from dataclasses import dataclass 

66from typing import Any, Dict, List, Optional, Tuple 

67 

68import torch 

69from torch import fx 

70 

71from .base import GraphPass 

72from .overlap.utils import is_view 

73from ..pass_config import PassConfig 

74 

75try: # Register the upstream ``ao::`` custom ops on import. 

76 import torch._functorch._activation_offloading.offload_ops # noqa: F401 

77 

78 _AO_OFFLOAD = torch.ops.ao.offload.default 

79 _AO_RELOAD = torch.ops.ao.reload.default 

80 _AO_WAIT = torch.ops.ao.wait_tensor.default 

81 _AO_OPS = frozenset((_AO_OFFLOAD, _AO_RELOAD, _AO_WAIT)) 

82 _AO_AVAILABLE = True 

83except (ImportError, AttributeError): 

84 _AO_OFFLOAD = _AO_RELOAD = _AO_WAIT = None 

85 _AO_OPS = frozenset() 

86 _AO_AVAILABLE = False 

87 

88_LOG = logging.getLogger(__name__) 

89 

90# Match the dynamic attention-swap policy: only independently-owned tensors of 

91# at least 1 MiB amortize the transfer; smaller views/activations stay on device. 

92_MIN_OFFLOAD_BYTES = 1024 * 1024 

93 

94_COLLECTIVE_NAMESPACES = frozenset(("_c10d_functional", "_c10d_functional_autograd")) 

95 

96 

97def _is_backward(node: fx.Node) -> bool: 

98 """Whether ``node`` belongs to the joint graph's backward half.""" 

99 return bool(node.meta.get("autograd_backward", False)) 

100 

101 

102def _stack_fqns(node: fx.Node) -> List[str]: 

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

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

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

106 

107 

108def _attention_roots( 

109 graph_module: fx.GraphModule, model: Optional[torch.nn.Module] 

110) -> frozenset: 

111 """Discover the ``self_attn`` module FQNs that define the offload scope. 

112 

113 The eager attention-swap discovery only accepts modules literally named 

114 ``self_attn``; graph mode mirrors that contract and additionally accepts 

115 roots recorded in the traced module stack when no live model is supplied. 

116 """ 

117 roots = set() 

118 if model is not None: 

119 roots.update( 

120 name 

121 for name, _ in model.named_modules() 

122 if name.rsplit(".", 1)[-1] == "self_attn" 

123 ) 

124 for node in graph_module.graph.nodes: 

125 for fqn in _stack_fqns(node): 

126 if fqn.rsplit(".", 1)[-1] == "self_attn": 

127 roots.add(fqn) 

128 return frozenset(roots) 

129 

130 

131def _in_attention_scope(node: fx.Node, roots: frozenset) -> bool: 

132 """Whether a node was produced inside one of the attention roots.""" 

133 if not roots: 

134 return False 

135 for fqn in _stack_fqns(node): 

136 if fqn.rsplit(".", 1)[-1] == "self_attn": 

137 return True 

138 if any(fqn == root or fqn.startswith(f"{root}.") for root in roots): 

139 return True 

140 return False 

141 

142 

143def _attention_root_of(node: fx.Node, roots: frozenset) -> Optional[str]: 

144 """Return the innermost attention root containing ``node`` (or ``None``).""" 

145 for fqn in _stack_fqns(node): 

146 for root in roots: 

147 if fqn == root or fqn.startswith(f"{root}."): 

148 return root 

149 return None 

150 

151 

152def _is_collective_or_wait(node: fx.Node) -> bool: 

153 """Whether the node is a functional collective launch or its wait. 

154 

155 Both the launch namespace and the autograd-compat namespace are covered so 

156 a collective result is never moved to host. 

157 """ 

158 namespace = getattr(node.target, "namespace", None) 

159 return namespace in _COLLECTIVE_NAMESPACES 

160 

161 

162def _is_dtensor(value: torch.Tensor) -> bool: 

163 """Recognize DTensor-like values that the ``ao::`` ops cannot transfer.""" 

164 return hasattr(value, "device_mesh") or hasattr(value, "to_local") 

165 

166 

167def _hint_to_int(value: Any) -> Optional[int]: 

168 """Convert an ``int``/``SymInt`` to a plain int without installing a guard. 

169 

170 Sizes only feed selection heuristics (the 1 MiB threshold and the host 

171 budget), so a symbolic value's trace-time hint is enough. ``int(SymInt)`` 

172 would specialize the graph and add a shape guard; unbacked (data-dependent) 

173 symbols have no hint and yield ``None`` so the candidate is skipped. 

174 """ 

175 try: 

176 if isinstance(value, torch.SymInt): 

177 value = value.node.hint 

178 return int(value) 

179 except Exception: # pylint: disable=broad-exception-caught 

180 return None 

181 

182 

183def _tensor_bytes(value: torch.Tensor) -> Optional[int]: 

184 """Return a hinted byte count, or ``None`` for unsupported values.""" 

185 try: 

186 nbytes = value.nelement() * value.element_size() 

187 except Exception: # pylint: disable=broad-exception-caught 

188 return None 

189 return _hint_to_int(nbytes) 

190 

191 

192def _is_contiguous(value: torch.Tensor) -> bool: 

193 """Check contiguity without tripping over symbolic fake tensors.""" 

194 try: 

195 return bool(value.is_contiguous()) 

196 except Exception: # pylint: disable=broad-exception-caught 

197 return False 

198 

199 

200def _storage_key(value: torch.Tensor) -> Optional[int]: 

201 """Return a stable identity for the tensor's storage, or ``None`` if unknown. 

202 

203 Aliasing is a storage property, not an operator property: replay-based view 

204 ops carry no schema ``alias_info`` (``aten::_unsafe_view``), and some 

205 schema-tagged views can still be full-size. FakeTensor preserves the 

206 underlying storage, so grouping by the real storage (``_cdata``) detects 

207 aliases that ``is_view`` misses without enumerating every op. 

208 

209 ``None`` means the storage identity cannot be proven. The Python object 

210 ``id()`` is deliberately not used as a fallback: it is not stable across 

211 ``untyped_storage()`` wrappers, and a mismatched identity could let an 

212 alias or an in-place write slip through. Returning ``None`` instead makes 

213 callers skip the candidate. 

214 """ 

215 try: 

216 storage = value.untyped_storage() 

217 except Exception: # pylint: disable=broad-exception-caught 

218 return None 

219 cdata = getattr(storage, "_cdata", None) 

220 return int(cdata) if cdata else None 

221 

222 

223def _shares_storage(left: fx.Node, right: fx.Node) -> bool: 

224 """Whether two nodes' values refer to the same storage (unknown => False).""" 

225 left_value = left.meta.get("val") 

226 right_value = right.meta.get("val") 

227 if not isinstance(left_value, torch.Tensor) or not isinstance( 

228 right_value, torch.Tensor 

229 ): 

230 return False 

231 left_key = _storage_key(left_value) 

232 return left_key is not None and left_key == _storage_key(right_value) 

233 

234 

235def _is_alias(user: fx.Node, producer: fx.Node) -> bool: 

236 """Whether ``user`` aliases ``producer``'s storage (schema view or shared).""" 

237 return is_view(user) or _shares_storage(user, producer) 

238 

239 

240def _is_mutable(node: fx.Node) -> bool: 

241 """Whether the node targets an in-place/mutating op.""" 

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

243 return schema is not None and schema.is_mutable 

244 

245 

246def _iter_nodes(value: Any) -> List[fx.Node]: 

247 """Collect the ``fx.Node`` leaves inside an argument (``Tensor[]`` etc.).""" 

248 nodes: List[fx.Node] = [] 

249 fx.map_arg(value, nodes.append) 

250 return nodes 

251 

252 

253def _storage_keys(nodes: List[fx.Node]) -> List[int]: 

254 """Resolvable storage keys of the tensor-valued ``nodes``.""" 

255 keys: List[int] = [] 

256 for node in nodes: 

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

258 if isinstance(value, torch.Tensor): 

259 key = _storage_key(value) 

260 if key is not None: 

261 keys.append(key) 

262 return keys 

263 

264 

265def _tensor_argument_keys(node: fx.Node) -> List[int]: 

266 """Storage keys of every tensor-valued argument to ``node``. 

267 

268 Conservative fallback for nodes whose write target cannot be pinned down: 

269 skipping an offload only costs memory, while missing a write corrupts the 

270 backward reload. 

271 """ 

272 return _storage_keys(_iter_nodes((node.args, node.kwargs))) 

273 

274 

275def _written_storage_keys(node: fx.Node) -> List[int]: 

276 """Storage keys that a mutating op writes through its arguments. 

277 

278 Schema ``alias_info.is_write`` marks the arguments an op modifies in place. 

279 A node's output value is not enough to detect this: ``aten._foreach_add_`` 

280 returns ``()`` but still writes its ``Tensor(a!)[]`` input, so grouping only 

281 returned tensors misses the write and could offload a value that the forward 

282 changes afterwards (making the backward reload stale). Written arguments are 

283 resolved by name first (keyword form) and then by positional index; 

284 ``fx.map_arg`` walks ``Tensor[]`` and nested containers to the nodes. 

285 

286 Two conservative fallbacks keep this sound when the exact write target is 

287 unknowable: 

288 

289 * a node without a schema (``call_method`` / ``call_module`` / plain 

290 callable) may mutate any argument, and 

291 * a mutable schema whose write argument cannot be resolved (vararg layout, 

292 an omitted positional default, or a missing ``meta['val']``). 

293 

294 Both return every tensor argument's storage instead of an empty list, so an 

295 unprovable write never lets a candidate be offloaded. 

296 """ 

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

298 if schema is None: 

299 return _tensor_argument_keys(node) 

300 if not schema.is_mutable: 

301 return [] 

302 

303 keys: List[int] = [] 

304 resolved = False 

305 unresolved = False 

306 position = 0 

307 for argument in schema.arguments: 

308 if argument.kwarg_only: 

309 index = None 

310 else: 

311 index = position 

312 position += 1 

313 alias_info = argument.alias_info 

314 if alias_info is None or not alias_info.is_write: 

315 continue 

316 if argument.name in node.kwargs: 

317 value = node.kwargs[argument.name] 

318 elif index is not None and index < len(node.args): 

319 value = node.args[index] 

320 else: 

321 unresolved = True 

322 continue 

323 written = _iter_nodes(value) 

324 if not written: 

325 continue 

326 resolved = True 

327 written_keys = _storage_keys(written) 

328 if not written_keys: 

329 unresolved = True 

330 keys.extend(written_keys) 

331 if unresolved or not resolved: 

332 return _tensor_argument_keys(node) 

333 return keys 

334 

335 

336def _group_storage_members(graph: fx.Graph) -> Dict[int, List[fx.Node]]: 

337 """Group tensor-valued nodes by the real identity of their storage.""" 

338 members: Dict[int, List[fx.Node]] = {} 

339 for node in graph.nodes: 

340 if node.op not in ("call_function", "placeholder", "get_attr"): 

341 continue 

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

343 if not isinstance(value, torch.Tensor): 

344 continue 

345 key = _storage_key(value) 

346 if key is not None: 

347 members.setdefault(key, []).append(node) 

348 return members 

349 

350 

351def _write_unsafe_keys(graph: fx.Graph) -> set: 

352 """Storage keys written in place anywhere in the graph. 

353 

354 Includes mutations whose return is not a tensor (``_foreach_add_`` returns 

355 ``()``): such a write still changes a storage that may already have been 

356 offloaded. ``call_method`` / ``call_module`` nodes carry no schema; their 

357 tensor arguments are marked conservatively because an unseen in-place write 

358 would make the backward reload stale. 

359 """ 

360 unsafe: set = set() 

361 for node in graph.nodes: 

362 if node.op in ("call_function", "call_method", "call_module"): 

363 unsafe.update(_written_storage_keys(node)) 

364 return unsafe 

365 

366 

367def _group_is_safe(nodes: List[fx.Node]) -> bool: 

368 """Whether a storage group has no external, mutating, or output sharer.""" 

369 if any(node.op != "call_function" for node in nodes): 

370 return False 

371 if any(_is_mutable(node) for node in nodes): 

372 return False 

373 if any(user.op == "output" for node in nodes for user in node.users): 

374 return False 

375 return True 

376 

377 

378@dataclass(frozen=True) 

379class _StorageIndex: 

380 """Real-storage groups used to prove a candidate can be offloaded safely. 

381 

382 A storage may be offloaded only when it is owned by a single forward 

383 ``call_function`` node (the ``owner``), no other node in the graph shares it 

384 through an external (placeholder/get_attr) or mutating op, and it is not a 

385 graph output. Writes are detected from the schema, not only from returned 

386 tensors, so an in-place op with no tensor return still poisons its input 

387 storage. Groups that cannot be proven safe are marked ``unsafe``. 

388 """ 

389 

390 owner: Dict[int, fx.Node] 

391 unsafe: frozenset 

392 

393 @classmethod 

394 def build(cls, graph: fx.Graph, order: Dict[fx.Node, int]) -> "_StorageIndex": 

395 """Group tensor-valued nodes by storage and validate each group.""" 

396 members = _group_storage_members(graph) 

397 unsafe = _write_unsafe_keys(graph) 

398 owner: Dict[int, fx.Node] = {} 

399 for key, nodes in members.items(): 

400 if key in unsafe: 

401 continue 

402 if not _group_is_safe(nodes): 

403 unsafe.add(key) 

404 continue 

405 owner[key] = min(nodes, key=lambda node: order[node]) 

406 return cls(owner=owner, unsafe=frozenset(unsafe)) 

407 

408 

409def _can_offload(node: fx.Node, storage: _StorageIndex) -> bool: 

410 """Whether a forward node's output is a profitable, safely-owned candidate. 

411 

412 Views are skipped because they alias their base, so offloading them frees 

413 nothing and, for ``getitem`` of a multi-output op, could release storage a 

414 sibling still uses. Shared storage is detected by real identity: only the 

415 group owner is eligible, and ``unsafe`` groups (external or mutating 

416 aliases) are left on device. 

417 """ 

418 if node.op != "call_function": 

419 return False 

420 if node.target in _AO_OPS: 

421 return False 

422 if is_view(node): 

423 return False 

424 if _is_collective_or_wait(node): 

425 return False 

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

427 if not isinstance(value, torch.Tensor) or _is_dtensor(value): 

428 return False 

429 if value.dim() < 2: 

430 return False 

431 if not _is_contiguous(value): 

432 return False 

433 nbytes = _tensor_bytes(value) 

434 if nbytes is None or nbytes < _MIN_OFFLOAD_BYTES: 

435 return False 

436 # Offload one storage at a time. Only the group owner (the node that 

437 # allocated the storage) is eligible; aliases are handled by the owner's 

438 # view chain, and external/mutating sharers mark the group unsafe. 

439 key = _storage_key(value) 

440 if key is None or key in storage.unsafe: 

441 return False 

442 if storage.owner.get(key) is not node: 

443 return False 

444 return True 

445 

446 

447def _view_chain( 

448 node: fx.Node, 

449) -> Tuple[List[fx.Node], List[fx.Node], List[Tuple[fx.Node, fx.Node]]]: 

450 """Collect forward aliases of ``node`` and their backward consumers. 

451 

452 Aliases are found by real storage identity as well as schema views, so 

453 replay ops without ``alias_info`` (``aten::_unsafe_view``) join the chain. 

454 

455 Returns: 

456 ``(lifetime_views, replay_views, view_bwd_redirects)``. ``lifetime_views`` 

457 are every forward alias view (needed to find the producer's true last 

458 forward use, so the storage is not released while a forward-only view is 

459 still live). ``replay_views`` is the subset that must be rebuilt on top 

460 of the reloaded tensor because it (or a descendant view) is read in 

461 backward. Each redirect pair ``(view_node, backward_user)`` must be 

462 repointed after replay. 

463 """ 

464 children: Dict[fx.Node, List[fx.Node]] = {} 

465 backward_users: Dict[fx.Node, List[fx.Node]] = {} 

466 lifetime_views: List[fx.Node] = [] 

467 visited: set = set() 

468 

469 def collect(current: fx.Node) -> None: 

470 """Walk alias users and record backward readers.""" 

471 for user in current.users: 

472 if user.op != "call_function": 

473 continue 

474 if _is_backward(user): 

475 backward_users.setdefault(current, []).append(user) 

476 elif _is_alias(user, current) and user not in visited: 

477 visited.add(user) 

478 children.setdefault(current, []).append(user) 

479 lifetime_views.append(user) 

480 collect(user) 

481 

482 collect(node) 

483 

484 needed: set = set() 

485 memo: Dict[fx.Node, bool] = {} 

486 

487 def is_read_in_backward(current: fx.Node) -> bool: 

488 """Whether a view subtree has any backward reader (memoized).""" 

489 cached = memo.get(current) 

490 if cached is not None: 

491 return cached 

492 needed_here = bool(backward_users.get(current)) 

493 for child in children.get(current, ()): 

494 if is_read_in_backward(child): 

495 needed_here = True 

496 memo[current] = needed_here 

497 if needed_here and current is not node: 

498 needed.add(current) 

499 return needed_here 

500 

501 is_read_in_backward(node) 

502 replay_views = [view for view in lifetime_views if view in needed] 

503 redirects = [ 

504 (view, user) for view in replay_views for user in backward_users.get(view, ()) 

505 ] 

506 return lifetime_views, replay_views, redirects 

507 

508 

509@dataclass 

510class _OffloadInfo: 

511 """Resolved backward consumers and forward lifetime of one candidate.""" 

512 

513 node: fx.Node 

514 direct_bwd_users: List[fx.Node] 

515 lifetime_views: List[fx.Node] 

516 replay_views: List[fx.Node] 

517 view_bwd_redirects: List[Tuple[fx.Node, fx.Node]] 

518 first_bwd_consumer: fx.Node 

519 last_forward_use: fx.Node 

520 last_use_arg: fx.Node 

521 

522 @property 

523 def all_bwd_users(self) -> List[fx.Node]: 

524 """Every backward consumer, direct or reached through a view.""" 

525 return self.direct_bwd_users + [user for _, user in self.view_bwd_redirects] 

526 

527 

528def _collect_offload_info( 

529 node: fx.Node, order: Dict[fx.Node, int] 

530) -> Optional[_OffloadInfo]: 

531 """Resolve one candidate's backward users and its last forward use.""" 

532 direct_bwd_users = [user for user in node.users if _is_backward(user)] 

533 lifetime_views, replay_views, view_bwd_redirects = _view_chain(node) 

534 all_bwd_users = direct_bwd_users + [user for _, user in view_bwd_redirects] 

535 if not all_bwd_users: 

536 return None 

537 

538 last_forward_use = node 

539 last_index = order[node] 

540 for chain_node in (node, *lifetime_views): 

541 for user in chain_node.users: 

542 if user.op != "call_function" or _is_backward(user): 

543 continue 

544 user_index = order.get(user) 

545 if user_index is not None and user_index > last_index: 

546 last_forward_use = user 

547 last_index = user_index 

548 

549 if last_forward_use is node: 

550 last_use_arg: fx.Node = node 

551 elif isinstance(last_forward_use.meta.get("val"), torch.Tensor): 

552 last_use_arg = last_forward_use 

553 else: 

554 # The forward ``ao.wait_tensor`` frees the producer's storage, so it must 

555 # carry a data-dependency edge on the last forward consumer 

556 # (``last_use_of_storage``). A tuple/list-returning consumer (split, 

557 # chunk, max, ...) has no single tensor value to pin that edge; skip the 

558 # candidate rather than risk a reordering pass hoisting the storage free 

559 # before its last forward read. 

560 return None 

561 

562 return _OffloadInfo( 

563 node=node, 

564 direct_bwd_users=direct_bwd_users, 

565 lifetime_views=lifetime_views, 

566 replay_views=replay_views, 

567 view_bwd_redirects=view_bwd_redirects, 

568 first_bwd_consumer=min(all_bwd_users, key=lambda item: order[item]), 

569 last_forward_use=last_forward_use, 

570 last_use_arg=last_use_arg, 

571 ) 

572 

573 

574def _collect_candidates( 

575 graph_module: fx.GraphModule, 

576 roots: frozenset, 

577 order: Dict[fx.Node, int], 

578) -> List[Tuple[fx.Node, int, _OffloadInfo]]: 

579 """Collect in-scope forward activations that can be offloaded.""" 

580 storage = _StorageIndex.build(graph_module.graph, order) 

581 candidates: List[Tuple[fx.Node, int, _OffloadInfo]] = [] 

582 for node in list(graph_module.graph.nodes): 

583 if _is_backward(node): 

584 continue 

585 if not _in_attention_scope(node, roots) or not _can_offload(node, storage): 

586 continue 

587 info = _collect_offload_info(node, order) 

588 if info is None: 

589 continue 

590 nbytes = _tensor_bytes(node.meta["val"]) 

591 if nbytes is None: 

592 continue 

593 candidates.append((node, nbytes, info)) 

594 return candidates 

595 

596 

597def _select_within_budget( 

598 candidates: List[Tuple[fx.Node, int, _OffloadInfo]], 

599 budget_bytes: float, 

600) -> Tuple[List[_OffloadInfo], int]: 

601 """Pick largest-first until the host budget is exhausted (0 means all).""" 

602 candidates.sort(key=lambda item: item[1], reverse=True) 

603 selected: List[_OffloadInfo] = [] 

604 total_bytes = 0 

605 for _, nbytes, info in candidates: 

606 if 0 < budget_bytes < total_bytes + nbytes: 

607 continue 

608 selected.append(info) 

609 total_bytes += nbytes 

610 return selected, total_bytes 

611 

612 

613def _copy_meta(source: fx.Node) -> Dict[str, Any]: 

614 """Copy node metadata suitable for an inserted offload-chain node.""" 

615 return { 

616 key: value 

617 for key, value in source.meta.items() 

618 if key not in ("val", "recompute") 

619 } 

620 

621 

622def _insert_forward_offload( 

623 graph_module: fx.GraphModule, 

624 info: _OffloadInfo, 

625) -> fx.Node: 

626 """Insert ``ao.offload`` and the deferred forward ``ao.wait_tensor``. 

627 

628 Returns: 

629 The wait node whose (CPU) result is the reload source. 

630 """ 

631 graph = graph_module.graph 

632 node = info.node 

633 value = node.meta["val"] 

634 src_meta = _copy_meta(node) 

635 

636 with graph.inserting_after(node): 

637 offload = graph.call_function(_AO_OFFLOAD, args=(node,)) 

638 offload.meta.update(src_meta) 

639 offload.meta["val"] = value.to(torch.device("cpu")) 

640 

641 if info.last_forward_use is node: 

642 anchor = offload 

643 else: 

644 anchor = info.last_forward_use 

645 

646 with graph.inserting_after(anchor): 

647 wait = graph.call_function( 

648 _AO_WAIT, 

649 args=(offload, node, info.last_use_arg), 

650 ) 

651 wait.meta.update(src_meta) 

652 wait.meta["val"] = offload.meta["val"] 

653 wait.meta["autograd_backward"] = False 

654 return wait 

655 

656 

657def _insert_backward_reload( 

658 graph_module: fx.GraphModule, 

659 info: _OffloadInfo, 

660 wait_offload: fx.Node, 

661 order: Dict[fx.Node, int], 

662) -> fx.Node: 

663 """Insert ``ao.reload`` + ``ao.wait_tensor`` and redirect backward readers.""" 

664 graph = graph_module.graph 

665 node = info.node 

666 value = node.meta["val"] 

667 src_meta = _copy_meta(node) 

668 first_bwd = info.first_bwd_consumer 

669 

670 with graph.inserting_before(first_bwd): 

671 # Candidates are contiguous, so the offloaded host tensor already carries 

672 # the runtime shape and contiguous stride. Omitting the explicit 

673 # size/stride keeps symbolic values out of the graph arguments: baking a 

674 # ``SymInt`` here would render an undefined symbol in the generated 

675 # forward (dynamic shapes). 

676 reload = graph.call_function( 

677 _AO_RELOAD, 

678 args=(wait_offload, value.device, None, None), 

679 ) 

680 reload.meta.update(src_meta) 

681 reload.meta["val"] = value 

682 reload.meta["autograd_backward"] = True 

683 

684 with graph.inserting_before(first_bwd): 

685 wait = graph.call_function(_AO_WAIT, args=(reload, wait_offload)) 

686 wait.meta.update(src_meta) 

687 wait.meta["val"] = value 

688 wait.meta["autograd_backward"] = True 

689 

690 for user in info.direct_bwd_users: 

691 user.replace_input_with(node, wait) 

692 

693 if info.replay_views: 

694 replay_map: Dict[fx.Node, fx.Node] = {node: wait} 

695 for view_node in sorted(info.replay_views, key=lambda item: order[item]): 

696 new_args = fx.map_arg( 

697 view_node.args, 

698 lambda arg: ( 

699 replay_map.get(arg, arg) if isinstance(arg, fx.Node) else arg 

700 ), 

701 ) 

702 new_kwargs = fx.map_arg( 

703 view_node.kwargs, 

704 lambda arg: ( 

705 replay_map.get(arg, arg) if isinstance(arg, fx.Node) else arg 

706 ), 

707 ) 

708 with graph.inserting_before(first_bwd): 

709 replayed = graph.call_function( 

710 view_node.target, args=new_args, kwargs=new_kwargs 

711 ) 

712 replayed.meta.update( 

713 {k: v for k, v in view_node.meta.items() if k != "recompute"} 

714 ) 

715 replayed.meta["autograd_backward"] = True 

716 replay_map[view_node] = replayed 

717 for consumed, user in info.view_bwd_redirects: 

718 replayed = replay_map.get(consumed) 

719 if replayed is not None: 

720 user.replace_input_with(consumed, replayed) 

721 return reload 

722 

723 

724def _region_anchors(graph_module: fx.GraphModule, roots: frozenset) -> List[fx.Node]: 

725 """Return the first backward node of each contiguous attention region.""" 

726 anchors: List[fx.Node] = [] 

727 previous_root: Optional[str] = None 

728 for node in graph_module.graph.nodes: 

729 if ( 

730 node.op != "call_function" 

731 or not _is_backward(node) 

732 or node.target in _AO_OPS 

733 ): 

734 continue 

735 root = _attention_root_of(node, roots) 

736 if root is not None and root != previous_root: 

737 anchors.append(node) 

738 previous_root = root 

739 return anchors 

740 

741 

742def _anchor_positions(anchors: List[fx.Node], roots: frozenset) -> Dict[str, List[int]]: 

743 """Map each attention root to its ordered region-anchor indices.""" 

744 positions: Dict[str, List[int]] = {} 

745 for index, anchor in enumerate(anchors): 

746 root = _attention_root_of(anchor, roots) 

747 if root is not None: 

748 positions.setdefault(root, []).append(index) 

749 return positions 

750 

751 

752def _reload_nodes(graph: fx.Graph) -> List[fx.Node]: 

753 """Return every backward ``ao.reload`` node in graph order.""" 

754 return [ 

755 node 

756 for node in graph.nodes 

757 if node.op == "call_function" 

758 and node.target is _AO_RELOAD 

759 and _is_backward(node) 

760 ] 

761 

762 

763def _move_reload( 

764 reload_node: fx.Node, 

765 roots: frozenset, 

766 anchors: List[fx.Node], 

767 anchor_positions: Dict[str, List[int]], 

768 order: Dict[fx.Node, int], 

769 n_regions: int, 

770) -> bool: 

771 """Move one reload earlier by ``n_regions`` regions when possible.""" 

772 root = _attention_root_of(reload_node, roots) 

773 positions = anchor_positions.get(root) if root is not None else None 

774 if not positions: 

775 return False 

776 reload_index = order[reload_node] 

777 # Latest occurrence of this root at or before the reload is its own 

778 # region when a module is reused across blocks; otherwise fall back to 

779 # the first occurrence (reload inserted before the region's first node). 

780 earlier = [pos for pos in positions if order[anchors[pos]] <= reload_index] 

781 region_index = earlier[-1] if earlier else positions[0] 

782 anchor = anchors[max(0, region_index - n_regions)] 

783 if order[anchor] >= reload_index: 

784 return False 

785 anchor.prepend(reload_node) 

786 return True 

787 

788 

789def _prefetch_reloads( 

790 graph_module: fx.GraphModule, 

791 n_regions: int, 

792 roots: frozenset, 

793) -> int: 

794 """Move backward reloads ``n_regions`` attention regions earlier. 

795 

796 Regions are located by graph position rather than by root name so a 

797 ``self_attn`` module that appears in several blocks (same FQN, or tied 

798 weights) still maps each reload to the region it actually belongs to. 

799 """ 

800 if n_regions <= 0 or not roots: 

801 return 0 

802 

803 order = {node: index for index, node in enumerate(graph_module.graph.nodes)} 

804 

805 # First backward node of each contiguous attention region. Tracking the 

806 # previous node's root (including ``None`` gaps) re-opens a region when the 

807 # same root resumes after out-of-scope backward work. 

808 anchors = _region_anchors(graph_module, roots) 

809 if not anchors: 

810 return 0 

811 

812 anchor_positions = _anchor_positions(anchors, roots) 

813 moved = 0 

814 for reload_node in _reload_nodes(graph_module.graph): 

815 if _move_reload( 

816 reload_node, roots, anchors, anchor_positions, order, n_regions 

817 ): 

818 moved += 1 

819 return moved 

820 

821 

822def _graph_shape_env(graph_module: fx.GraphModule): 

823 """Return the ``ShapeEnv`` backing the graph's fake tensors, if any.""" 

824 for node in graph_module.graph.nodes: