perf(moe): reuse the DeepEP dispatch layout on activation-checkpoint recompute

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

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(