diff --git a/nemo_automodel/components/distributed/activation_checkpointing.py b/nemo_automodel/components/distributed/activation_checkpointing.py index 81f72a287d..34932a4e05 100644 --- a/nemo_automodel/components/distributed/activation_checkpointing.py +++ b/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 diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index 87654ea3f1..c188415b58 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/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 diff --git a/tests/unit_tests/distributed/test_activation_checkpointing.py b/tests/unit_tests/distributed/test_activation_checkpointing.py index 7f7b29f179..f4fd60af2c 100644 --- a/tests/unit_tests/distributed/test_activation_checkpointing.py +++ b/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