Repository: huggingface/transformers
Pull request: [testing] Fix `JetMoeIntegrationTest`

# What does this PR do?

> tests/models/jetmoe/test_modeling_jetmoe.py::JetMoeIntegrationTest::test_model_8b_generation 

cause the pytest process being killed (and sometimes hangs for hours), see the log below.

It loads the model on  `cpu`. Let's use "auto".

The outputs changes after #41324 (which is a fix of #40132), so the tests still fail but not causing the pytest process being killed and job hanging. 

There is a new warning showing 

>  The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.


cc @ArthurZucker for this output change issue. 

Log when running before this PR: 

> 2025-09-30T16:33:22.8891725Z tests/models/jetmoe/test_modeling_jetmoe.py::JetMoeIntegrationTest::test_model_8b_batched_generation 
2025-09-30T16:33:22.8892381Z -------------------------------- live log call ---------------------------------
2025-09-30T16:33:22.8893314Z WARNING  transformers.generation.configuration_utils:logging.py:328 The following generation flags are not valid and may be ignored: ['temperature']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
2025-09-30T16:33:23.3206200Z FAILED                                                                   [ 99%]
2025-09-30T16:39:21.4805460Z Killed
2025-09-30T16:39:21.4847256Z tests/models/jetmoe/test_modeling_jetmoe.py::JetMoeIntegrationTest::test_model_8b_generation 

Changed files:
filename: tests/models/jetmoe/test_modeling_jetmoe.py
patch: @@ -13,14 +13,13 @@
 # limitations under the License.
 """Testing suite for the PyTorch JetMoe model."""
 
-import gc
 import unittest
 
 import pytest
 
 from transformers import AutoTokenizer, is_torch_available
 from transformers.testing_utils import (
-    backend_empty_cache,
+    cleanup,
     require_flash_attn,
     require_torch,
     require_torch_gpu,
@@ -127,41 +126,39 @@ def test_flash_attn_2_inference_equivalence_right_padding(self):
 
 @require_torch
 class JetMoeIntegrationTest(unittest.TestCase):
+    def setUp(self):
+        cleanup(torch_device, gc_collect=True)
+
+    def tearDown(self):
+        cleanup(torch_device, gc_collect=True)
+
     @slow
     def test_model_8b_logits(self):
         input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
-        model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b")
+        model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b", device_map="auto", torch_dtype=torch.bfloat16)
         input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
         with torch.no_grad():
             out = model(input_ids).logits.float().cpu()
         # Expected mean on dim = -1
-        EXPECTED_MEAN = torch.tensor([[0.2507, -2.7073, -1.3445, -1.9363, -1.7216, -1.7370, -1.9054, -1.9792]])
+        EXPECTED_MEAN = torch.tensor([[0.1943, -2.7299, -1.3466, -1.9385, -1.7457, -1.7472, -1.8647, -1.8547]])
         torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, rtol=1e-2, atol=1e-2)
         # slicing logits[0, 0, 0:30]
-        EXPECTED_SLICE = torch.tensor([-3.3689,  5.9006,  5.7450, -1.7012, -4.7072, -4.7071, -4.7071, -4.7071, -4.7072, -4.7072, -4.7072, -4.7071,  3.8321,  9.1746, -4.7071, -4.7072, -4.7071, -4.7072, -4.7071, -4.7072, -4.7071, -4.7071, -4.7071, -4.7071, -4.7071, -4.7071, -4.7071, -4.7071, -4.7071, -4.7071])  # fmt: skip
+        EXPECTED_SLICE = torch.tensor([-3.4844, 6.0625, 5.8750, -1.6875, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, 3.8750, 9.3750, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812, -4.7812])  # fmt: skip
         torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, rtol=1e-4, atol=1e-4)
 
-        del model
-        backend_empty_cache(torch_device)
-        gc.collect()
-
     @slow
     def test_model_8b_generation(self):
         EXPECTED_TEXT_COMPLETION = """My favourite condiment is ....\nI love ketchup. I love"""
         prompt = "My favourite condiment is "
         tokenizer = AutoTokenizer.from_pretrained("jetmoe/jetmoe-8b", use_fast=False)
-        model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b")
+        model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b", device_map="auto", torch_dtype=torch.bfloat16)
         input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.model.embed_tokens.weight.device)
 
         # greedy generation outputs
         generated_ids = model.generate(input_ids, max_new_tokens=10, temperature=0)
         text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
         self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
 
-        del model
-        backend_empty_cache(torch_device)
-        gc.collect()
-
     @slow
     def test_model_8b_batched_generation(self):
         EXPECTED_TEXT_COMPLETION = [
@@ -173,14 +170,10 @@ def test_model_8b_batched_generation(self):
             "My favourite ",
         ]
         tokenizer = AutoTokenizer.from_pretrained("jetmoe/jetmoe-8b", use_fast=False)
-        model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b")
+        model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b", device_map="auto", torch_dtype=torch.bfloat16)
         input_ids = tokenizer(prompt, return_tensors="pt", padding=True).to(model.model.embed_tokens.weight.device)
 
         # greedy generation outputs
         generated_ids = model.generate(**input_ids, max_new_tokens=10, temperature=0)
         text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
         self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
-
-        del model
-        backend_empty_cache(torch_device)
-        gc.collect()