[`Moe`] Fix grouped mm fallback

@@ -233,22 +233,43 @@ def _grouped_mm_fallback_setup_context(ctx, inputs, output):     ctx.offs = inputs[2]  -def _grouped_mm_fallback_backward(ctx, grad_output):-    """Backward pass for `_grouped_mm_fallback`. Computes grad_input and grad_weight per expert group; offs has no gradient."""-    input, weight = ctx.saved_tensors+def _grouped_mm_fallback_bwd(+    input: torch.Tensor,+    weight: torch.Tensor,+    offs: torch.Tensor,+    grad_output: torch.Tensor,+) -> tuple[torch.Tensor, torch.Tensor]:+    """Backward implementation for the grouped matmul fallback."""     grad_input = torch.zeros_like(input)     grad_weight = torch.zeros_like(weight)      start = 0     # single cpu<->gpu sync point here,     # avoids multiple syncs inside the loop-    for i, end in enumerate(ctx.offs.tolist()):+    for i, end in enumerate(offs.tolist()):         if start == end:             continue         torch.mm(grad_output[start:end], weight[i].T, out=grad_input[start:end])         torch.mm(input[start:end].T, grad_output[start:end], out=grad_weight[i])         start = end +    return grad_input, grad_weight+++def _grouped_mm_fallback_backward_fake(+    input: torch.Tensor,+    weight: torch.Tensor,+    offs: torch.Tensor,+    grad_output: torch.Tensor,+) -> tuple[torch.Tensor, torch.Tensor]:+    """Shape/dtype inference for the backward custom operator."""+    return torch.empty_like(input), torch.empty_like(weight)+++def _grouped_mm_fallback_backward(ctx, grad_output):+    """Backward pass for `_grouped_mm_fallback`."""+    input, weight = ctx.saved_tensors+    grad_input, grad_weight = torch.ops.transformers.grouped_mm_fallback_backward(input, weight, ctx.offs, grad_output)     return grad_input, grad_weight, None  @@ -265,6 +286,17 @@ def _grouped_mm_fallback_backward(ctx, grad_output):         _grouped_mm_fallback_backward,         setup_context=_grouped_mm_fallback_setup_context,     )+    # extra bwd registration to make autograd avoid looking into the experts loop+    torch.library.custom_op(+        "transformers::grouped_mm_fallback_backward",+        _grouped_mm_fallback_bwd,+        mutates_args=(),+        schema="(Tensor input, Tensor weight, Tensor offs, Tensor grad_output) -> (Tensor, Tensor)",+    )+    torch.library.register_fake(+        "transformers::grouped_mm_fallback_backward",+        _grouped_mm_fallback_backward_fake,+    )   def _can_use_grouped_mm(input: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> bool:

@@ -21,11 +21,12 @@  import types import unittest+from unittest.mock import patch  import torch  from transformers.activations import ACT2FN-from transformers.integrations.moe import batched_mm_experts_forward+from transformers.integrations.moe import batched_mm_experts_forward, grouped_mm_experts_forward from transformers.testing_utils import require_torch, torch_device  @@ -90,3 +91,60 @@ def test_sentinel_slots_do_not_reach_the_output(self):         moved = self.top_k_index.masked_fill(self.top_k_index >= NUM_EXPERTS, NUM_EXPERTS + 1)         out_moved = batched_mm_experts_forward(self.experts, self.hidden_states, moved, self._weights())         torch.testing.assert_close(out, out_moved)+++@require_torch+class GroupedMmFallbackCompileTest(unittest.TestCase):+    def setUp(self):+        torch.manual_seed(0)++        self.experts = make_experts(num_experts=NUM_EXPERTS, hidden=8, inter=16, is_expert_parallel=False)+        self.experts.gate_up_proj.requires_grad_(True)+        self.experts.down_proj.requires_grad_(True)++        self.hidden_states = torch.randn(3, 8, device=torch_device, requires_grad=True)+        # Leave experts 1 and 3 empty+        self.top_k_index = torch.tensor([[0, 2], [2, 0], [0, 2]], device=torch_device)+        self.top_k_weights = torch.tensor(+            [[0.7, 0.3], [0.4, 0.6], [0.2, 0.8]],+            device=torch_device,+        )++    def test_forward_backward_fullgraph(self):+        def forward():+            return grouped_mm_experts_forward(+                self.experts,+                self.hidden_states,+                self.top_k_index,+                self.top_k_weights,+            )++        def run(fn):+            self.hidden_states.grad = None+            self.experts.gate_up_proj.grad = None+            self.experts.down_proj.grad = None++            output = fn()+            output.square().sum().backward()  # arbitrary we only want the bwd to be triggered++            return (+                output.detach().clone(),+                self.hidden_states.grad.detach().clone(),+                self.experts.gate_up_proj.grad.detach().clone(),+                self.experts.down_proj.grad.detach().clone(),+            )++        # Force the fallback to be selected+        with patch(+            "transformers.integrations.moe._can_use_grouped_mm",+            new=lambda *args: False,+        ):+            eager_results = run(forward)+            compiled_results = run(torch.compile(forward, fullgraph=True))++        for eager, compiled in zip(eager_results, compiled_results):+            torch.testing.assert_close(compiled, eager, rtol=1e-4, atol=1e-4)++        # Unused experts should have zero gradients+        for grad in compiled_results[2:]:+            self.assertTrue(torch.all(grad[1::2] == 0))