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