> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.distributed.activation_checkpointing Selective activation checkpointing core. TorchTitan-style selective activation checkpointing: the policy decides, per op, whether to save or recompute an activation, saving the expensive ops (attention, half of the matmuls, comm collectives) while recomputing the cheap ones. This module holds the parts of the AC implementation that do not depend on the rest of `parallelizer.py` (notably the heavy, transformers-aware `_extract_model_layers`). `parallelizer.py` imports from here -- never the other way around -- so the dependency stays one-directional and the central parallelizer file stays small. ## Module Contents ### Functions | Name | Description | | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- | | [`_build_selective_ac_save_ops`](#nemo_automodel-components-distributed-activation_checkpointing-_build_selective_ac_save_ops) | Build the set of ops whose activations are always saved under selective AC. | | [`_default_compute_intensive_ops`](#nemo_automodel-components-distributed-activation_checkpointing-_default_compute_intensive_ops) | Compute-intensive aten ops from PyTorch's partitioner, or `()` if unavailable. | | [`_disable_dynamo_lru_cache`](#nemo_automodel-components-distributed-activation_checkpointing-_disable_dynamo_lru_cache) | Best-effort disable of TorchDynamo's LRU cache for selective AC + compile. | | [`_existing_ops`](#nemo_automodel-components-distributed-activation_checkpointing-_existing_ops) | - | | [`_ffpa_forward_ops`](#nemo_automodel-components-distributed-activation_checkpointing-_ffpa_forward_ops) | FFPA forward ops (dense + varlen); import the CuTeDSL kernel so they register first, else `()`. | | [`_get_transformer_engine_attention_backend_cache`](#nemo_automodel-components-distributed-activation_checkpointing-_get_transformer_engine_attention_backend_cache) | Return Transformer Engine's attention-backend cache when available. | | [`_is_cuda_to_cpu_copy`](#nemo_automodel-components-distributed-activation_checkpointing-_is_cuda_to_cpu_copy) | - | | [`_maybe_trace_selective_ac_decision`](#nemo_automodel-components-distributed-activation_checkpointing-_maybe_trace_selective_ac_decision) | Log a selective-AC decision once per op (no-op unless tracing is enabled). | | [`_registered_child_name`](#nemo_automodel-components-distributed-activation_checkpointing-_registered_child_name) | Return the registered name for a child reached through an attribute. | | [`_replace_child_module`](#nemo_automodel-components-distributed-activation_checkpointing-_replace_child_module) | Replace `target` with `replacement` in `root`'s module tree. | | [`_resolve_op_attr`](#nemo_automodel-components-distributed-activation_checkpointing-_resolve_op_attr) | Resolve a dotted attribute path from `root`, or `None` if any part is absent. | | [`_resolve_torch_op`](#nemo_automodel-components-distributed-activation_checkpointing-_resolve_torch_op) | Resolve a dotted `torch.ops` path, defaulting to the `default` overload. | | [`_restore_sdpa_state`](#nemo_automodel-components-distributed-activation_checkpointing-_restore_sdpa_state) | Temporarily restore the SDPA callable and backend set captured during forward. | | [`_restore_transformer_engine_attention_backend_cache`](#nemo_automodel-components-distributed-activation_checkpointing-_restore_transformer_engine_attention_backend_cache) | Restore the forward-entry TE attention cache only for recomputation. | | [`_wrap_first_existing_attr`](#nemo_automodel-components-distributed-activation_checkpointing-_wrap_first_existing_attr) | Checkpoint-wrap the first matching registered child attr on `module`. | | [`apply_selective_checkpointing_to_layers`](#nemo_automodel-components-distributed-activation_checkpointing-apply_selective_checkpointing_to_layers) | Wrap whole transformer blocks with the selective-AC policy. | | [`apply_submodule_checkpointing`](#nemo_automodel-components-distributed-activation_checkpointing-apply_submodule_checkpointing) | Wrap a transformer block's sub-modules with `checkpoint_wrapper`. | | [`detect_kv_sharing_and_maybe_disable_cache`](#nemo_automodel-components-distributed-activation_checkpointing-detect_kv_sharing_and_maybe_disable_cache) | Detect KV-sharing and disable `use_cache` for non-KV-shared models. | | [`ensure_fsdp_ops_sac_ignored`](#nemo_automodel-components-distributed-activation_checkpointing-ensure_fsdp_ops_sac_ignored) | Keep FSDP2 parameter-lifecycle ops out of SAC's saved-op replay. | | [`ensure_profiler_ops_sac_ignored`](#nemo_automodel-components-distributed-activation_checkpointing-ensure_profiler_ops_sac_ignored) | Keep `torch.ops.profiler` record-function ops out of SAC's op replay. | | [`ignore_sac_ops`](#nemo_automodel-components-distributed-activation_checkpointing-ignore_sac_ops) | Add available operators to PyTorch's selective-AC ignore set. | | [`is_selective_activation_checkpointing`](#nemo_automodel-components-distributed-activation_checkpointing-is_selective_activation_checkpointing) | Return whether the config value selects selective activation checkpointing. | | [`make_selective_checkpoint_context_fn`](#nemo_automodel-components-distributed-activation_checkpointing-make_selective_checkpoint_context_fn) | Build a TorchTitan-style selective activation checkpointing context. | | [`sdpa_backend_snapshot_context_fn`](#nemo_automodel-components-distributed-activation_checkpointing-sdpa_backend_snapshot_context_fn) | Snapshot the ambient SDPA state and restore it on checkpoint recompute. | | [`transformer_engine_attention_backend_snapshot_context_fn`](#nemo_automodel-components-distributed-activation_checkpointing-transformer_engine_attention_backend_snapshot_context_fn) | Snapshot Transformer Engine's attention-backend cache for checkpoint replay. | | [`unwrap_checkpoint_wrapper`](#nemo_automodel-components-distributed-activation_checkpointing-unwrap_checkpoint_wrapper) | Return the activation-checkpointed module, or the input module if it is not wrapped. | ### Data [`SELECTIVE_AC_WRAPPER_FLAG`](#nemo_automodel-components-distributed-activation_checkpointing-SELECTIVE_AC_WRAPPER_FLAG) [`_SELECTIVE_AC_MATMUL_OPS`](#nemo_automodel-components-distributed-activation_checkpointing-_SELECTIVE_AC_MATMUL_OPS) [`_SELECTIVE_AC_MUST_SAVE_OPS`](#nemo_automodel-components-distributed-activation_checkpointing-_SELECTIVE_AC_MUST_SAVE_OPS) [`_SELECTIVE_AC_TO_COPY_OP`](#nemo_automodel-components-distributed-activation_checkpointing-_SELECTIVE_AC_TO_COPY_OP) [`_SELECTIVE_AC_TRACE`](#nemo_automodel-components-distributed-activation_checkpointing-_SELECTIVE_AC_TRACE) [`_SELECTIVE_AC_TRACE_SEEN`](#nemo_automodel-components-distributed-activation_checkpointing-_SELECTIVE_AC_TRACE_SEEN) [`_TORCH_PROFILER_SAC_IGNORE_MIN_VERSION`](#nemo_automodel-components-distributed-activation_checkpointing-_TORCH_PROFILER_SAC_IGNORE_MIN_VERSION) [`logger`](#nemo_automodel-components-distributed-activation_checkpointing-logger) ### API ```python nemo_automodel.components.distributed.activation_checkpointing._build_selective_ac_save_ops() -> frozenset ``` Build the set of ops whose activations are always saved under selective AC. The set is seeded from PyTorch's compute-intensive op list and supplemented with attention variants, low-precision/reduction ops, the compiled HOP, and communication collectives whose outputs are expensive to recompute. ```python nemo_automodel.components.distributed.activation_checkpointing._default_compute_intensive_ops() -> tuple ``` Compute-intensive aten ops from PyTorch's partitioner, or `()` if unavailable. Mirrors TorchTitan: seeding from PyTorch's own `compute_intensive_ops` list keeps the save-set in sync with upstream rather than relying on a frozen, hand-maintained list. `torch._functorch.partitioners` is a private API, so any failure falls back to the curated supplement in :func:`_build_selective_ac_save_ops`. ```python nemo_automodel.components.distributed.activation_checkpointing._disable_dynamo_lru_cache() -> None ``` Best-effort disable of TorchDynamo's LRU cache for selective AC + compile. With multiple pipeline microbatches, dynamo may compile a second graph with dynamic shapes and then select it over the static graph whose compiled-HOP output SAC cached for microbatch 0, tripping a missing-symint assertion. Selecting graphs in insertion order avoids this. Mirrors TorchTitan. The underlying API is private, so failures are swallowed. ```python nemo_automodel.components.distributed.activation_checkpointing._existing_ops( ops = () ) ``` ```python nemo_automodel.components.distributed.activation_checkpointing._ffpa_forward_ops() -> tuple ``` FFPA forward ops (dense + varlen); import the CuTeDSL kernel so they register first, else `()`. Their ops register only on that import, which otherwise lands after this save-set is frozen -- hence the eager import rather than a bare resolve. ```python nemo_automodel.components.distributed.activation_checkpointing._get_transformer_engine_attention_backend_cache() -> dict | None ``` Return Transformer Engine's attention-backend cache when available. ```python nemo_automodel.components.distributed.activation_checkpointing._is_cuda_to_cpu_copy( func, args, kwargs ) -> bool ``` ```python nemo_automodel.components.distributed.activation_checkpointing._maybe_trace_selective_ac_decision( func, decision, is_alternating: bool, is_recompute: bool ) -> None ``` Log a selective-AC decision once per op (no-op unless tracing is enabled). **Parameters:** **`func`** The op the policy was queried about. --- **`decision`** The `CheckpointPolicy` the policy returned for `func`. --- **`is_alternating`** `bool` Whether `func` is an alternating-save matmul op. --- **`is_recompute`** `bool` Whether the policy was queried during the recompute pass; decisions are only logged on the forward pass to avoid duplicates. --- ```python nemo_automodel.components.distributed.activation_checkpointing._registered_child_name( module: torch.nn.Module, attr: str, child: torch.nn.Module ) -> str | None ``` Return the registered name for a child reached through an attribute. ```python nemo_automodel.components.distributed.activation_checkpointing._replace_child_module( root: torch.nn.Module, target: torch.nn.Module, replacement: torch.nn.Module ) -> bool ``` Replace `target` with `replacement` in `root`'s module tree. ```python nemo_automodel.components.distributed.activation_checkpointing._resolve_op_attr( root: object, dotted_path: str ) ``` Resolve a dotted attribute path from `root`, or `None` if any part is absent. Used for ops that live outside `torch.ops` (higher-order ops, optional custom backends such as DeepEP/HybridEP). Missing namespaces/ops raise `AttributeError` on access, so they are swallowed and reported as `None`. ```python nemo_automodel.components.distributed.activation_checkpointing._resolve_torch_op( dotted_path: str ) ``` Resolve a dotted `torch.ops` path, defaulting to the `default` overload. ```python nemo_automodel.components.distributed.activation_checkpointing._restore_sdpa_state( sdpa: collections.abc.Callable, backends: list[torch.nn.attention.SDPBackend] ) ``` Temporarily restore the SDPA callable and backend set captured during forward. ```python nemo_automodel.components.distributed.activation_checkpointing._restore_transformer_engine_attention_backend_cache( cache: dict, captured_cache: dict, recompute_context: contextlib.AbstractContextManager ) ``` Restore the forward-entry TE attention cache only for recomputation. ```python nemo_automodel.components.distributed.activation_checkpointing._wrap_first_existing_attr( module: torch.nn.Module, attr_names: tuple[str, ...], skip: bool = False, context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None = None ) -> int ``` Checkpoint-wrap the first matching registered child attr on `module`. ```python nemo_automodel.components.distributed.activation_checkpointing.apply_selective_checkpointing_to_layers( model: torch.nn.Module, layers: typing.List[torch.nn.Module], has_kv_sharing: bool, enable_compile: bool = False ) -> None ``` Wrap whole transformer blocks with the selective-AC policy. KV-shared models cannot checkpoint attention through the `DynamicCache`, so they fall back to sub-module checkpointing. `layers` is mutated in place so callers that retain the list (e.g. for subsequent FSDP sharding) see the wrapped modules. Works without FSDP/distributed, so it is shared by the FSDP2 strategy and the single-GPU path. ```python nemo_automodel.components.distributed.activation_checkpointing.apply_submodule_checkpointing( layers: typing.List[torch.nn.Module], has_kv_sharing: bool, context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None = sdpa_backend_snapshot_conte... ) -> None ``` Wrap a transformer block's sub-modules with `checkpoint_wrapper`. This is the sub-module granularity path used both as the default (non-compile) behavior and as the fallback for selective activation checkpointing on KV-shared models, which cannot checkpoint the whole block. `self_attn` is skipped for KV-shared models: recomputing attention during backward would double-write to the `DynamicCache`, corrupting the K/V entries that later shared layers depend on. **Parameters:** **`layers`** `List[nn.Module]` Transformer decoder layers to wrap (mutated in place). --- **`has_kv_sharing`** `bool` Whether the model reuses K/V across layers via the cache. --- **`context_fn`** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None` — default: sdpa\_backend\_snapshot\_context\_fn Factory returning `(forward_ctx, recompute_ctx)` for the attention and MLP checkpoint wrappers. Defaults to restoring the forward-time SDPA state; pass `None` to disable it. Norm wrappers stay plain because they dispatch no SDPA. --- ```python nemo_automodel.components.distributed.activation_checkpointing.detect_kv_sharing_and_maybe_disable_cache( model: torch.nn.Module ) -> bool ``` Detect KV-sharing and disable `use_cache` for non-KV-shared models. Models with KV-shared layers (e.g. Gemma4 2B/4B) pass K/V from earlier layers to later layers through the `DynamicCache`; disabling the cache breaks that dependency, so `use_cache` is left untouched for them. **Returns:** `bool` Whether the model uses KV-sharing. ```python nemo_automodel.components.distributed.activation_checkpointing.ensure_fsdp_ops_sac_ignored() -> None ``` Keep FSDP2 parameter-lifecycle ops out of SAC's saved-op replay. FSDP2 can prefetch an inner module's parameters before a checkpointed forward, while backward-time recomputation must unshard them from inside the checkpointed region. Its copy-in/copy-out and c10d all-gather ops are runtime parameter management, not model activations, and must execute normally when needed instead of being indexed against the forward's selective-AC op stream. ```python nemo_automodel.components.distributed.activation_checkpointing.ensure_profiler_ops_sac_ignored() -> None ``` Keep `torch.ops.profiler` record-function ops out of SAC's op replay. torch 2.13's FSDP2 runs its pre/post-forward hooks under `torch.autograd.profiler.record_function`, which emits dispatchable `torch.ops.profiler._record_function_*` ops. When an FSDP module boundary sits inside a selective-activation-checkpointed region (e.g. MoE experts sharded separately inside a checkpointed decoder block), those hooks fire a different number of times during the backward recompute than during the forward. SAC replays the forward op stream by per-op invocation index, so the extra profiler op shifts the stream and training fails with `profiler._record_function_enter_new.default invocation index N encountered during backward but not found in storage`. Range ops carry no tensors SAC could cache or restore; adding them to `SAC_IGNORED_OPS` only removes them from the replay accounting (they still execute). No-op before torch 2.13 and on torch builds without `SAC_IGNORED_OPS` or the profiler op namespace. ```python nemo_automodel.components.distributed.activation_checkpointing.ignore_sac_ops( ops: list[object | None] ) -> None ``` Add available operators to PyTorch's selective-AC ignore set. **Parameters:** **`ops`** `list[object | None]` Operators that should execute outside SAC replay accounting. Entries may be `None` when an optional operator is unavailable. --- ```python nemo_automodel.components.distributed.activation_checkpointing.is_selective_activation_checkpointing( activation_checkpointing: object ) -> bool ``` Return whether the config value selects selective activation checkpointing. **Parameters:** **`activation_checkpointing`** `object` The configured value (bool or string such as `"selective"`/`"full"`). --- **Returns:** `bool` `True` only for the string `"selective"` (case- and ```python nemo_automodel.components.distributed.activation_checkpointing.make_selective_checkpoint_context_fn() ``` Build a TorchTitan-style selective activation checkpointing context. ```python nemo_automodel.components.distributed.activation_checkpointing.sdpa_backend_snapshot_context_fn() -> tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager] ``` Snapshot the ambient SDPA state and restore it on checkpoint recompute. A `context_fn` for non-reentrant `checkpoint_wrapper`: torch's non-reentrant checkpoint invokes it at region entry on every checkpointed forward, so the state read here is exactly what the forward runs under. Both the enabled backend set and `F.scaled_dot_product_attention` are restored during recompute. The callable matters for context parallelism: a VLM vision tower temporarily suspends CP's ring-SDPA monkeypatch, and checkpoint recompute occurs after that forward-only suspension has exited. Replaying under the captured callable keeps bidirectional vision attention local while restoring the backward-time CP dispatcher after recompute. Re-pinning the forward-time backend set also prevents checkpoint metadata mismatches when ambient backend forcing (an `sdpa_kernel` pin or module-level backend toggling) is active at forward time but does not span recompute. State toggled *inside* the checkpointed region between attention calls is not captured because the snapshot is taken once at region entry. **Returns:** `AbstractContextManager` `(forward_ctx, recompute_ctx)`: a no-op context for the checkpoint ```python nemo_automodel.components.distributed.activation_checkpointing.transformer_engine_attention_backend_snapshot_context_fn( context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None = None ) -> tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager] ``` Snapshot Transformer Engine's attention-backend cache for checkpoint replay. Transformer Engine caches the parameters and result of attention-backend selection in module-global state. A checkpointed forward can populate that cache, so backward-time recomputation would otherwise enter a different parameter-comparison branch and dispatch a different aten op sequence. The cache is restored to its forward-entry state for recomputation, then reset to the state owned by the surrounding backward pass. **Parameters:** **`context_fn`** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None` — default: None Optional checkpoint context factory to compose inside the cache restoration, such as a selective activation-checkpoint policy. --- **Returns:** `AbstractContextManager` `(forward_ctx, recompute_ctx)` with the supplied forward context and a ```python nemo_automodel.components.distributed.activation_checkpointing.unwrap_checkpoint_wrapper( module: torch.nn.Module ) -> torch.nn.Module ``` Return the activation-checkpointed module, or the input module if it is not wrapped. **Parameters:** **`module`** `nn.Module` Module that may have been wrapped by `checkpoint_wrapper`. --- **Returns:** `nn.Module` The inner checkpointed module when present, otherwise `module`. ```python nemo_automodel.components.distributed.activation_checkpointing.SELECTIVE_AC_WRAPPER_FLAG = '_nemo_selective_ac' ``` ```python nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_MATMUL_OPS = _existing_ops(_resolve_torch_op('aten.mm'), _resolve_torch_op('aten.linear'), _r... ``` ```python nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_MUST_SAVE_OPS = _build_selective_ac_save_ops() ``` ```python nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_TO_COPY_OP = _resolve_torch_op('aten._to_copy') ``` ```python nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_TRACE = os.environ.get('NEMO_SELECTIVE_AC_TRACE', '0').lower() not in ('0', '', 'false',... ``` ```python nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_TRACE_SEEN: set[str] = set() ``` ```python nemo_automodel.components.distributed.activation_checkpointing._TORCH_PROFILER_SAC_IGNORE_MIN_VERSION = (2, 13) ``` ```python nemo_automodel.components.distributed.activation_checkpointing.logger = logging.getLogger(__name__) ```