Repository: huggingface/transformers
Pull request: Optimize MoEs for decoding using batched_mm

# What does this PR do?

batched_mm is more performant on "smaller inputs", especially when used with torch.compile.
smaller inputs: up to 2^6-2^9 tokens depending on the num experts and num topk

<img width="2000" height="2000" alt="crossover_plots" src="https://github.com/user-attachments/assets/21a9f333-b3d5-40ae-90a6-a060da98eb3b" />
<img width="1168" height="243" alt="generation benchmark" src="https://github.com/user-attachments/assets/54daeb58-e7dc-45ca-a29f-351f12c8b13d" />



## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/generation/utils.py
patch: @@ -19,6 +19,7 @@
 import os
 import warnings
 from collections.abc import Callable
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional, Union
 
@@ -2192,17 +2193,24 @@ def _valid_auto_compile_criteria(self, model_kwargs: dict[str, Any], generation_
                     )
                     generation_config.compile_config.fullgraph = False
 
-            # If we use grouped_mm and dtype different than bfloat16, we fallback to batched_mm
-            if self.config._experts_implementation == "grouped_mm":
-                if self.dtype != torch.bfloat16:
-                    logger.warning_once(
-                        "torch._grouped_mm currently only supports bfloat16 when being compiled with torch.compile. "
-                        "Falling back to batched_mm implementation for compilation."
-                    )
-                    self.set_experts_implementation("batched_mm")
-
         return can_compile
 
+    @contextmanager
+    def _optimize_model_for_decode(self):
+        original_experts_implementation = self.config._experts_implementation
+        if original_experts_implementation == "grouped_mm":
+            logger.info_once(
+                "We will be switching to 'batched_mm' for the decoding stage as it is much more performant than 'grouped_mm' on smaller inputs. "
+                "If you experience any issues with this, please open an issue on the Hugging Face Transformers GitHub repository.",
+            )
+            self.set_experts_implementation("batched_mm")
+
+        try:
+            yield
+        finally:
+            if original_experts_implementation == "grouped_mm":
+                self.set_experts_implementation(original_experts_implementation)
+
     def _get_deprecated_gen_repo(
         self,
         generation_mode: GenerationMode,
@@ -2860,7 +2868,8 @@ def _sample(
         while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
             if prefill_consumed:
                 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-                outputs = model_forward(**model_inputs, return_dict=True)
+                with self._optimize_model_for_decode():
+                    outputs = model_forward(**model_inputs, return_dict=True)
             prefill_consumed = True
             model_kwargs = self._update_model_kwargs_for_generation(
                 outputs,

filename: tests/models/deepseek_v2/test_modeling_deepseek_v2.py
patch: @@ -16,8 +16,6 @@
 
 import unittest
 
-import pytest
-
 from transformers import BitsAndBytesConfig, Cache, is_torch_available
 from transformers.testing_utils import require_read_token, require_torch, require_torch_accelerator, slow, torch_device
 
@@ -135,11 +133,6 @@ def test_model_rope_scaling_frequencies(self):
         with self.assertRaises(AssertionError):
             torch.testing.assert_close(yarn_freqs_cis_long, original_freqs_cis_long)
 
-    @unittest.skip("Dynamic control flow in MoE")
-    @pytest.mark.torch_compile_test
-    def test_torch_compile_for_training(self):
-        pass
-
     def test_tp_plan_matches_params(self):
         """Need to overwrite as the plan contains keys that are valid but depend on some configs flags and cannot
         be valid all at the same time"""

filename: tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py
patch: @@ -161,10 +161,7 @@ def get_large_model(cls):
     @classmethod
     def get_small_model(cls):
         cls.model = Ernie4_5_MoeForCausalLM.from_pretrained(
-            "hf-internal-testing/ERNIE-4.5-Small-Moe",
-            device_map="auto",
-            dtype="auto",
-            experts_implementation="eager",
+            "hf-internal-testing/ERNIE-4.5-Small-Moe", device_map="auto", dtype="auto", experts_implementation="eager"
         )
 
         return cls.model

filename: tests/models/flex_olmo/test_modeling_flex_olmo.py
patch: @@ -16,8 +16,6 @@
 
 import unittest
 
-import pytest
-
 from transformers import is_torch_available
 from transformers.models.auto.tokenization_auto import AutoTokenizer
 from transformers.testing_utils import (
@@ -57,11 +55,6 @@ class FlexOlmoModelTest(CausalLMModelTest, unittest.TestCase):
     # used in `test_torch_compile_for_training`
     _torch_compile_train_cls = FlexOlmoForCausalLM if is_torch_available() else None
 
-    @unittest.skip("Dynamic control flow in MoE")
-    @pytest.mark.torch_compile_test
-    def test_torch_compile_for_training(self):
-        pass
-
 
 @require_torch
 class FlexOlmoIntegrationTest(unittest.TestCase):

filename: tests/models/qwen2_moe/test_modeling_qwen2_moe.py
patch: @@ -119,7 +119,7 @@ class Qwen2MoeIntegrationTest(unittest.TestCase):
     def get_model(cls):
         if cls.model is None:
             cls.model = Qwen2MoeForCausalLM.from_pretrained(
-                "Qwen/Qwen1.5-MoE-A2.7B", device_map="auto", dtype=torch.float16
+                "Qwen/Qwen1.5-MoE-A2.7B", device_map="auto", dtype=torch.float16, experts_implementation="eager"
             )
         return cls.model