fix(dist): keep profiler record-function ops out of SAC replay accounting

nemo_automodel/components/distributed/activation_checkpointing.py

@@ -40,8 +40,12 @@
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
 
+from nemo_automodel.shared.import_utils import get_torch_version
+
 logger = logging.getLogger(__name__)
 
+_TORCH_PROFILER_SAC_IGNORE_MIN_VERSION = (2, 13)
+
 
 def unwrap_checkpoint_wrapper(module: nn.Module) -> nn.Module:
     """Return the activation-checkpointed module, or the input module if it is not wrapped.
@@ -253,8 +257,43 @@ def _maybe_trace_selective_ac_decision(func, decision, is_alternating: bool, *,
     logger.info("[selective-ac] %s -> %s", key, verdict)
 
 
+def 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.
+    """
+    if get_torch_version().release < _TORCH_PROFILER_SAC_IGNORE_MIN_VERSION:
+        return
+
+    sac_ignored = getattr(torch.utils.checkpoint, "SAC_IGNORED_OPS", None)
+    profiler_ops = getattr(torch.ops, "profiler", None)
+    if sac_ignored is None or profiler_ops is None:
+        return
+    for packet_name in ("_record_function_enter", "_record_function_enter_new", "_record_function_exit"):
+        packet = getattr(profiler_ops, packet_name, None)
+        if packet is None:
+            continue
+        for overload_name in packet.overloads():
+            sac_ignored.add(getattr(packet, overload_name))
+
+
 def make_selective_checkpoint_context_fn():
     """Build a TorchTitan-style selective activation checkpointing context."""
+    ensure_profiler_ops_sac_ignored()
 
     def selective_checkpointing_context_fn():
         # Count matmuls separately for the forward and recompute passes. torch

nemo_automodel/components/moe/parallelizer.py

@@ -492,6 +492,10 @@ def _custom_policy(ctx, func, *args, **kwargs):
     def selective_checkpointing_context_fn():
         return create_selective_checkpoint_contexts(_custom_policy)
 
+    from nemo_automodel.components.distributed.activation_checkpointing import ensure_profiler_ops_sac_ignored
+
+    ensure_profiler_ops_sac_ignored()
+
     # Weight-tied (use_repeated_layer) MTP head blocks must NOT be activation
     # checkpointed: the single physical block is recomputed once per MTP depth in
     # backward, and FSDP2 cannot re-unshard the *shared* EP-sharded experts param

tests/unit_tests/distributed/test_activation_checkpointing.py

@@ -20,6 +20,7 @@
 import pytest
 import torch
 import torch.nn.functional as F
+from packaging.version import Version
 from torch import nn
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper, checkpoint_wrapper
 from torch.nn.attention import SDPBackend, sdpa_kernel
@@ -330,3 +331,79 @@ def test_detect_kv_sharing_leaves_cache_enabled_for_kv_shared_models():
     assert has_kv_sharing is True
     assert model.config.use_cache is True
     assert model.config.text_config.use_cache is True
+
+
+def _sac_context_factory():
+    from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
+
+    def policy(ctx, func, *args, **kwargs):
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return lambda: create_selective_checkpoint_contexts(policy)
+
+
+def test_profiler_ops_sac_ignore_skips_before_torch_2_13(monkeypatch):
+    sac_ignored = set()
+    monkeypatch.setattr(ac, "get_torch_version", lambda: Version("2.12.0"))
+    monkeypatch.setattr(torch.utils.checkpoint, "SAC_IGNORED_OPS", sac_ignored, raising=False)
+
+    ac.ensure_profiler_ops_sac_ignored()
+
+    assert sac_ignored == set()
+
+
+def test_profiler_ops_sac_ignore_includes_torch_2_13_alpha(monkeypatch):
+    op = object()
+    packet = type("FakePacket", (), {"default": op, "overloads": lambda self: ("default",)})()
+    profiler_ops = type(
+        "FakeProfilerOps",
+        (),
+        {
+            "_record_function_enter": packet,
+            "_record_function_enter_new": packet,
+            "_record_function_exit": packet,
+        },
+    )()
+    sac_ignored = set()
+    monkeypatch.setattr(ac, "get_torch_version", lambda: Version("2.13.0a0+8145d630e8"))
+    monkeypatch.setattr(torch.utils.checkpoint, "SAC_IGNORED_OPS", sac_ignored, raising=False)
+    monkeypatch.setattr(torch.ops, "profiler", profiler_ops, raising=False)
+
+    ac.ensure_profiler_ops_sac_ignored()
+
+    assert sac_ignored == {op}
+
+
+def test_sac_replay_tolerates_recompute_only_record_function(monkeypatch):
+    """A profiler range entered only during backward recompute must not desync SAC replay.
+
+    torch 2.13's FSDP2 runs its hooks under ``record_function``; with an FSDP
+    boundary inside a SAC region the range ops fire a different number of times
+    in the recompute than in the forward, which shifts SAC's per-op replay
+    index and raises ``... encountered during backward but not found in
+    storage``. ``ensure_profiler_ops_sac_ignored`` keeps profiler ops out of
+    the replay accounting.
+    """
+    if not hasattr(torch.utils.checkpoint, "SAC_IGNORED_OPS"):
+        pytest.skip("torch build without SAC_IGNORED_OPS")
+
+    monkeypatch.setattr(ac, "get_torch_version", lambda: Version("2.13.0a0+8145d630e8"))
+    ac.ensure_profiler_ops_sac_ignored()
+    assert torch.ops.profiler._record_function_enter_new.default in torch.utils.checkpoint.SAC_IGNORED_OPS
+
+    linear = nn.Linear(4, 4)
+    calls = {"n": 0}
+
+    def fn(x):
+        calls["n"] += 1
+        if calls["n"] > 1:  # backward-time recompute takes a different hook path
+            with torch.autograd.profiler.record_function("recompute-only-range"):
+                return linear(x)
+        return linear(x)
+
+    x = torch.randn(2, 4, requires_grad=True)
+    out = torch.utils.checkpoint.checkpoint(fn, x, use_reentrant=False, context_fn=_sac_context_factory())
+    out.sum().backward()
+
+    assert calls["n"] == 2
+    assert x.grad is not None