Repository: huggingface/transformers
Pull request: Fix failed unit tests for moonshine_streaming model

@ydshieh , pls help review, thx!

Changed files:
filename: src/transformers/models/moonshine_streaming/modeling_moonshine_streaming.py
patch: @@ -822,7 +822,7 @@ def forward(
         position_embeddings = self.pos_emb(
             torch.arange(encoder_hidden_states.shape[1], device=encoder_hidden_states.device)
         )
-        encoder_hidden_states += position_embeddings
+        encoder_hidden_states += position_embeddings.to(encoder_hidden_states.device)
         encoder_hidden_states = self.proj(encoder_hidden_states)
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

filename: src/transformers/models/moonshine_streaming/modular_moonshine_streaming.py
patch: @@ -376,7 +376,7 @@ def forward(
         position_embeddings = self.pos_emb(
             torch.arange(encoder_hidden_states.shape[1], device=encoder_hidden_states.device)
         )
-        encoder_hidden_states += position_embeddings
+        encoder_hidden_states += position_embeddings.to(encoder_hidden_states.device)
         encoder_hidden_states = self.proj(encoder_hidden_states)
 
         return super().forward(

filename: tests/models/moonshine_streaming/test_modeling_moonshine_streaming.py
patch: @@ -17,7 +17,7 @@
 import unittest
 
 from transformers import MoonshineStreamingConfig, MoonshineStreamingEncoderConfig, is_torch_available
-from transformers.testing_utils import cleanup, require_torch, slow, torch_device
+from transformers.testing_utils import Expectations, cleanup, require_torch, slow, torch_device
 
 from ...test_configuration_common import ConfigTester
 from ...test_modeling_common import (
@@ -162,6 +162,10 @@ def test_can_init_all_missing_weights(self):
     def test_init_weights_can_init_buffers(self):
         self.skipTest("MoonshineStreaming uses special buffer initialization that conflicts with this test")
 
+    # MoonshineStreaming test model's pos_emb layer is large relative to total model size,
+    # requiring higher GPU allocation percentage to enable device splitting
+    model_split_percents = [0.5, 0.9, 0.95]
+
     def test_attention_outputs(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.return_dict = True
@@ -544,14 +548,26 @@ def test_small_logits_batch(self):
         outputs = model.generate(**inputs, max_new_tokens=1, return_dict_in_generate=True, output_logits=True)
 
         # fmt: off
-        EXPECTED_LOGITS = torch.tensor(
-            [
-                [-9.596293449401855, -1.297331690788269, 2.817121982574463, -9.826224327087402, -9.802359580993652, -9.802471160888672, -9.81285285949707, -9.82018756866455, -9.801692962646484, -9.809906005859375],
-                [-9.602995872497559, 0.32756108045578003, 3.0864665508270264, -9.754168510437012, -9.803014755249023, -9.832489013671875, -9.785274505615234, -9.750894546508789, -9.827933311462402, -9.816366195678711],
-                [-10.247313499450684, -0.4231721254699707, 3.1179518699645996, -9.989541053771973, -10.001238822937012, -10.040529251098633, -9.996538162231445, -10.052029609680176, -9.986088752746582, -10.036115646362305],
-                [-9.98245906829834, -1.4063411709259033, 3.539100170135498, -9.433758735656738, -9.444565773010254, -9.49752426147461, -9.452383995056152, -9.457331657409668, -9.432816505432129, -9.439447402954102],
-            ]
-        )
+        EXPECTED_LOGITS = Expectations(
+            {
+                ("xpu", None): torch.tensor(
+                    [
+                        [-9.5963, -1.2973, 2.8171, -9.8262, -9.8024, -9.8025, -9.8129, -9.8202, -9.8017, -9.8099],
+                        [-9.6030, 0.3276, 3.0865, -9.7542, -9.8030, -9.8325, -9.7853, -9.7509, -9.8279, -9.8164],
+                        [-10.2473, -0.4238, 3.1180, -9.9895, -10.0012, -10.0405, -9.9965, -10.0520, -9.9861, -10.0361],
+                        [-9.9825, -1.4059, 3.5391, -9.4338, -9.4446, -9.4975, -9.4524, -9.4573, -9.4328, -9.4394],
+                    ]
+                ),
+                (None, None): torch.tensor(
+                    [
+                        [-9.5963, -1.2973, 2.8171, -9.8262, -9.8024, -9.8025, -9.8129, -9.8202, -9.8017, -9.8099],
+                        [-9.6030, 0.3276, 3.0865, -9.7542, -9.8030, -9.8325, -9.7853, -9.7509, -9.8279, -9.8164],
+                        [-10.2473, -0.4232, 3.1180, -9.9895, -10.0012, -10.0405, -9.9965, -10.0520, -9.9861, -10.0361],
+                        [-9.9825, -1.4063, 3.5391, -9.4338, -9.4446, -9.4975, -9.4524, -9.4573, -9.4328, -9.4394],
+                    ]
+                ),
+            }
+        ).get_expectation()
         # fmt: on
         torch.testing.assert_close(outputs.logits[0][:, :10].cpu(), EXPECTED_LOGITS, rtol=2e-4, atol=2e-4)
 
@@ -565,14 +581,26 @@ def test_medium_logits_batch(self):
         outputs = model.generate(**inputs, max_new_tokens=1, return_dict_in_generate=True, output_logits=True)
 
         # fmt: off
-        EXPECTED_LOGITS = torch.tensor(
-            [
-                [-9.423518180847168, -1.6021490097045898, 1.3190011978149414, -10.032197952270508, -10.08576774597168, -10.04221248626709, -10.057312965393066, -10.089818000793457, -10.141901969909668, -10.003352165222168],
-                [-9.891376495361328, -2.268763542175293, 2.4474310874938965, -10.193374633789062, -10.256990432739258, -10.184536933898926, -10.223142623901367, -10.29221248626709, -10.325952529907227, -10.256648063659668],
-                [-9.396651268005371, -0.7291030287742615, 2.299491403982544, -9.815659523010254, -9.854050636291504, -9.821599006652832, -9.81181812286377, -9.838842391967773, -9.854424476623535, -9.855895042419434],
-                [-8.918790817260742, -0.6990604400634766, 1.3242177963256836, -8.931782722473145, -9.016800880432129, -8.92956829071045, -8.945950508117676, -8.984317779541016, -8.983695030212402, -8.945679664611816],
-            ]
-        )
+        EXPECTED_LOGITS = Expectations(
+            {
+                ("xpu", None): torch.tensor(
+                    [
+                        [-9.4235, -1.6021, 1.3190, -10.0322, -10.0858, -10.0422, -10.0573, -10.0898, -10.1419, -10.0034],
+                        [-9.8914, -2.2688, 2.4474, -10.1934, -10.2570, -10.1845, -10.2231, -10.2922, -10.3259, -10.2566],
+                        [-9.3967, -0.7291, 2.2988, -9.8157, -9.8541, -9.8216, -9.8118, -9.8388, -9.8544, -9.8559],
+                        [-8.9188, -0.6991, 1.3242, -8.9318, -9.0168, -8.9296, -8.9460, -8.9843, -8.9837, -8.9457],
+                    ]
+                ),
+                (None, None): torch.tensor(
+                    [
+                        [-9.4235, -1.6021, 1.3190, -10.0322, -10.0858, -10.0422, -10.0573, -10.0898, -10.1419, -10.0034],
+                        [-9.8914, -2.2688, 2.4474, -10.1934, -10.2570, -10.1845, -10.2231, -10.2922, -10.3260, -10.2566],
+                        [-9.3967, -0.7291, 2.2995, -9.8157, -9.8541, -9.8216, -9.8118, -9.8388, -9.8544, -9.8559],
+                        [-8.9188, -0.6991, 1.3242, -8.9318, -9.0168, -8.9296, -8.9460, -8.9843, -8.9837, -8.9457],
+                    ]
+                ),
+            }
+        ).get_expectation()
         # fmt: on
         torch.testing.assert_close(outputs.logits[0][:, :10].cpu(), EXPECTED_LOGITS, rtol=2e-4, atol=2e-4)