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`	Build the set of ops whose activations are always saved under selective AC.
`_default_compute_intensive_ops`	Compute-intensive aten ops from PyTorch’s partitioner, or `()` if unavailable.
`_disable_dynamo_lru_cache`	Best-effort disable of TorchDynamo’s LRU cache for selective AC + compile.
`_existing_ops`	-
`_ffpa_forward_ops`	FFPA forward ops (dense + varlen); import the CuTeDSL kernel so they register first, else `()`.
`_get_transformer_engine_attention_backend_cache`	Return Transformer Engine’s attention-backend cache when available.
`_is_cuda_to_cpu_copy`	-
`_maybe_trace_selective_ac_decision`	Log a selective-AC decision once per op (no-op unless tracing is enabled).
`_registered_child_name`	Return the registered name for a child reached through an attribute.
`_replace_child_module`	Replace `target` with `replacement` in `root`’s module tree.
`_resolve_op_attr`	Resolve a dotted attribute path from `root`, or `None` if any part is absent.
`_resolve_torch_op`	Resolve a dotted `torch.ops` path, defaulting to the `default` overload.
`_restore_sdpa_state`	Temporarily restore the SDPA callable and backend set captured during forward.
`_restore_transformer_engine_attention_backend_cache`	Restore the forward-entry TE attention cache only for recomputation.
`_wrap_first_existing_attr`	Checkpoint-wrap the first matching registered child attr on `module`.
`apply_selective_checkpointing_to_layers`	Wrap whole transformer blocks with the selective-AC policy.
`apply_submodule_checkpointing`	Wrap a transformer block’s sub-modules with `checkpoint_wrapper`.
`detect_kv_sharing_and_maybe_disable_cache`	Detect KV-sharing and disable `use_cache` for non-KV-shared models.
`ensure_fsdp_ops_sac_ignored`	Keep FSDP2 parameter-lifecycle ops out of SAC’s saved-op replay.
`ensure_profiler_ops_sac_ignored`	Keep `torch.ops.profiler` record-function ops out of SAC’s op replay.
`ignore_sac_ops`	Add available operators to PyTorch’s selective-AC ignore set.
`is_selective_activation_checkpointing`	Return whether the config value selects selective activation checkpointing.
`make_selective_checkpoint_context_fn`	Build a TorchTitan-style selective activation checkpointing context.
`sdpa_backend_snapshot_context_fn`	Snapshot the ambient SDPA state and restore it on checkpoint recompute.
`transformer_engine_attention_backend_snapshot_context_fn`	Snapshot Transformer Engine’s attention-backend cache for checkpoint replay.
`unwrap_checkpoint_wrapper`	Return the activation-checkpointed module, or the input module if it is not wrapped.

Data

SELECTIVE_AC_WRAPPER_FLAG

_SELECTIVE_AC_MATMUL_OPS

_SELECTIVE_AC_MUST_SAVE_OPS

_SELECTIVE_AC_TO_COPY_OP

_SELECTIVE_AC_TRACE

_SELECTIVE_AC_TRACE_SEEN

_TORCH_PROFILER_SAC_IGNORE_MIN_VERSION

logger

API

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.

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.

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.

nemo_automodel.components.distributed.activation_checkpointing._existing_ops(
    ops = ()
)

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.

nemo_automodel.components.distributed.activation_checkpointing._get_transformer_engine_attention_backend_cache() -> dict | None

Return Transformer Engine’s attention-backend cache when available.

nemo_automodel.components.distributed.activation_checkpointing._is_cuda_to_cpu_copy(
    func,
    args,
    kwargs
) -> bool

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.

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.

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.

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.

nemo_automodel.components.distributed.activation_checkpointing._resolve_torch_op(
    dotted_path: str
)

Resolve a dotted torch.ops path, defaulting to the default overload.

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.

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.

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.

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.

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]] | NoneDefaults to 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.

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.

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.

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.

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.

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

nemo_automodel.components.distributed.activation_checkpointing.make_selective_checkpoint_context_fn()

Build a TorchTitan-style selective activation checkpointing context.

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

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]] | NoneDefaults to 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

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.

nemo_automodel.components.distributed.activation_checkpointing.SELECTIVE_AC_WRAPPER_FLAG = '_nemo_selective_ac'

nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_MATMUL_OPS = _existing_ops(_resolve_torch_op('aten.mm'), _resolve_torch_op('aten.linear'), _r...

nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_MUST_SAVE_OPS = _build_selective_ac_save_ops()

nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_TO_COPY_OP = _resolve_torch_op('aten._to_copy')

nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_TRACE = os.environ.get('NEMO_SELECTIVE_AC_TRACE', '0').lower() not in ('0', '', 'false',...

nemo_automodel.components.distributed.activation_checkpointing._SELECTIVE_AC_TRACE_SEEN: set[str] = set()

nemo_automodel.components.distributed.activation_checkpointing._TORCH_PROFILER_SAC_IGNORE_MIN_VERSION = (2, 13)

nemo_automodel.components.distributed.activation_checkpointing.logger = logging.getLogger(__name__)