Repository: huggingface/transformers
Pull request: [`Moe`] Post interface fixes

As per title, a few tests fail with `grouped_mm` implementation. Way less integration tests fail than I initially expected. 
Related failures: https://huggingface.co/datasets/hf-internal-testing/transformers_daily_ci/raw/3fed21c3fdb63a1b19d6030d493139a947135fd3/2026-01-06/ci_results_run_models_gpu/new_failures_with_bad_commit_grouped_by_authors.json (note that there are still some offloading issues but they are unrelated to moe but a buffer issue)

Kind of a general follow up to #43117 and #42697

Changed files:
filename: tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py
patch: @@ -152,6 +152,7 @@ def get_large_model(cls):
         cls.model = Ernie4_5_MoeForCausalLM.from_pretrained(
             "baidu/ERNIE-4.5-21B-A3B-PT",
             device_map="auto",
+            experts_implementation="eager",
             quantization_config=BitsAndBytesConfig(load_in_4bit=True),
         )
 
@@ -163,6 +164,7 @@ def get_small_model(cls):
             "hf-internal-testing/ERNIE-4.5-Small-Moe",
             device_map="auto",
             dtype="auto",
+            experts_implementation="eager",
         )
 
         return cls.model

filename: tests/models/lfm2_moe/test_modeling_lfm2_moe.py
patch: @@ -164,7 +164,10 @@ def tearDown(self):
     def get_model(cls):
         if cls.model is None:
             cls.model = Lfm2MoeForCausalLM.from_pretrained(
-                "LiquidAI/LFM2-8B-A1B", device_map="auto", dtype=torch.bfloat16
+                "LiquidAI/LFM2-8B-A1B",
+                device_map="auto",
+                dtype=torch.bfloat16,
+                experts_implementation="eager",
             )
         return cls.model
 

filename: tests/test_modeling_common.py
patch: @@ -3566,7 +3566,7 @@ def test_sdpa_can_compile_dynamic(self):
 
             with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_pretrained(tmpdirname)
-                model = model_class.from_pretrained(tmpdirname, dtype=torch.float16, attn_implementation="sdpa")
+                model = model_class.from_pretrained(tmpdirname, dtype=torch.bfloat16, attn_implementation="sdpa")
                 model.to(torch_device)
 
                 # For PyTorch 2.1 - 2.3.0 set `dynamic=True`. In the future setting `dynamic=None` and using `torch._dynamo.mark_dynamic()`
@@ -3577,7 +3577,7 @@ def test_sdpa_can_compile_dynamic(self):
                 inputs_dict.pop("decoder_attention_mask", None)
                 for name, inp in inputs_dict.items():
                     if isinstance(inp, torch.Tensor) and inp.dtype in [torch.float32, torch.float16]:
-                        inputs_dict[name] = inp.to(torch.float16)
+                        inputs_dict[name] = inp.to(torch.bfloat16)
 
                 # use no_grad to save some memory
                 with torch.no_grad():
@@ -3914,7 +3914,12 @@ def test_torch_compile_for_training(self):
         if attn_implementation is not None:
             config._attn_implementation = attn_implementation
 
-        model = cls(config).to(torch_device)
+        model = cls(config).to(device=torch_device)
+
+        # torch._grouped_mm still only supports bfloat16 when used with torch.compile
+        # bfloat16 is problematic with precisions so we keep an implementation with full precision
+        if model.config._experts_implementation == "grouped_mm":
+            model.set_experts_implementation("batched_mm")
 
         inputs = {
             "input_ids": torch.randint(low=1, high=model.config.vocab_size, size=(2, 10), device=torch_device),