Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / activation_checkpoint.py: 86%
456 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""Activation checkpointing helpers for distributed model components."""
16# pylint: disable=forbidden-backend-import
18import fnmatch
19import logging
20import weakref
21from collections.abc import Callable
22from dataclasses import dataclass
23from typing import Any, Optional
25import torch
26from torch import nn
28from hyper_parallel.core.activation_memory.api import (
29 CheckpointPolicy,
30 checkpoint_exclude_wrapper,
31 checkpoint_wrapper,
32 create_selective_checkpoint_contexts,
33 ignore_sac_ops as _ignore_sac_ops,
34)
35from hyper_parallel.core.activation_memory.swap import (
36 SwapManager,
37 _teardown_wired_swap_layers,
38)
39from hyper_parallel.models.adapter_spec import RecomputePolicy
40from hyper_parallel.models.registry import get_model_adapter
42logger = logging.getLogger(__name__)
45def _resolve_torch_op(dotted_path: str):
46 """Resolve a torch operator to its default overload when available."""
47 if dotted_path.startswith("ops."):
48 dotted_path = dotted_path[4:]
49 if dotted_path.count(".") == 1:
50 dotted_path = f"{dotted_path}.default"
51 obj = torch.ops
52 try:
53 for part in dotted_path.split("."):
54 obj = getattr(obj, part)
55 except AttributeError:
56 return None
57 return obj
60def _resolve_op_attr(root: object, dotted_path: str):
61 """Resolve an optional operator outside the regular ``torch.ops`` tree."""
62 obj = root
63 try:
64 for part in dotted_path.split("."):
65 obj = getattr(obj, part)
66 except AttributeError:
67 return None
68 return obj
71def _existing_ops(*ops):
72 """Return the available operators, omitting optional torch operators."""
73 return frozenset(op for op in ops if op is not None)
76# Matmul operators alternate between saving and recomputing their outputs. The
77# counter is scoped to one checkpoint region by ``make_selective_checkpoint_context_fn``.
78_SELECTIVE_AC_MATMUL_OPS = _existing_ops(
79 _resolve_torch_op("aten.matmul"),
80 _resolve_torch_op("aten.mm"),
81 _resolve_torch_op("aten.linear"),
82 _resolve_torch_op("aten._grouped_mm"),
83 _resolve_torch_op("aten._scaled_grouped_mm"),
84)
86# Some model implementations mutate these operator outputs in-place. Keeping
87# the output by reference would either trip SAC's version check or replay the
88# mutation on an already-mutated tensor during recomputation.
89_SELECTIVE_AC_FORCE_RECOMPUTE_OPS = _existing_ops(
90 _resolve_torch_op("aten.topk"),
91)
94def _default_compute_intensive_ops() -> tuple:
95 """Get PyTorch's compute-intensive operator list when the private API exists."""
96 try:
97 # This private PyTorch API is unavailable on some supported versions.
98 from torch._functorch.partitioners import get_default_op_list # pylint: disable=import-outside-toplevel
100 return tuple(op.default for op in get_default_op_list().compute_intensive_ops)
101 except (ImportError, AttributeError, RuntimeError):
102 return ()
105def _ffpa_forward_ops() -> tuple:
106 """Resolve optional FFPA forward operators after their extension registers."""
107 try:
108 # The optional extension registers its operators only when imported.
109 import ffpa_attn.cute # pylint: disable=C0415, W0611
110 except (ImportError, OSError, RuntimeError):
111 return ()
112 return (
113 _resolve_op_attr(torch.ops, "ffpa_attn._fwd_cute.default"),
114 _resolve_op_attr(torch.ops, "ffpa_attn._varlen_fwd_cute.default"),
115 )
118_SELECTIVE_AC_COMPUTE_OP_NAMES = (
119 "aten.mm",
120 "aten.addmm",
121 "aten.bmm",
122 "aten.linear",
123 "aten._scaled_mm",
124 "aten._scaled_dot_product_cudnn_attention",
125 "aten._scaled_dot_product_efficient_attention",
126 "aten._scaled_dot_product_flash_attention",
127 "aten._scaled_dot_product_flash_attention_for_cpu",
128 "aten._scaled_dot_product_fused_attention_overrideable",
129 "aten.scaled_dot_product_attention",
130 "npu.npu_fusion_attention_v3",
131 "aten._flex_attention",
132 "aten.topk",
133 "aten.max",
134)
136_SELECTIVE_AC_COMM_OP_NAMES = (
137 "aten.all_to_all_single",
138 "aten.reduce_scatter_tensor",
139 "_c10d_functional.all_to_all_single",
140 "_c10d_functional.reduce_scatter_tensor",
141 "c10d.allreduce_",
142)
145def _build_selective_ac_must_save_ops():
146 """Build the expensive/communication operator set for selective AC."""
147 save_ops = set(_default_compute_intensive_ops())
148 compute_ops = _existing_ops(
149 *(_resolve_torch_op(name) for name in _SELECTIVE_AC_COMPUTE_OP_NAMES),
150 _resolve_op_attr(torch, "_higher_order_ops.flex_attention"),
151 _resolve_op_attr(torch, "_higher_order_ops.inductor_compiled_code"),
152 _resolve_op_attr(torch.ops, "torch_attn._varlen_attn.default"),
153 *_ffpa_forward_ops(),
154 )
155 comm_ops = _existing_ops(
156 *(_resolve_torch_op(name) for name in _SELECTIVE_AC_COMM_OP_NAMES),
157 _resolve_op_attr(torch.ops, "deepep.dispatch.default"),
158 _resolve_op_attr(torch.ops, "deepep.combine.default"),
159 _resolve_op_attr(torch.ops, "hybridep.dispatch.default"),
160 _resolve_op_attr(torch.ops, "hybridep.combine.default"),
161 )
162 save_ops.update(compute_ops)
163 save_ops.update(comm_ops)
164 save_ops.difference_update(_SELECTIVE_AC_FORCE_RECOMPUTE_OPS)
165 return frozenset(save_ops)
168_SELECTIVE_AC_MUST_SAVE_OPS = _build_selective_ac_must_save_ops()
171@dataclass(frozen=True)
172class _TransformerBlockInfo:
173 """A transformer block discovered below an HF checkpointing owner.
175 ``parent`` is the registered repeated-block container (for example, a
176 ``ModuleList``), while ``child_name`` is the actual registered key. Keeping
177 both lets callers replace a block without assuming that the key is a
178 contiguous integer or that the container is a specific PyTorch class.
179 """
181 fqn: str
182 module: nn.Module
183 parent: nn.Module
184 child_name: str
185 container_fqn: str
188@dataclass(frozen=True)
189class _LayerContainerInfo:
190 """One repeated-block container and the blocks selected from it."""
192 container: nn.Module
193 path: str
194 blocks: tuple[_TransformerBlockInfo, ...]
197def _get_checkpoint_wrapped_module(module: nn.Module) -> Optional[nn.Module]:
198 """Return the module directly held by a supported checkpoint wrapper."""
199 for attr_name in ("_wrapped_module", "_checkpoint_wrapped_module"):
200 wrapped_module = getattr(module, attr_name, None)
201 if isinstance(wrapped_module, nn.Module):
202 return wrapped_module
203 return None
206def _find_transformer_block_modules(
207 model: nn.Module,
208) -> tuple[list[_TransformerBlockInfo], set[int]]:
209 """Find transformer blocks below modules with HF checkpointing support.
211 HuggingFace model components expose ``gradient_checkpointing`` on modules
212 that own a repeated block container. The marker is the structural
213 contract; no model class name or conventional ``layers`` path is needed.
214 Each direct child of a marked repeated container is returned with its
215 registered name so callers can safely replace it in-place.
217 Args:
218 model: Model whose transformer blocks should be located.
220 Returns:
221 A list of block metadata and the IDs of blocks already selected. The
222 ID set is useful to callers that need to add other wrap targets while
223 avoiding duplicate modules.
224 """
225 block_infos = []
226 discovered_block_ids = set()
227 for owner_fqn, owner in model.named_modules():
228 # Match FSDP2's behavior for nested checkpointing owners: once an
229 # owner has itself been selected as a block, do not scan inside it.
230 if id(owner) in discovered_block_ids:
231 continue
232 if not hasattr(owner, "gradient_checkpointing"):
233 continue
235 for container_name, container in owner.named_children():
236 children = list(container.named_children())
237 if not children:
238 continue
239 container_fqn = (
240 f"{owner_fqn}.{container_name}" if owner_fqn else container_name
241 )
242 for block_name, block in children:
243 if id(block) in discovered_block_ids:
244 continue
245 discovered_block_ids.add(id(block))
246 wrapped_module = _get_checkpoint_wrapped_module(block)
247 if wrapped_module is not None:
248 # A checkpoint wrapper may proxy attributes from its inner
249 # block, including ``gradient_checkpointing``. Treat the
250 # wrapper and inner module as one logical transformer block
251 # during the rest of the module-tree traversal.
252 discovered_block_ids.add(id(wrapped_module))
253 block_fqn = f"{container_fqn}.{block_name}"
254 block_infos.append(
255 _TransformerBlockInfo(
256 fqn=block_fqn,
257 module=block,
258 parent=container,
259 child_name=block_name,
260 container_fqn=container_fqn,
261 )
262 )
263 return block_infos, discovered_block_ids
266def ignore_sac_ops(ops: list[object | None]) -> None:
267 """Exclude available runtime operators from selective-AC replay accounting.
269 Args:
270 ops: Backend operators to ignore. ``None`` entries represent optional
271 operators that are unavailable in the installed PyTorch version.
272 """
273 _ignore_sac_ops(ops)
276def ensure_profiler_ops_sac_ignored() -> None:
277 """Keep profiler record-function operators out of selective-AC replay.
279 FSDP hooks run under ``record_function`` and may execute a different number
280 of profiler range operators in the original forward and recomputation. The
281 range operators carry no activations, so excluding them only removes them
282 from replay accounting while preserving their execution.
283 """
284 profiler_ops = getattr(torch.ops, "profiler", None)
285 if profiler_ops is None:
286 return
288 ops_to_ignore = []
289 for packet_name in (
290 "_record_function_enter",
291 "_record_function_enter_new",
292 "_record_function_exit",
293 ):
294 packet = getattr(profiler_ops, packet_name, None)
295 if packet is None:
296 continue
297 for overload_name in packet.overloads():
298 ops_to_ignore.append(getattr(packet, overload_name))
299 ignore_sac_ops(ops_to_ignore)
302_FSDP_SAC_IGNORED_OP_NAMES = (
303 "fsdp.all_gather_copy_in",
304 "fsdp.split_with_sizes_copy",
305 "fsdp.chunk_cat",
306 "fsdp.copy_",
307 "c10d._allgather_base_",
308 "aten.empty.memory_format",
309 "aten.empty_like",
310 "aten.view",
311)
313_FSDP_SAC_IGNORED_OPS = [_resolve_torch_op(name) for name in _FSDP_SAC_IGNORED_OP_NAMES]
316def ensure_fsdp_ops_sac_ignored() -> None:
317 """Keep FSDP parameter-lifecycle operators out of selective-AC replay.
319 Forward prefetch may unshard parameters before a checkpoint region, while
320 recomputation may need to unshard them inside that region. These allocation,
321 copy and collective operators manage parameters rather than model
322 activations, and therefore must not be matched against the forward replay.
323 """
324 ignore_sac_ops(_FSDP_SAC_IGNORED_OPS)
327def compile_selective_checkpoint_policy(
328 ctx: Any,
329 func: Any,
330 *args: Any,
331 **kwargs: Any,
332) -> CheckpointPolicy:
333 """Choose a stateless selective-checkpoint policy for compile mode.
335 Args:
336 ctx: Checkpoint context supplied by the native selective-checkpoint API.
337 func: The operator being traced.
338 *args: Operator arguments (unused).
339 **kwargs: Operator keyword arguments (unused).
341 Returns:
342 The checkpoint policy for ``func``.
343 """
344 del ctx, args, kwargs
345 if func in _SELECTIVE_AC_FORCE_RECOMPUTE_OPS:
346 return CheckpointPolicy.MUST_RECOMPUTE
347 if func in _SELECTIVE_AC_MATMUL_OPS:
348 return CheckpointPolicy.MUST_SAVE
349 if func in _SELECTIVE_AC_MUST_SAVE_OPS:
350 return CheckpointPolicy.MUST_SAVE
351 return CheckpointPolicy.MUST_RECOMPUTE
354def _make_selective_checkpoint_policy_fn() -> Callable:
355 """Create an isolated eager selective activation checkpointing policy."""
356 matmul_counts = {False: 0, True: 0}
358 def selective_checkpointing_policy(
359 ctx: Any,
360 func: Any,
361 *args: Any,
362 **kwargs: Any,
363 ) -> CheckpointPolicy:
364 """Decide whether ``func``'s output is saved or recomputed.
366 Follows the selective-activation-checkpointing policy contract: matmuls
367 alternate between save and recompute, expensive/communication ops are
368 always saved, and everything else is recomputed.
370 Args:
371 ctx: Checkpoint context carrying the ``is_recompute`` phase flag.
372 func: The operator being traced.
373 *args: Operator arguments (unused).
374 **kwargs: Operator keyword arguments (unused).
376 Returns:
377 The checkpoint policy for ``func``.
378 """
379 del args, kwargs
380 if func in _SELECTIVE_AC_FORCE_RECOMPUTE_OPS:
381 return CheckpointPolicy.MUST_RECOMPUTE
382 if func in _SELECTIVE_AC_MATMUL_OPS:
383 matmul_count = matmul_counts.get(ctx.is_recompute, 0) + 1
384 matmul_counts[ctx.is_recompute] = matmul_count
385 if matmul_count % 2:
386 return CheckpointPolicy.MUST_SAVE
387 return CheckpointPolicy.MUST_RECOMPUTE
388 if func in _SELECTIVE_AC_MUST_SAVE_OPS:
389 return CheckpointPolicy.MUST_SAVE
390 return CheckpointPolicy.MUST_RECOMPUTE
392 return selective_checkpointing_policy
395def make_selective_checkpoint_context_fn(
396 allow_cache_entry_mutation: bool = False,
397) -> Callable[[], tuple[object, object]]:
398 """Create a per-checkpoint-region selective activation policy context.
400 Expensive operations are saved, ordinary operations are recomputed, and
401 matmul operations alternate between the two decisions. A new counter is
402 created every time the returned factory is invoked, matching the
403 ``context_fn`` contract of non-reentrant checkpointing.
405 Args:
406 allow_cache_entry_mutation: Whether to accept deterministic in-place
407 rewrites of cached tensors instead of raising.
409 Returns:
410 A no-argument factory that creates the forward and recompute contexts.
411 """
412 ensure_profiler_ops_sac_ignored()
413 ensure_fsdp_ops_sac_ignored()
415 def selective_checkpoint_context_fn() -> tuple[object, object]:
416 """Create a fresh pair of forward/recompute selective-AC contexts.
418 Returns:
419 The ``(forward_context, recompute_context)`` pair expected by the
420 non-reentrant checkpointing ``context_fn`` contract.
421 """
422 return create_selective_checkpoint_contexts(
423 _make_selective_checkpoint_policy_fn(),
424 allow_cache_entry_mutation=allow_cache_entry_mutation,
425 )
427 return selective_checkpoint_context_fn
430def _find_transformer_layer_container_infos(
431 model: nn.Module,
432) -> list[_LayerContainerInfo]:
433 """Find repeated block containers without model-specific path rules.
435 Containers are derived from :func:`_find_transformer_block_modules`, so a
436 model can freely name its towers and layer attributes. A container is
437 included once even when blocks are shared or the module tree exposes an
438 alias to it.
440 Args:
441 model: Model whose transformer layer containers should be located.
443 Returns:
444 A list of discovered containers, ordered by the model's registration
445 order. An empty list means no module advertises HF checkpointing
446 support for a repeated block container.
447 """
448 block_infos, _ = _find_transformer_block_modules(model)
449 containers = []
450 blocks_by_container = {}
451 seen_container_ids = set()
452 for block_info in block_infos:
453 container_id = id(block_info.parent)
454 blocks_by_container.setdefault(container_id, []).append(block_info)
456 for block_info in block_infos:
457 container_id = id(block_info.parent)
458 if container_id in seen_container_ids:
459 continue
460 seen_container_ids.add(container_id)
461 containers.append(
462 _LayerContainerInfo(
463 container=block_info.parent,
464 path=block_info.container_fqn,
465 blocks=tuple(blocks_by_container[container_id]),
466 )
467 )
468 return containers
471def _flatten_layer_container_infos(
472 containers: list[_LayerContainerInfo],
473) -> list[nn.Module]:
474 """Expand discovered containers into individual transformer blocks."""
475 return [
476 block_info.module
477 for container in containers
478 for block_info in container.blocks
479 ]
482def _should_use_hf_native_gradient_checkpointing(
483 model: nn.Module,
484 layers: list[nn.Module],
485 *,
486 enable_compile: bool = False,
487) -> bool:
488 """Return whether full checkpointing can use HuggingFace's native API."""
489 if enable_compile:
490 return False
492 if not layers or any(
493 not any(parameter.requires_grad for parameter in layer.parameters())
494 for layer in layers
495 ):
496 return False
498 try:
499 # Guarded import: transformers is optional and older versions may not
500 # expose GradientCheckpointingLayer.
501 from transformers.modeling_layers import GradientCheckpointingLayer # pylint: disable=import-outside-toplevel
502 except ImportError:
503 return False
505 return (
506 all(isinstance(layer, GradientCheckpointingLayer) for layer in layers)
507 and getattr(model, "supports_gradient_checkpointing", False)
508 and hasattr(model, "gradient_checkpointing_enable")
509 )
512def _wrap_layer_containers(
513 containers: list[_LayerContainerInfo],
514 wrapper: Callable,
515 *,
516 context_fn: Optional[Callable[[], tuple[object, object]]] = None,
517) -> int:
518 """Wrap every layer in the discovered containers and return the count."""
519 wrapped_count = 0
520 for container_info in containers:
521 for block_info in container_info.blocks:
522 checkpoint_kwargs = {} if context_fn is None else {"context_fn": context_fn}
523 current_block = getattr(block_info.parent, block_info.child_name, None)
524 if not isinstance(current_block, nn.Module) or _is_checkpoint_wrapped(current_block):
525 continue
526 setattr(
527 block_info.parent,
528 block_info.child_name,
529 wrapper(current_block, **checkpoint_kwargs),
530 )
531 wrapped_count += 1
532 return wrapped_count
535def _is_checkpoint_wrapped(module: nn.Module) -> bool:
536 """Return whether a module is already wrapped by a supported checkpoint wrapper."""
537 return hasattr(module, "_wrapped_module") or hasattr(
538 module,
539 "_checkpoint_wrapped_module",
540 )
543def _warn_if_nothing_wrapped(
544 wrapped_count: int,
545 activation_checkpoint: str,
546 containers: list[_LayerContainerInfo],
547) -> None:
548 """Warn when a checkpointing request wrapped no module at all.
550 Every wrapping strategy either probes for known submodule names or skips
551 modules that are already wrapped, so a zero count silently leaves the model
552 running fully eager and memory-unbounded.
553 """
554 if wrapped_count != 0 or not containers:
555 return
556 layer_count = sum(len(container.blocks) for container in containers)
557 logger.warning(
558 "%s activation checkpointing wrapped no module on %d layer(s) in %s; the "
559 "model is running without activation checkpointing. Expected submodule "
560 "names may not match this architecture, or the layers are already wrapped.",
561 activation_checkpoint.capitalize(),
562 layer_count,
563 ", ".join(container.path for container in containers),
564 )
567def _find_checkpoint_wrappers(module: nn.Module, prefix: str = "") -> dict[str, nn.Module]:
568 """Find outermost checkpoint wrappers by their relative module paths."""
569 if _is_checkpoint_wrapped(module):
570 return {prefix: module}
572 wrappers = {}
573 for child_name, child in module.named_children():
574 child_path = f"{prefix}.{child_name}" if prefix else child_name
575 wrappers.update(_find_checkpoint_wrappers(child, child_path))
576 return wrappers
579def _register_forward_prefetch_layers(containers: list[_LayerContainerInfo]) -> None:
580 """Register swap prefetch chains within each repeated-block container."""
581 swap_manager = SwapManager()
582 for container_info in containers:
583 wrapper_chains = {}
584 for block_info in container_info.blocks:
585 current_block = getattr(block_info.parent, block_info.child_name, None)
586 if not isinstance(current_block, nn.Module):
587 continue
588 for relative_path, wrapper in _find_checkpoint_wrappers(current_block).items():
589 wrapper_chains.setdefault(relative_path, []).append(wrapper)
591 wired_modules: list[nn.Module] = []
592 for wrappers in wrapper_chains.values():
593 for current_wrapper, next_wrapper in zip(wrappers, wrappers[1:]):
594 swap_manager.set_forward_prefetch_layer(current_wrapper, next_wrapper)
595 wired_modules.extend((current_wrapper, next_wrapper))
596 wired_modules = list(dict.fromkeys(wired_modules)) # dedupe, keep order
598 # swap layers registered above live in the process-wide SwapManager
599 # singleton and are never torn down otherwise, so a long-lived process
600 # that rebuilds or discards models leaks SwapGroup entries and hook
601 # handles. Tear them down when the container is collected. The
602 # container must not be a finalizer argument (that would keep it alive);
603 # only the wired child modules are captured, and the callback receives
604 # them directly because the container is unreachable at that point.
605 if wired_modules:
606 weakref.finalize(
607 container_info.container,
608 _teardown_wired_swap_layers,
609 wired_modules,
610 )
613def _wrap_first_existing_attr(
614 module: nn.Module,
615 attr_names: tuple[str, ...],
616 wrapper: Callable,
617 *,
618 skip: bool = False,
619) -> int:
620 """Checkpoint-wrap the first registered child matching ``attr_names``."""
621 if skip:
622 return 0
624 for attr in attr_names:
625 child = getattr(module, attr, None)
626 if not isinstance(child, nn.Module):
627 continue
628 child_name = next(
629 (
630 name
631 for name, registered_child in module.named_children()
632 if registered_child is child
633 ),
634 None,
635 )
636 if child_name is None:
637 continue
638 if _is_checkpoint_wrapped(child):
639 return 0
640 setattr(module, child_name, wrapper(child))
641 return 1
642 return 0
645def _eager_checkpoint_kwargs(swap_inputs: bool) -> dict[str, Any]:
646 """Build the eager-mode checkpoint keyword arguments.
648 Compile mode drops ``swap_inputs`` entirely, but eager checkpointing accepts
649 it in both states: ``False`` keeps ordinary recomputation, while ``True``
650 offloads the checkpoint inputs to host memory.
651 """
652 return {"swap_inputs": swap_inputs}
655def apply_submodule_checkpointing(
656 layers: list[nn.Module],
657 has_kv_sharing: bool,
658 enable_compile: bool = False,
659 swap_inputs: bool = False,
660) -> int:
661 """Apply full activation checkpointing to transformer-layer submodules.
663 Args:
664 layers: Transformer layers whose selected submodules should be wrapped.
665 has_kv_sharing: Whether attention submodules must remain outside the
666 recomputation regions.
667 enable_compile: Whether the wrapped regions will be compiled.
668 swap_inputs: Whether checkpoint inputs should be offloaded in eager
669 execution. This is ignored in compile mode.
670 """
671 checkpoint_kwargs = {} if enable_compile else _eager_checkpoint_kwargs(swap_inputs)
673 def submodule_checkpoint_wrapper(module: nn.Module) -> nn.Module:
674 """Wrap one selected layer submodule with checkpointing."""
675 return checkpoint_wrapper(module, **checkpoint_kwargs)
677 wrapped_count = 0
678 for layer in layers:
679 wrapped_count += _wrap_first_existing_attr(
680 layer,
681 ("mlp", "feed_forward", "ffn"),
682 submodule_checkpoint_wrapper,
683 )
684 wrapped_count += _wrap_first_existing_attr(
685 layer,
686 ("self_attn", "attention", "attn", "linear_attn"),
687 submodule_checkpoint_wrapper,
688 skip=has_kv_sharing,
689 )
690 wrapped_count += _wrap_first_existing_attr(
691 layer,
692 ("input_layernorm", "attention_norm", "layer_norm1", "norm1"),
693 submodule_checkpoint_wrapper,
694 )
695 wrapped_count += _wrap_first_existing_attr(
696 layer,
697 ("post_attention_layernorm", "ffn_norm", "layer_norm2", "norm2"),
698 submodule_checkpoint_wrapper,
699 )
700 for attr in (
701 "mlp_moe_gen",
702 "input_layernorm_moe_gen",
703 "post_attention_layernorm_moe_gen",
704 ):
705 child = getattr(layer, attr, None)
706 if isinstance(child, nn.Module) and not _is_checkpoint_wrapped(child):
707 setattr(layer, attr, submodule_checkpoint_wrapper(child))
708 wrapped_count += 1
710 logger.info(
711 "Applied submodule activation checkpointing to %d layer(s), wrapping %d submodule(s)",
712 len(layers),
713 wrapped_count,
714 )
715 return wrapped_count
718def _detect_kv_sharing_and_maybe_disable_cache(model: nn.Module) -> bool:
719 """Detect cross-layer KV sharing and disable ordinary model caches."""
720 config = getattr(model, "config", None)
721 text_config = getattr(config, "text_config", None) or config
722 has_kv_sharing = getattr(text_config, "num_kv_shared_layers", 0) > 0
723 if has_kv_sharing or config is None:
724 return has_kv_sharing
726 sub_config_names = getattr(type(config), "sub_configs", None) or {
727 "text_config": None,
728 }
729 sub_configs = (getattr(config, name, None) for name in sub_config_names)
730 for sub_config in (config, *sub_configs):
731 if sub_config is None or (
732 sub_config is not config and not hasattr(sub_config, "use_cache")
733 ):
734 continue
735 if getattr(sub_config, "use_cache", None) is not False:
736 _try_disable_use_cache(sub_config)
737 return False
740def _try_disable_use_cache(sub_config: Any) -> None:
741 """Best-effort disable of ``use_cache`` on one config object."""
742 try:
743 setattr(sub_config, "use_cache", False)
744 except Exception: # pylint: disable=broad-exception-caught
745 # Configuration objects may reject assignment with custom errors.
746 pass
749def _validate_activation_checkpoint_config(
750 activation_checkpoint: Optional[str],
751 swap_inputs: bool,
752 enable_compile: bool,
753) -> None:
754 """Validate the activation-checkpoint options and warn about ignored ones.
756 Args:
757 activation_checkpoint: Requested recomputation mode.
758 swap_inputs: Whether checkpoint inputs should be offloaded.
759 enable_compile: Whether wrapped regions will be compiled.
761 Raises:
762 ValueError: The mode or ``swap_inputs`` has an unsupported value.
763 """
764 if activation_checkpoint not in ("full", "selective"):
765 raise ValueError(
766 "activation_checkpoint.mode must be 'full' or 'selective', but got "
767 f"{activation_checkpoint!r}"
768 )
769 if not isinstance(swap_inputs, bool):
770 raise ValueError(
771 "activation_checkpoint.swap_inputs must be bool, but got "
772 f"{type(swap_inputs).__name__}"
773 )
775 if swap_inputs and enable_compile:
776 logger.warning(
777 "activation_checkpoint.swap_inputs is not supported with torch.compile; "
778 "input swapping will be disabled."
779 )
782def _adapter_recompute_spec(model: nn.Module):
783 """Resolve a model-owned recompute declaration without family branches."""
784 config = getattr(model, "config", None)
785 identities = [getattr(config, "model_type", None)]
786 identities.extend(getattr(config, "architectures", None) or ())
787 adapter_spec = None
788 for identity in identities:
789 if not identity:
790 continue
791 adapter_spec = get_model_adapter(identity)
792 if adapter_spec is not None:
793 break
794 provider = getattr(adapter_spec, "recompute", None)
795 if provider is None or not callable(provider):
796 raise ValueError(
797 "activation_checkpoint.selection.source=model_adapter_safe_regions "
798 "requires ModelAdapterSpec.recompute"
799 )
800 recompute_policy = provider() # pylint: disable=not-callable
801 if not isinstance(recompute_policy, RecomputePolicy):
802 raise TypeError("ModelAdapterSpec.recompute must return RecomputePolicy")
803 return recompute_policy
806def _validate_selected_layer_indices(
807 layer_indices: Any,
808 blocks: list,
809) -> None:
810 """Validate an explicit list of layer indices against the blocks."""
811 if not isinstance(layer_indices, (list, tuple)) or not layer_indices:
812 raise ValueError(
813 "model-adapter safe recompute layer_indices must be a non-empty sequence"
814 )
815 if any(
816 isinstance(index, bool) or not isinstance(index, int) or index < 0
817 for index in layer_indices
818 ):
819 raise ValueError(
820 "model-adapter safe recompute layer_indices must contain only "
821 "non-negative integers"
822 )
823 if len(set(layer_indices)) != len(layer_indices):
824 raise ValueError(
825 "model-adapter safe recompute layer_indices must not contain duplicates"
826 )
827 if max(layer_indices) >= len(blocks):
828 raise ValueError(
829 "model-adapter safe recompute layer index "
830 f"{max(layer_indices)} exceeds the maximum index {len(blocks) - 1}"
831 )
834def _validate_selected_layer_count(
835 layer_count: Any,
836 blocks: list,
837) -> None:
838 """Validate a leading layer count against the blocks."""
839 if isinstance(layer_count, bool) or not isinstance(layer_count, int) or layer_count < 0:
840 raise ValueError(
841 "model-adapter safe recompute layer_count must be a non-negative integer"
842 )
843 if layer_count > len(blocks):
844 raise ValueError(
845 f"adapter recompute layer_count {layer_count} exceeds discovered layers {len(blocks)}"
846 )
849def _selected_checkpoint_layer_fqns(
850 containers: list[_LayerContainerInfo],
851 selection: Any,
852) -> set[str]:
853 """Select a leading layer count or exact layer indices."""
854 blocks = [block for container in containers for block in container.blocks]
855 layer_count = getattr(selection, "layer_count", None)
856 layer_indices = getattr(selection, "layer_indices", None)
857 if (layer_count is None) == (layer_indices is None):
858 raise ValueError(
859 "model-adapter safe recompute requires exactly one of "
860 "selection.layer_count or selection.layer_indices"
861 )
862 if layer_indices is not None:
863 _validate_selected_layer_indices(layer_indices, blocks)
864 return {blocks[index].fqn for index in layer_indices}
866 _validate_selected_layer_count(layer_count, blocks)
867 if layer_count == 0:
868 return set()
869 return {blocks[index].fqn for index in range(layer_count)}
872def _select_layer_region_roots(
873 module_by_fqn: dict,
874 selected_layer_fqns: set[str],
875 recompute_spec: Any,
876) -> list:
877 """Return the outermost checkpoint region roots inside the selected layers."""
878 matched = [
879 (module_fqn, module)
880 for module_fqn, module in module_by_fqn.items()
881 if module_fqn
882 and any(
883 fnmatch.fnmatchcase(module_fqn, pattern)
884 for pattern in recompute_spec.region_patterns
885 )
886 ]
887 outermost = [
888 (module_fqn, module)
889 for module_fqn, module in matched
890 if not any(
891 module_fqn != other_fqn and module_fqn.startswith(f"{other_fqn}.")
892 for other_fqn, _ in matched
893 )
894 ]
895 return [
896 (module_fqn, module)
897 for module_fqn, module in outermost
898 if any(
899 module_fqn == layer_fqn or module_fqn.startswith(f"{layer_fqn}.")
900 for layer_fqn in selected_layer_fqns
901 )
902 ]
905def _matches_exclusion(module_fqn: str, module: nn.Module, recompute_spec: Any) -> bool:
906 """Return whether one module is declared to stay outside the recompute region."""
907 if any(
908 fnmatch.fnmatchcase(module_fqn, pattern)
909 for pattern in recompute_spec.exclude_patterns
910 ):
911 return True
912 predicate = recompute_spec.exclude_predicate
913 return bool(predicate is not None and predicate(module_fqn, module))
916def _select_excluded_targets(
917 module_by_fqn: dict,
918 region_roots: list,
919 recompute_spec: Any,
920) -> list:
921 """Collect the declared descendants that must stay outside the recompute region."""
922 root_fqns = tuple(root_fqn for root_fqn, _ in region_roots)
923 selected_targets = []
924 for module_fqn, module in module_by_fqn.items():
925 if not any(
926 module_fqn == root_fqn or module_fqn.startswith(f"{root_fqn}.")
927 for root_fqn in root_fqns
928 ):
929 continue
930 if _matches_exclusion(module_fqn, module, recompute_spec):
931 selected_targets.append((module_fqn, module))
932 return selected_targets
935def _wrap_excluded_targets(module_by_fqn: dict, exclude_targets: list) -> int:
936 """Install exclusion wrappers deepest-first; return the wrapped count."""
937 wrapped_count = 0
938 for module_fqn, module in sorted(
939 exclude_targets,
940 key=lambda item: item[0].count("."),
941 reverse=True,
942 ):
943 if _is_checkpoint_wrapped(module):
944 continue
945 parent_fqn, child_name = module_fqn.rsplit(".", 1)
946 setattr(
947 module_by_fqn[parent_fqn],
948 child_name,
949 checkpoint_exclude_wrapper(module),
950 )
951 wrapped_count += 1
952 return wrapped_count
955def _validate_region_exclusions(roots: list, exclude_targets: list) -> None:
956 """Reject region shapes that would recompute an excluded module or swallow it."""
957 root_fqns = {root_fqn for root_fqn, _ in roots}
958 for module_fqn, _ in exclude_targets:
959 if module_fqn in root_fqns:
960 raise ValueError(
961 f"recompute region root {module_fqn!r} is also excluded "
962 "from the recompute"
963 )
964 root_ids = {id(module) for _, module in roots}
965 for module_fqn, module in exclude_targets:
966 if any(
967 child is not module and id(child) in root_ids
968 for child in module.modules()
969 ):
970 raise ValueError(
971 f"excluded recompute region {module_fqn!r} contains a recompute region root"
972 )
975def _apply_model_adapter_regions_checkpointing(
976 model: nn.Module,
977 containers: list[_LayerContainerInfo],
978 ac_layers: list[nn.Module],
979 selection: Any,
980 *,
981 enable_compile: bool,
982 swap_inputs: bool,
983) -> nn.Module:
984 """Wrap the adapter-declared recompute regions of the selected layers.
986 The outermost pattern match inside every selected layer becomes one
987 checkpoint region, so a layer pattern keeps a single saved copy of the layer
988 input. Declared exclusions stay outside the recompute region, which is what
989 lets a cross-layer state producer or consumer live inside a checkpointed
990 layer.
991 """
992 recompute_spec = _adapter_recompute_spec(model)
993 if not recompute_spec.region_patterns:
994 raise ValueError(
995 "activation_checkpoint.selection.source=model_adapter_safe_regions "
996 "requires RecomputePolicy.region_patterns"
997 )
998 selected_layer_fqns = _selected_checkpoint_layer_fqns(containers, selection)
999 if not selected_layer_fqns:
1000 logger.info("Model-adapter recompute selected zero layers")
1001 return model
1002 if hasattr(model, "gradient_checkpointing_disable"):
1003 model.gradient_checkpointing_disable()
1004 checkpoint_kwargs = {"swap_inputs": swap_inputs} if not enable_compile else {}
1006 module_by_fqn = dict(model.named_modules())
1007 roots = _select_layer_region_roots(module_by_fqn, selected_layer_fqns, recompute_spec)
1008 if not roots:
1009 raise ValueError(
1010 "model adapter recompute patterns matched no region in the selected layers"
1011 )
1012 excluded_count = _wrap_declared_exclusions(module_by_fqn, roots, recompute_spec)
1013 wrapped_count = _wrap_region_roots(model, module_by_fqn, roots, checkpoint_kwargs)
1014 if swap_inputs and not enable_compile:
1015 _register_forward_prefetch_layers(containers)
1016 if recompute_spec.on_applied is not None:
1017 recompute_spec.on_applied(model)
1018 logger.info(
1019 "Model-adapter recompute wrapped %d region(s) and excluded %d region(s) "
1020 "across %d/%d discovered layers",
1021 wrapped_count,
1022 excluded_count,
1023 len(selected_layer_fqns),
1024 len(ac_layers),
1025 )
1026 return model
1029def _wrap_declared_exclusions(module_by_fqn: dict, roots: list, recompute_spec: Any) -> int:
1030 """Validate the declared exclusions of the selected regions and wrap them."""
1031 exclude_targets = _select_excluded_targets(module_by_fqn, roots, recompute_spec)
1032 declared_exclusions = bool(
1033 recompute_spec.exclude_patterns or recompute_spec.exclude_predicate is not None
1034 )
1035 if declared_exclusions and not exclude_targets:
1036 raise ValueError(
1037 "model adapter recompute exclusions matched no module "
1038 "inside the selected regions"
1039 )
1040 _validate_region_exclusions(roots, exclude_targets)
1041 return _wrap_excluded_targets(module_by_fqn, exclude_targets)
1044def _wrap_region_roots(model: nn.Module, module_by_fqn: dict, roots: list, checkpoint_kwargs: dict) -> int:
1045 """Wrap every selected region root, deepest first, and return the count."""
1046 wrapped_count = 0
1047 for module_fqn, module in sorted(
1048 roots,
1049 key=lambda item: item[0].count("."),
1050 reverse=True,
1051 ):
1052 if _is_checkpoint_wrapped(module):
1053 continue
1054 parent_fqn, child_name = module_fqn.rsplit(".", 1)
1055 parent = module_by_fqn[parent_fqn] if parent_fqn else model
1056 setattr(parent, child_name, checkpoint_wrapper(module, **checkpoint_kwargs))
1057 wrapped_count += 1
1058 if wrapped_count == 0:
1059 raise ValueError("model adapter recompute matched no region to wrap")
1060 return wrapped_count
1063def _resolve_activation_checkpoint_containers(
1064 model: nn.Module,
1065) -> list[_LayerContainerInfo]:
1066 """Discover the repeated-block containers that activation checkpointing targets.
1068 Args:
1069 model: Model to inspect.
1071 Returns:
1072 The discovered containers, in registration order.
1074 Raises:
1075 ValueError: No container advertises HuggingFace checkpointing support.
1076 """
1077 containers = _find_transformer_layer_container_infos(model)
1078 if not containers:
1079 raise ValueError(
1080 f"{type(model).__name__} has no module with a 'gradient_checkpointing' "
1081 "attribute and a non-empty repeated block container"
1082 )
1083 return containers
1086def _apply_selective_checkpointing(
1087 containers: list[_LayerContainerInfo],
1088 ac_layers: list[nn.Module],
1089 has_kv_sharing: bool,
1090 *,
1091 enable_compile: bool,
1092 swap_inputs: bool,
1093 allow_cache_entry_mutation: bool = False,
1094) -> int:
1095 """Wrap every discovered layer for selective recomputation.
1097 KV-shared models fall back to submodule checkpointing so attention does not
1098 write the cache again during backward recomputation.
1100 Args:
1101 containers: The repeated-block containers to wrap.
1102 ac_layers: Every block selected from ``containers``.
1103 has_kv_sharing: Whether attention submodules must stay outside the
1104 recomputation regions.
1105 enable_compile: Whether the wrapped regions will be compiled.
1106 swap_inputs: Whether checkpoint inputs should be offloaded in eager
1107 execution. Ignored in compile mode.
1108 allow_cache_entry_mutation: Whether to accept deterministic in-place
1109 rewrites of cached tensors instead of raising.
1111 Returns:
1112 The number of layers, or submodules in the KV-shared fallback, wrapped.
1113 """
1114 if has_kv_sharing:
1115 logger.warning(
1116 "Selective activation checkpointing is not supported for KV-shared models; "
1117 "falling back to submodule activation checkpointing."
1118 )
1119 return apply_submodule_checkpointing(
1120 ac_layers,
1121 has_kv_sharing,
1122 enable_compile=enable_compile,
1123 swap_inputs=swap_inputs,
1124 )
1126 if enable_compile:
1127 def compile_checkpoint_wrapper(layer: nn.Module) -> nn.Module:
1128 """Wrap one layer with the compile-compatible SAC policy."""
1129 return checkpoint_wrapper(
1130 layer,
1131 policy_fn=compile_selective_checkpoint_policy,
1132 )
1134 return _wrap_layer_containers(containers, compile_checkpoint_wrapper)
1136 def eager_checkpoint_wrapper(layer: nn.Module, **checkpoint_kwargs: Any) -> nn.Module:
1137 """Wrap one layer with eager selective checkpointing."""
1138 return checkpoint_wrapper(
1139 layer,
1140 swap_inputs=swap_inputs,
1141 **checkpoint_kwargs,
1142 )
1144 return _wrap_layer_containers(
1145 containers,
1146 eager_checkpoint_wrapper,
1147 context_fn=make_selective_checkpoint_context_fn(
1148 allow_cache_entry_mutation),
1149 )
1152def _apply_full_checkpointing(
1153 model: nn.Module,
1154 containers: list[_LayerContainerInfo],
1155 ac_layers: list[nn.Module],
1156 has_kv_sharing: bool,
1157 *,
1158 enable_compile: bool,
1159 swap_inputs: bool,
1160) -> Optional[int]:
1161 """Wrap discovered layers for full recomputation.
1163 Prefers the HuggingFace-native implementation when every eligibility check
1164 passes, and otherwise uses Hyper Parallel's wrappers.
1166 Args:
1167 model: Model whose layers should be wrapped.
1168 containers: The repeated-block containers to wrap.
1169 ac_layers: Every block selected from ``containers``.
1170 has_kv_sharing: Whether attention submodules must stay outside the
1171 recomputation regions.
1172 enable_compile: Whether the wrapped regions will be compiled.
1173 swap_inputs: Whether checkpoint inputs should be offloaded in eager
1174 execution. Ignored in compile mode.
1176 Returns:
1177 The number of wrapped submodules, or ``None`` when the HuggingFace-native
1178 implementation was enabled instead, in which case no per-layer wrapping
1179 happened and swap prefetch chains must not be registered.
1180 """
1181 if not has_kv_sharing and _should_use_hf_native_gradient_checkpointing(
1182 model,
1183 ac_layers,
1184 enable_compile=enable_compile,
1185 ):
1186 if swap_inputs:
1187 logger.warning(
1188 "activation_checkpoint.swap_inputs is not supported by Hugging Face native "
1189 "gradient checkpointing for now; input swapping will be disabled."
1190 )
1191 model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True})
1192 logger.info("Using HuggingFace native gradient checkpointing for discovered layers.")
1193 return None
1195 if has_kv_sharing:
1196 return apply_submodule_checkpointing(
1197 ac_layers,
1198 has_kv_sharing,
1199 enable_compile=enable_compile,
1200 swap_inputs=swap_inputs,
1201 )
1203 def full_checkpoint_wrapper(layer: nn.Module) -> nn.Module:
1204 """Wrap one complete transformer layer for full recomputation."""
1205 if enable_compile:
1206 return checkpoint_wrapper(layer)
1207 return checkpoint_wrapper(layer, **_eager_checkpoint_kwargs(swap_inputs))
1209 return _wrap_layer_containers(containers, full_checkpoint_wrapper)
1212def _apply_activation_checkpointing(
1213 model: nn.Module,
1214 activation_checkpoint: Optional[str],
1215 enable_compile: bool = False,
1216 swap_inputs: bool = False,
1217 selection: Optional[Any] = None,
1218 allow_cache_entry_mutation: bool = False,
1219) -> nn.Module:
1220 """Apply full or selective recomputation to discovered transformer layers.
1222 Validates the requested options, discovers the repeated-block containers that
1223 advertise HuggingFace checkpointing support, then delegates to the
1224 mode-specific wrapper. Swap prefetch chains are registered afterwards for
1225 every configuration that wrapped the layers itself.
1226 """
1227 _validate_activation_checkpoint_config(
1228 activation_checkpoint,
1229 swap_inputs,
1230 enable_compile,
1231 )
1232 containers = _resolve_activation_checkpoint_containers(model)
1233 ac_layers = _flatten_layer_container_infos(containers)
1234 has_kv_sharing = _detect_kv_sharing_and_maybe_disable_cache(model)
1236 selection_source = getattr(selection, "source", "default")
1237 if selection_source == "model_adapter_safe_regions":
1238 return _apply_model_adapter_regions_checkpointing(
1239 model,
1240 containers,
1241 ac_layers,
1242 selection,
1243 enable_compile=enable_compile,
1244 swap_inputs=swap_inputs,
1245 )
1247 if hasattr(model, "gradient_checkpointing_disable"):
1248 model.gradient_checkpointing_disable()
1250 # Selective recomputation normally wraps whole layers. KV-shared models
1251 # instead use submodule checkpointing so attention does not write the cache
1252 # again during backward recomputation.
1253 if activation_checkpoint == "selective":
1254 wrapped_count = _apply_selective_checkpointing(
1255 containers,
1256 ac_layers,
1257 has_kv_sharing,
1258 enable_compile=enable_compile,
1259 swap_inputs=swap_inputs,
1260 allow_cache_entry_mutation=allow_cache_entry_mutation,
1261 )
1262 _warn_if_nothing_wrapped(wrapped_count, activation_checkpoint, containers)
1263 logger.info(
1264 "Selective activation checkpointing applied to %d layer(s) in: %s",
1265 wrapped_count,
1266 ", ".join(container.path for container in containers),
1267 )
1268 else:
1269 wrapped_count = _apply_full_checkpointing(
1270 model,
1271 containers,
1272 ac_layers,
1273 has_kv_sharing,
1274 enable_compile=enable_compile,
1275 swap_inputs=swap_inputs,
1276 )
1277 if wrapped_count is None:
1278 # The HuggingFace-native path returns early in the original flow and
1279 # therefore skips both the warning and the swap prefetch registration.
1280 return model
1281 _warn_if_nothing_wrapped(wrapped_count, activation_checkpoint, containers)
1282 logger.info(
1283 "%s activation checkpointing wrapped %d submodule(s) in: %s",
1284 activation_checkpoint.capitalize(),
1285 wrapped_count,
1286 ", ".join(container.path for container in containers),
1287 )
1289 if swap_inputs and not enable_compile:
1290 _register_forward_prefetch_layers(containers)
1291 return model