diff --git a/nemo_automodel/components/moe/megatron/fused_a2a.py b/nemo_automodel/components/moe/megatron/fused_a2a.py index cec285cc1e..1e2c42eaf2 100644 --- a/nemo_automodel/components/moe/megatron/fused_a2a.py +++ b/nemo_automodel/components/moe/megatron/fused_a2a.py @@ -18,6 +18,8 @@ # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE import os +import threading +from contextlib import contextmanager try: from deep_ep import Buffer @@ -27,6 +29,87 @@ except ImportError: HAVE_DEEP_EP = False + +# ── DeepEP dispatch replay across activation-checkpoint recompute ──────────── +# +# Activation checkpointing replays a block's forward during backward, which +# re-runs its MoE dispatch from scratch: DeepEP recomputes the routing layout +# (`get_dispatch_layout` -> `notify_dispatch`) and then moves the tokens. The +# layout is pure a function of the routing, and the routing is identical on the +# replay (checkpointing restores the same inputs and the AC policy pins the +# router's top-k), so that recomputation is redundant. +# +# DeepEP already exposes the cheap path: passing the `handle` returned by a +# previous dispatch skips the layout exchange entirely -- that is what +# `FusedDispatch.backward` does, and it is why the backward dispatch costs +# ~4 ms against ~4.4 s for the layout-computing forward dispatches. +# +# Recording is scoped to one checkpoint frame: the AC wrapper builds a recorder +# per checkpointed call, the forward appends each dispatch's handle and routing +# metadata in call order, and the recompute consumes them in the same order. +# Only the handle and the (small) per-token routing metadata are retained -- +# the dispatched activations themselves are still re-communicated, so the +# memory that activation checkpointing saves is preserved. +class DispatchReplayRecorder: + """Per-checkpoint-frame log of DeepEP dispatch results, replayed on recompute.""" + + def __init__(self) -> None: + self._records: list = [] + self._cursor = 0 + self.replay_misses = 0 + + def record(self, entry) -> None: + self._records.append(entry) + + def take(self): + """Next recorded dispatch, or None when the replay outruns the log.""" + if self._cursor >= len(self._records): + # Recompute issued more dispatches than the forward did. Fall back to + # a full dispatch rather than replaying a mismatched layout. + self.replay_misses += 1 + return None + entry = self._records[self._cursor] + self._cursor += 1 + return entry + + def rewind(self) -> None: + self._cursor = 0 + + +class _DispatchReplayState(threading.local): + """Thread-local replay binding. ``threading.local`` runs ``__init__`` once per + thread, so both attributes always exist on every thread that touches it.""" + + def __init__(self) -> None: + self.recorder: DispatchReplayRecorder | None = None + self.mode: str | None = None + + +_dispatch_replay_state = _DispatchReplayState() + + +def _replay_mode() -> str | None: + return _dispatch_replay_state.mode + + +def _active_recorder() -> "DispatchReplayRecorder | None": + return _dispatch_replay_state.recorder + + +@contextmanager +def dispatch_replay_scope(recorder: "DispatchReplayRecorder | None", mode: str): + """Bind ``recorder`` for the enclosed region. ``mode`` is 'record' or 'replay'.""" + prev_recorder = _dispatch_replay_state.recorder + prev_mode = _dispatch_replay_state.mode + _dispatch_replay_state.recorder = recorder + _dispatch_replay_state.mode = mode if recorder is not None else None + try: + yield + finally: + _dispatch_replay_state.recorder = prev_recorder + _dispatch_replay_state.mode = prev_mode + + try: import importlib.util @@ -154,8 +237,32 @@ def forward( previous_event = None if async_finish: previous_event = EventOverlap(EventHandle()) - # Calculate layout before actual dispatch buffer = get_buffer(group, get_hidden_bytes(x)) + + # Activation-checkpoint replay: reuse the layout this dispatch computed + # on the original forward instead of recomputing it. Cached-mode dispatch + # returns only recv_x, so the routing metadata comes from the record. + recorder = _active_recorder() + if recorder is not None and _replay_mode() == "replay": + replayed = recorder.take() + if replayed is not None: + cached_handle, recv_token_indices, recv_token_probs, tokens_per_expert = replayed + recv_x, _, _, _, _, after_event_overlap = buffer.dispatch( + x, + handle=cached_handle, + previous_event=previous_event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + if async_finish: + after_event_overlap.current_stream_wait() + ctx.group = group + ctx.handle = cached_handle + ctx.async_finish = async_finish + ctx.allocate_on_comm_stream = allocate_on_comm_stream + return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, cached_handle) + + # Calculate layout before actual dispatch ( num_tokens_per_rank, num_tokens_per_rdma_rank, @@ -204,6 +311,11 @@ def forward( ctx.allocate_on_comm_stream = allocate_on_comm_stream tokens_per_expert = torch.tensor(num_recv_tokens_per_expert_list) + if recorder is not None and _replay_mode() == "record": + # Keep the handle and routing metadata (small) so the recompute can + # skip the layout exchange; recv_x is deliberately not retained. + recorder.record((handle, recv_token_indices, recv_token_probs, tokens_per_expert)) + return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle) @staticmethod diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index e95ed49221..2f26d1c69c 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -442,6 +442,59 @@ def _apply_multimodal_tower_ac(model: nn.Module, scopes: tuple[str, ...]) -> Non apply_submodule_checkpointing(tower_layers, has_kv_sharing=False, context_fn=sdpa_backend_snapshot_context_fn) +def _replay_deepep_dispatch_on_recompute( + context_fn: Callable[[], tuple[AbstractContextManager, AbstractContextManager]], +) -> Callable[[], tuple[AbstractContextManager, AbstractContextManager]]: + """Let a checkpointed block's recompute reuse the DeepEP layout it already computed. + + Activation checkpointing replays the block during backward, and the replayed + MoE dispatch recomputes its routing layout from scratch even though the + routing is identical to the forward's. DeepEP skips that exchange when handed + the previous dispatch's handle -- the same shortcut its backward already + takes, which is why the backward dispatch is ~1000x cheaper than a + layout-computing forward one. + + The recorder is built inside ``checkpoint_context_fn``, so each checkpointed + call gets its own: a block called once per forward pass keeps its passes' + dispatches separate, and the recompute consumes them in the order the + matching forward produced them. + """ + from nemo_automodel.components.moe.megatron.fused_a2a import ( + DispatchReplayRecorder, + dispatch_replay_scope, + ) + + def checkpoint_context_fn() -> tuple[AbstractContextManager, AbstractContextManager]: + forward_context, recompute_context = context_fn() + recorder = DispatchReplayRecorder() + + @contextmanager + def scoped(mode, inner): + if mode == "replay": + recorder.rewind() + with dispatch_replay_scope(recorder, mode), inner: + yield + + return scoped("record", forward_context), scoped("replay", recompute_context) + + return checkpoint_context_fn + + +def _uses_hybridep_dispatch(model: nn.Module) -> bool: + """True when any expert module dispatches tokens through HybridEP.""" + modules = getattr(model, "modules", None) + if not callable(modules): + # Duck-typed models (tests, custom wrappers) need not expose + # nn.Module.modules(); this check only gates a warning, so skip it + # rather than making apply_ac require more of the model than the + # checkpointing itself does. + return False + return any( + isinstance(m, (GroupedExpertsDeepEP, GroupedExpertsTE)) and m.dispatcher_backend == "hybridep" + for m in modules() + ) + + def apply_ac( model: nn.Module, ignore_router: bool = True, @@ -484,6 +537,21 @@ def apply_ac( scopes = normalize_activation_checkpointing_scope(activation_checkpointing_scope) checkpoint_decoder = "all" in scopes or "language" in scopes + if checkpoint_decoder and _uses_hybridep_dispatch(model): + logger.warning( + "Activation checkpointing is enabled with the HybridEP token dispatcher. " + "HybridEP's dispatch is not reproducible under checkpoint recompute: replaying it " + "with bit-identical routing can return a different number of tokens than the " + "forward pass produced, which surfaces as torch.utils.checkpoint.CheckpointError " + "('Recomputed values ... have different metadata', e.g. [2791, hidden] vs " + "[2701, hidden]). Verified on 8xH100 with DiffusionGemma-26B-A4B: the router's " + "top-k selection matched exactly on every rank across forward and recompute, so " + "the drift originates in the dispatch, not in routing -- ignore_router_for_ac " + "cannot prevent it. If this run dies with CheckpointError, either switch to " + "dispatcher='deepep' (reproducible under recompute) or disable activation " + "checkpointing (HybridEP is fastest in that mode)." + ) + if checkpoint_decoder and not selective and not ignore_router: logger.warning( "Activation checkpointing is enabled with ignore_router_for_ac=False. The MoE " @@ -614,14 +682,20 @@ def _with_attention_backend_snapshot(context_fn=None): if mtp_repeated and id(block) in mtp_block_ids: continue if ignore_router: + # Only this branch pins routing across recompute (the policy saves the + # router projection and top-k), which is what makes replaying the + # DeepEP layout sound. Do not extend the replay to the else-branch: + # there the router is recomputed and may route differently, so a + # replayed layout would silently mis-route instead of failing loudly. + block_context_fn = _preserve_gate_load_during_recompute( + block, + _with_attention_backend_snapshot(selective_checkpointing_context_fn), + ) block = ptd_checkpoint_wrapper( block, preserve_rng_state=True, determinism_check=_register_moe_checkpoint_determinism_check(), - context_fn=_preserve_gate_load_during_recompute( - block, - _with_attention_backend_snapshot(selective_checkpointing_context_fn), - ), + context_fn=_replay_deepep_dispatch_on_recompute(block_context_fn), ) else: block = ptd_checkpoint_wrapper( diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index d4770b5793..0b93fc6900 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -14,11 +14,16 @@ import sys import types -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from unittest.mock import MagicMock, patch import pytest +# torch.utils.checkpoint.create_selective_checkpoint_contexts returns +# ``(forward_ctx, recompute_ctx)``, and apply_ac's context wrappers unpack it, +# so the stubs below must return a pair rather than a bare sentinel. +SELECTIVE_CTX = ("CTX_FORWARD", "CTX_RECOMPUTE") + class DummyParam: """Mock parameter with requires_grad attribute.""" @@ -258,7 +263,7 @@ class CheckpointPolicy: PREFER_RECOMPUTE = 2 def create_selective_checkpoint_contexts(policy_factory): - return "CTX" + return SELECTIVE_CTX utils_checkpoint_stub.CheckpointPolicy = CheckpointPolicy utils_checkpoint_stub._allowed_determinism_checks_to_fns = {"default": object(), "none": object()} @@ -472,6 +477,31 @@ def reject_unsupported_mtp_cp_pp(model): activation_checkpointing_stub, ) + # apply_ac wraps each block's context in the DeepEP dispatch-replay recorder, + # which imports fused_a2a lazily. Stub it so this module does not depend on + # some earlier test having imported it under the real torch. + fused_a2a_stub = types.ModuleType("nemo_automodel.components.moe.megatron.fused_a2a") + + class _StubDispatchReplayRecorder: + def __init__(self): + self.replay_misses = 0 + self.rewind_count = 0 + + def rewind(self): + self.rewind_count += 1 + + @contextmanager + def _stub_dispatch_replay_scope(recorder, mode): + yield + + fused_a2a_stub.DispatchReplayRecorder = _StubDispatchReplayRecorder + fused_a2a_stub.dispatch_replay_scope = _stub_dispatch_replay_scope + monkeypatch.setitem( + sys.modules, + "nemo_automodel.components.moe.megatron.fused_a2a", + fused_a2a_stub, + ) + distributed_config_stub = types.ModuleType("nemo_automodel.components.distributed.config") distributed_config_stub.normalize_activation_checkpointing_scope = lambda value: ( (value,) if isinstance(value, str) else tuple(value or ("all",)) @@ -668,7 +698,7 @@ def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=N return wrapper_returns.pop(0) wrapper_mock = MagicMock(side_effect=fake_wrapper) - ctx_mock = MagicMock(return_value="CTX") + ctx_mock = MagicMock(return_value=SELECTIVE_CTX) monkeypatch.setattr(P, "ptd_checkpoint_wrapper", wrapper_mock) monkeypatch.setattr(P, "create_selective_checkpoint_contexts", ctx_mock) @@ -698,7 +728,7 @@ def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=N def test_apply_ac_warns_when_router_is_recomputed(monkeypatch): P = _import_parallelizer_with_stubs(monkeypatch) monkeypatch.setattr(P, "ptd_checkpoint_wrapper", MagicMock(side_effect=lambda block, **kw: block)) - monkeypatch.setattr(P, "create_selective_checkpoint_contexts", MagicMock(return_value="CTX")) + monkeypatch.setattr(P, "create_selective_checkpoint_contexts", MagicMock(return_value=SELECTIVE_CTX)) logger_mock = MagicMock() monkeypatch.setattr(P, "logger", logger_mock) @@ -747,12 +777,16 @@ def test_apply_ac_custom_policy_saves_router_projection_and_topk(monkeypatch): def fake_create_selective_checkpoint_contexts(policy_cb): nonlocal captured_policy captured_policy = policy_cb - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): assert preserve_rng_state is True assert callable(context_fn) - assert context_fn() == "CTX" + # The selective contexts are wrapped for DeepEP dispatch replay, so what + # comes back is the wrapper's pair rather than SELECTIVE_CTX itself. + forward_ctx, recompute_ctx = context_fn() + assert hasattr(forward_ctx, "__enter__") + assert hasattr(recompute_ctx, "__enter__") return block monkeypatch.setattr(P, "create_selective_checkpoint_contexts", fake_create_selective_checkpoint_contexts) @@ -2164,7 +2198,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): break if captured_hidden_size is not None: break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2240,7 +2274,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): result = policy_cb(None, torch_stub.ops.aten.mm.default, object(), rhs) if result == P.CheckpointPolicy.MUST_SAVE: captured_num_experts = ne - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2282,7 +2316,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): if result == P.CheckpointPolicy.MUST_SAVE: captured_hidden_size = 512 captured_num_experts = 32 - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2321,7 +2355,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): if result == P.CheckpointPolicy.MUST_SAVE: captured_hidden_size = 1024 captured_num_experts = 64 - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2369,7 +2403,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): break if captured_hidden_size is not None: break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2417,7 +2451,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): break if captured_hidden_size is not None: break - return "CTX" + return SELECTIVE_CTX monkeypatch.setattr(P, "create_selective_checkpoint_contexts", fake_create_selective_checkpoint_contexts) monkeypatch.setattr( @@ -2476,7 +2510,7 @@ def vlm_get_text_module(m): return m.language_model if hasattr(m, "language_model") else m monkeypatch.setattr(P, "get_text_module", vlm_get_text_module) - monkeypatch.setattr(P, "create_selective_checkpoint_contexts", MagicMock(return_value="CTX")) + monkeypatch.setattr(P, "create_selective_checkpoint_contexts", MagicMock(return_value=SELECTIVE_CTX)) monkeypatch.setattr(P, "ptd_checkpoint_wrapper", MagicMock(side_effect=lambda b, **kw: b)) lm_blocks = [DummyBlock(), DummyBlock()] @@ -2738,7 +2772,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): if result == P.CheckpointPolicy.MUST_SAVE: captured_num_experts = ne break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2781,7 +2815,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): if result == P.CheckpointPolicy.MUST_SAVE: captured_num_experts = ne break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2825,7 +2859,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): if result == P.CheckpointPolicy.MUST_SAVE: captured_num_experts = ne break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -2873,7 +2907,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): if result == P.CheckpointPolicy.MUST_SAVE: captured_num_experts = ne break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: @@ -3074,7 +3108,7 @@ def fake_create_selective_checkpoint_contexts(policy_cb): break if captured_hidden_size is not None: break - return "CTX" + return SELECTIVE_CTX def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=None): if context_fn is not None: