Fix for PR-2142

megatron/core/transformer/moe/moe_layer.py

@@ -437,11 +437,12 @@ def custom_forward(hidden_states, intermediate_tensors, padding_mask=None):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     parallel_state.get_tensor_model_parallel_group(),
                     hidden_states,
+                    intermediate_tensors,
                     padding_mask,
                 )
             else:
                 outputs = tensor_parallel.checkpoint(
-                    custom_forward, False, hidden_states, padding_mask
+                    custom_forward, False, hidden_states, intermediate_tensors, padding_mask
                 )
         else:
             outputs = custom_forward(hidden_states, intermediate_tensors, padding_mask)

tests/unit_tests/transformer/moe/test_moe_layer.py

@@ -276,3 +276,123 @@ def test_moe_layer_fp16_forward_backward(
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
+
+
+class TestMoELayerRecompute:
+    """Test MoE layer with recompute enabled (activation checkpointing).
+
+    Tests both code paths:
+    - fp8=False: uses tensor_parallel.checkpoint
+    - fp8=True: uses te_checkpoint (requires TE >= 1.7.0)
+    """
+
+    def setup_method(self, method):
+        pass
+
+    @pytest.mark.parametrize("moe_token_dispatcher_type", ["allgather", "alltoall"])
+    @pytest.mark.parametrize("num_moe_experts", [2, 4])
+    @pytest.mark.parametrize("with_padding_mask", [True, False])
+    @pytest.mark.parametrize("tp_size,ep_size", [(1, 1), (4, 2)])
+    @pytest.mark.parametrize("fp8", [False, True])
+    def test_moe_layer_recompute_forward_backward(
+        self, num_moe_experts, moe_token_dispatcher_type, with_padding_mask, tp_size, ep_size, fp8
+    ):
+        """Test MoE layer forward and backward pass with recompute enabled.
+
+        When fp8=False, uses tensor_parallel.checkpoint.
+        When fp8=True, uses te_checkpoint (requires TE >= 1.7.0).
+        """
+        # Skip fp8 tests if TE version is not sufficient
+        if fp8 and not is_te_min_version("1.7.0.dev0"):
+            pytest.skip("FP8 MoE recompute requires TE 1.7.0 and later.")
+
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size, expert_model_parallel_size=ep_size
+        )
+        _set_random_seed(seed_=123, data_parallel_random_init=False)
+
+        hidden_size = 64
+        sequence_length = 32
+        micro_batch_size = 2
+
+        transformer_config = TransformerConfig(
+            num_layers=1,
+            hidden_size=hidden_size,
+            num_attention_heads=4,
+            num_moe_experts=num_moe_experts,
+            use_cpu_initialization=False,
+            moe_token_dispatcher_type=moe_token_dispatcher_type,
+            moe_router_load_balancing_type="aux_loss",
+            moe_router_topk=2,
+            moe_aux_loss_coeff=0.01,
+            moe_grouped_gemm=False,
+            moe_ffn_hidden_size=256,
+            add_bias_linear=False,
+            # Enable recompute for MoE layer
+            recompute_granularity="selective",
+            recompute_modules=["moe"],
+            tensor_model_parallel_size=tp_size,
+            expert_model_parallel_size=ep_size,
+            sequence_parallel=tp_size > 1,
+            fp8=fp8,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+        )
+
+        # Use TE spec for fp8, local spec otherwise
+        if fp8:
+            transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+                num_experts=num_moe_experts, moe_grouped_gemm=False
+            )
+        else:
+            transformer_layer_spec = get_gpt_layer_local_spec(
+                num_experts=num_moe_experts, moe_grouped_gemm=False
+            )
+
+        moe_layer = MoELayer(
+            transformer_config, transformer_layer_spec.submodules.mlp.submodules
+        ).cuda()
+
+        hidden_states = torch.randn(
+            sequence_length,
+            micro_batch_size,
+            hidden_size,
+            device=torch.cuda.current_device(),
+            dtype=torch.bfloat16,
+            requires_grad=True,
+        )
+
+        # Create padding mask if needed: shape [batch_size, sequence_length]
+        padding_mask = None
+        if with_padding_mask:
+            padding_mask = torch.ones(
+                micro_batch_size,
+                sequence_length,
+                device=torch.cuda.current_device(),
+                dtype=torch.bool,
+            )
+            # Mark last 4 tokens as padding for each batch
+            padding_mask[:, -4:] = False
+
+        output, _ = moe_layer(hidden_states, padding_mask=padding_mask)
+
+        assert output.dtype == torch.bfloat16, f"Expected bf16 output, got {output.dtype}"
+        assert output.shape == hidden_states.shape, f"Output shape mismatch"
+
+        # Backward pass - this is where recompute/checkpoint is actually used
+        loss = output.sum()
+        loss.backward()
+
+        assert hidden_states.grad is not None, "Input gradients should exist"
+        assert (
+            hidden_states.grad.dtype == torch.bfloat16
+        ), f"Expected bf16 gradients, got {hidden_states.grad.dtype}"
+
+        for name, param in moe_layer.named_parameters():
+            if param.requires_grad:
+                assert param.grad is not None, f"Gradient for {name} should exist"
+
+        Utils.destroy_model_parallel()
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()