[MoE] Apply router scores before W2 projection

tests/unit_tests/cpu/models/common/test_token_dispatcher.py

@@ -0,0 +1,267 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""The local combine remat region, shared by the standard dispatchers at EP=1.
+
+``LocalTokenDispatcher.combine`` -- also the AllToAll dispatcher's EP=1
+fallback -- runs as one never-replayed ``token_dispatcher.combine`` region,
+like the DeepEP combine: the expert outputs are read only by that saved body,
+so nothing pins them for a replay that never comes.
+
+These tests pin the three consequences on CPU:
+
+- the combine body runs exactly once (never replayed) under every policy,
+  while the policy still decides the w2 projection's replay;
+- with absorbed router scores and w2 saved, the w2 output's storage is freed
+  after the forward, since no persistence marker pins it anymore;
+- forward and backward match the un-checkpointed run bitwise.
+"""
+
+import gc
+import unittest
+from unittest.mock import patch
+
+import torch
+import torch.nn as nn
+import torch_remat as remat
+from torch.multiprocessing.reductions import StorageWeakRef
+
+import torchtitan.models.common.moe.token_dispatcher as token_dispatcher_module
+from torchtitan.distributed.activation_remat import (
+    ActivationRematConfig,
+    DefaultActivationRematConfig,
+    FullActivationRematConfig,
+)
+from torchtitan.models.base.module import Module, ModuleDict
+from torchtitan.models.common.moe.token_dispatcher import (
+    AllToAllTokenDispatcher,
+    LocalTokenDispatcher,
+)
+
+_MODEL_DIM = 8
+_NUM_TOKENS = 16
+_NUM_EXPERTS = 4
+_TOP_K = 2
+
+
+class _W2Region(Module):
+    """Stand-in for the w2 grouped projection: one policy-controlled region.
+
+    The matmul allocates a fresh output and, like a grouped GEMM, saves its
+    input and weight (not its output) for backward, so only the remat
+    machinery could retain the output. ``num_forwards`` counts region-body
+    executions, so a skipped replay does not count.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(_MODEL_DIM, _MODEL_DIM))
+        self.num_forwards = 0
+
+    def forward(self, x_RD: torch.Tensor) -> torch.Tensor:
+        return remat.region(
+            self._project,
+            self.remat_region_name("grouped_mm"),
+            recompute=self.remat_should_recompute("grouped_mm"),
+        )(x_RD)
+
+    def _project(self, x_RD: torch.Tensor) -> torch.Tensor:
+        self.num_forwards += 1
+        return x_RD @ self.weight
+
+
+class _DispatcherBlock(Module):
+    """Routed-experts block: dispatch, one projection, combine.
+
+    The router scores depend on the gate, so a wrong score placement or a
+    misaligned combine changes the gradients, as in training.
+    """
+
+    def __init__(self, dispatcher_config: LocalTokenDispatcher.Config):
+        super().__init__()
+        self.token_dispatcher = dispatcher_config.build()
+        self.w2 = _W2Region()
+        self.gate = nn.Linear(_MODEL_DIM, _NUM_EXPERTS, bias=False)
+        self.num_experts = _NUM_EXPERTS
+        self.top_k = dispatcher_config.top_k
+
+    def forward(self, x_TD: torch.Tensor) -> torch.Tensor:
+        num_tokens = x_TD.shape[0]
+        token_ids_T1 = torch.arange(num_tokens, device=x_TD.device).unsqueeze(-1)
+        offsets_1K = torch.arange(self.top_k, device=x_TD.device).unsqueeze(0)
+        expert_ids_TK = (token_ids_T1 + offsets_1K) % self.num_experts
+        scores_TK = torch.sigmoid(self.gate(x_TD)).gather(1, expert_ids_TK)
+        num_tokens_per_expert_E = torch.bincount(
+            expert_ids_TK.flatten(), minlength=self.num_experts
+        )
+        routed_input_RD, _num_tokens_E, metadata = self.token_dispatcher.dispatch(
+            x_TD, scores_TK, expert_ids_TK, num_tokens_per_expert_E
+        )
+        routed_scores_R = getattr(metadata, "routed_scores_R", None)
+        if routed_scores_R is not None:
+            # Absorbed scores scale the projection input, as GroupedLinear
+            # does before its grouped_mm region.
+            remat.recompute_needs_tensor(routed_scores_R)
+            routed_input_RD = (
+                routed_input_RD.float() * routed_scores_R.reshape(-1, 1)
+            ).to(routed_input_RD.dtype)
+        routed_output_RD = self.w2(routed_input_RD)
+        out_TD = self.token_dispatcher.combine(routed_output_RD, metadata, x_TD)
+        # The loss is a bare consumer of the combine output.
+        remat.recompute_needs_tensor(out_TD)
+        return out_TD.double().sum()
+
+
+class _Model(Module):
+    def __init__(self, block: Module):
+        super().__init__()
+        self.layers = ModuleDict({"0": block})
+
+    def forward(self, x_TD: torch.Tensor) -> torch.Tensor:
+        return self.layers["0"](x_TD)
+
+
+def _forward_backward(
+    model: Module, x_TD: torch.Tensor
+) -> tuple[torch.Tensor, list[torch.Tensor]]:
+    """Run one forward/backward; return the loss and input/parameter grads."""
+    model.zero_grad(set_to_none=True)
+    input_TD = x_TD.detach().clone().requires_grad_(True)
+    loss = model(input_TD)
+    loss.backward()
+    assert input_TD.grad is not None
+    grads = [input_TD.grad.detach().clone()]
+    grads += [parameter.grad.detach().clone() for parameter in model.parameters()]
+    return loss.detach().clone(), grads
+
+
+class TestLocalCombineRematRegion(unittest.TestCase):
+    @staticmethod
+    def _dispatcher_cases():
+        # name, dispatcher Config. Without a sparse mesh the AllToAll
+        # dispatcher falls back to the local dispatch/combine methods.
+        return (
+            ("local", LocalTokenDispatcher.Config),
+            ("alltoall_ep1", AllToAllTokenDispatcher.Config),
+        )
+
+    @staticmethod
+    def _build(config_cls, absorb_router_scores: bool) -> _Model:
+        torch.manual_seed(42)
+        return _Model(
+            _DispatcherBlock(
+                config_cls(
+                    num_experts=_NUM_EXPERTS,
+                    top_k=_TOP_K,
+                    absorb_router_scores=absorb_router_scores,
+                )
+            )
+        )
+
+    def test_combine_never_replays_and_matches_eager(self):
+        for name, config_cls in self._dispatcher_cases():
+            for absorb in (False, True):
+                for policy, expected_w2_forwards in (
+                    (FullActivationRematConfig(), 2),
+                    (DefaultActivationRematConfig(), 1),
+                ):
+                    with self.subTest(
+                        dispatcher=name,
+                        absorb=absorb,
+                        policy=type(policy).__qualname__,
+                    ):
+                        baseline = self._build(config_cls, absorb)
+                        remat_model = self._build(config_cls, absorb)
+                        remat_model.load_state_dict(baseline.state_dict())
+                        policy.apply(remat_model)
+                        block = remat_model.layers["0"]
+
+                        num_scatter_adds = 0
+                        original_scatter_add = (
+                            token_dispatcher_module.deterministic_scatter_add
+                        )
+
+                        def counted_scatter_add(*args, **kwargs):
+                            nonlocal num_scatter_adds
+                            num_scatter_adds += 1
+                            return original_scatter_add(*args, **kwargs)
+
+                        torch.manual_seed(0)
+                        x_TD = torch.randn(_NUM_TOKENS, _MODEL_DIM)
+                        expected = _forward_backward(baseline, x_TD)
+                        with patch.object(
+                            token_dispatcher_module,
+                            "deterministic_scatter_add",
+                            counted_scatter_add,
+                        ):
+                            actual = _forward_backward(remat_model, x_TD)
+
+                        # The combine body runs only on the original forward;
+                        # the policy still decides the w2 projection's replay.
+                        self.assertEqual(num_scatter_adds, 1)
+                        self.assertEqual(block.w2.num_forwards, expected_w2_forwards)
+                        torch.testing.assert_close(
+                            actual[0], expected[0], rtol=0, atol=0
+                        )
+                        for actual_grad, expected_grad in zip(actual[1], expected[1]):
+                            torch.testing.assert_close(
+                                actual_grad, expected_grad, rtol=0, atol=0
+                            )
+
+    def test_absorbed_combine_does_not_pin_w2_output(self):
+        # With the router scores absorbed, no backward consumer needs the w2
+        # output (the scatter-add saves only indices), and the combine region
+        # reads it only from its saved body, so saving w2 must not keep the
+        # w2 output alive.
+        for name, config_cls in self._dispatcher_cases():
+            with self.subTest(dispatcher=name):
+                baseline = self._build(config_cls, absorb_router_scores=True)
+                remat_model = self._build(config_cls, absorb_router_scores=True)
+                remat_model.load_state_dict(baseline.state_dict())
+                ActivationRematConfig(save_regions=["w2.grouped_mm"]).apply(remat_model)
+                w2 = remat_model.layers["0"].w2
+                w2_output_storages = []
+                original_forward = _W2Region.forward
+
+                def recorded_forward(module, *args, **kwargs):
+                    output = original_forward(module, *args, **kwargs)
+                    # The forward also runs (with its region skipped) during
+                    # replay.
+                    if module is w2 and not remat.is_recomputing():
+                        # remat retains a detached alias, so track the storage.
+                        w2_output_storages.append(
+                            StorageWeakRef(output.untyped_storage())
+                        )
+                    return output
+
+                torch.manual_seed(0)
+                x_TD = torch.randn(_NUM_TOKENS, _MODEL_DIM)
+                expected = _forward_backward(baseline, x_TD)
+                with patch.object(
+                    _W2Region, "forward", autospec=True, side_effect=recorded_forward
+                ):
+                    input_TD = x_TD.detach().clone().requires_grad_(True)
+                    loss = remat_model(input_TD)
+                    gc.collect()
+                    retained = not w2_output_storages[0].expired()
+                    loss.backward()
+
+                self.assertFalse(retained)
+                self.assertEqual(len(w2_output_storages), 1)
+                torch.testing.assert_close(loss.detach(), expected[0], rtol=0, atol=0)
+                actual_grads = [input_TD.grad.detach().clone()]
+                actual_grads += [
+                    parameter.grad.detach().clone()
+                    for parameter in remat_model.parameters()
+                ]
+                for actual_grad, expected_grad in zip(actual_grads, expected[1]):
+                    torch.testing.assert_close(
+                        actual_grad, expected_grad, rtol=0, atol=0
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit_tests/cpu/recipes/test_debug_config_defaults.py

@@ -26,6 +26,8 @@
 )
 from torchtitan_recipes.tests.models.deepseek_v3 import (
     deepseek_v3_debugmodel,
+    deepseek_v3_debugmodel_alltoall,
+    deepseek_v3_debugmodel_deepep,
     deepseek_v3_debugmodel_hybridep,
     deepseek_v3_debugmodel_mtp,
     deepseek_v3_debugmodel_mxfp8,
@@ -90,6 +92,8 @@
 
 _DEBUG_CONFIG_FACTORIES: tuple[DebugConfigFactory, ...] = (
     deepseek_v3_debugmodel,
+    deepseek_v3_debugmodel_alltoall,
+    deepseek_v3_debugmodel_deepep,
     deepseek_v3_debugmodel_hybridep,
     deepseek_v3_debugmodel_mtp,
     deepseek_v3_debugmodel_mxfp8,

tests/unit_tests/gpu/distributed/test_deepep_activation_checkpoint.py

@@ -128,7 +128,8 @@ def test_dispatch_and_combine_are_saved_not_replayed(self):
         # DeepEP fills receive slots with atomics, so a replayed dispatch can
         # receive rows in another order than the forward's handle records.
         # Backward would then pair each token's gradient with another token's
-        # activations. These policies save dispatch and combine instead.
+        # activations. These policies save the dispatch communication, and the
+        # combine region is never replayed regardless.
         self._check_deepep_activation_checkpointing(
             (
                 DefaultActivationRematConfig(),
@@ -139,9 +140,13 @@ def test_dispatch_and_combine_are_saved_not_replayed(self):
 
     @with_comms
     def test_region_ac_replay_is_exact(self):
-        # The full policy and policies without ep_communication saved replay dispatch and
-        # combine. DeepEP then uses a deterministic buffer, which receives rows in
-        # the forward's order, so the gradients match the run without AC bitwise.
+        # The full policy and policies without ep_communication saved replay
+        # the dispatch communication. DeepEP then uses a deterministic buffer,
+        # which receives rows in the forward's order, so the gradients match
+        # the run without AC bitwise. The combine region is never replayed:
+        # it folds the score multiply, unpermute and combine kernel into one
+        # saved body, so combine runs only in the forward and the dispatch
+        # backward.
         self._check_deepep_activation_checkpointing(
             (
                 FullActivationRematConfig(),
@@ -230,7 +235,9 @@ def call(*args, **kwargs):
                 self.assertEqual(baseline_calls, {"dispatch": 2, "combine": 2})
                 self.assertEqual(
                     num_calls,
-                    {"dispatch": 2 + num_replays, "combine": 2 + num_replays},
+                    # Only the dispatch communication replays; the combine
+                    # region is never replayed.
+                    {"dispatch": 2 + num_replays, "combine": 2},
                     msg=f"unexpected DeepEP replays: {num_calls}",
                 )
                 if num_replays:

tests/unit_tests/gpu/distributed/test_remat_all_to_all.py

@@ -156,14 +156,15 @@ def test_collective_replay_follows_region_policy(self):
         )
 
         # dispatch holds the count exchange, its device-to-host sync, and the
-        # dispatch all-to-all; combine holds the combine all-to-all. The default
-        # policy recomputes the routed-expert projections but must retain both
-        # dispatcher regions under the same routed_experts prefix.
+        # dispatch all-to-all; combine holds the combine all-to-all. The
+        # dispatch region follows the policy; the combine region is never
+        # replayed (like DeepEP's), so only the dispatch collectives replay,
+        # and a combine-only save pattern changes nothing.
         dispatch = "routed_experts.token_dispatcher.dispatch"
         combine = "routed_experts.token_dispatcher.combine"
         for policy_config, expected_replay_collectives, expected_replay_syncs in (
-            (ActivationRematConfig(save_regions=[]), 3, 1),
-            (ActivationRematConfig(save_regions=[dispatch]), 1, 0),
+            (ActivationRematConfig(save_regions=[]), 2, 1),
+            (ActivationRematConfig(save_regions=[dispatch]), 0, 0),
             (ActivationRematConfig(save_regions=[combine]), 2, 1),
             (ActivationRematConfig(save_regions=[dispatch, combine]), 0, 0),
             (DefaultActivationRematConfig(), 0, 0),
@@ -223,9 +224,10 @@ def counted_sync(*args, **kwargs):
     @with_comms
     def test_saved_w2_output_is_not_retained_under_ep(self):
         # Unpermute and the combine all-to-all do not need the w2 output in
-        # backward, so saving w2 and the combine must not keep the w2 output
-        # alive. bf16 activations, as in mixed-precision training, need no cast
-        # after w2.
+        # backward, so saving w2 must not keep the w2 output alive: the
+        # combine region is always saved and reads it only from its body, so
+        # it needs no persisting for a replay. bf16 activations, as in
+        # mixed-precision training, need no cast after w2.
         mesh = init_device_mesh(
             self.device_type,
             (self.world_size,),
@@ -249,10 +251,7 @@ def test_saved_w2_output_is_not_retained_under_ep(self):
             )
             remat_model.load_state_dict(baseline.state_dict())
             ActivationRematConfig(
-                save_regions=[
-                    "routed_experts.w2.grouped_mm",
-                    "routed_experts.token_dispatcher.combine",
-                ]
+                save_regions=["routed_experts.w2.grouped_mm"]
             ).apply(remat_model)
             w2 = remat_model.layers["0"].routed_experts.w2
             w2_output_refs = []