Repository: huggingface/transformers
Pull request: fix device dismatch issue for pe_audio_video model parallelism


Changed files:
filename: src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
patch: @@ -86,7 +86,8 @@ class TimmWrapperPreTrainedModel(PreTrainedModel):
     main_input_name = "pixel_values"
     input_modalities = ("image",)
     config: TimmWrapperConfig
-    _no_split_modules = []
+    # add WA here as `timm` does not support model parallelism
+    _no_split_modules = ["TimmWrapperModel"]
     model_tags = ["timm"]
 
     # used in Trainer to avoid passing `loss_kwargs` to model forward

filename: tests/models/pe_audio/test_modeling_pe_audio.py
patch: @@ -101,7 +101,8 @@ def seq_length(self):
 
     def prepare_config_and_inputs(self):
         input_values = floats_tensor([self.batch_size, self.num_channels, self.audio_seq_length])
-        valid_lengths = ids_tensor([self.batch_size], self.audio_seq_length)
+        # Generate valid_lengths in range [1, self.audio_seq_length] to ensure at least one valid frame
+        valid_lengths = ids_tensor([self.batch_size], self.audio_seq_length - 1) + 1
         padding_mask = torch.arange(self.audio_seq_length, device=torch_device)[None, :] < valid_lengths[:, None]
         padding_mask = padding_mask.int()
         config = self.get_config()

filename: tests/models/pe_audio_video/test_modeling_pe_audio_video.py
patch: @@ -149,7 +149,8 @@ def seq_length(self):
 
     def prepare_config_and_inputs(self):
         input_values = floats_tensor([self.batch_size, self.num_audio_channels, self.audio_seq_length])
-        valid_audio_lengths = ids_tensor([self.batch_size], self.audio_seq_length)
+        # Generate valid_lengths in range [1, self.audio_seq_length] to ensure at least one valid frame
+        valid_audio_lengths = ids_tensor([self.batch_size], self.audio_seq_length - 1) + 1
         padding_mask = torch.arange(self.audio_seq_length, device=torch_device)[None, :] < valid_audio_lengths[:, None]
         padding_mask = padding_mask.int()
 
@@ -162,7 +163,8 @@ def prepare_config_and_inputs(self):
                 self.config_kwargs["video_config"]["vision_config"]["model_args"]["img_size"][1],
             ]
         )
-        valid_video_lengths = ids_tensor([self.batch_size], self.num_frames)
+        # Generate valid_lengths in range [1, self.num_frames] to ensure at least one valid frame
+        valid_video_lengths = ids_tensor([self.batch_size], self.num_frames - 1) + 1
         padding_mask_videos = (
             torch.arange(self.num_frames, device=torch_device)[None, :] < valid_video_lengths[:, None]
         )
@@ -258,6 +260,10 @@ def test_feed_forward_chunking(self):
     def test_save_load(self):
         pass
 
+    @unittest.skip(reason="TimmWrapperModel does not support model parallelism")
+    def test_model_parallelism(self):
+        pass
+
     @unittest.skip(reason="@eustlb this is not really expected")
     def test_batching_equivalence(self):
         pass

filename: tests/models/pe_video/test_modeling_pe_video.py
patch: @@ -103,10 +103,9 @@ def prepare_config_and_inputs(self):
                 self.config_kwargs["vision_config"]["model_args"]["img_size"][1],
             ]
         )
-        valid_lengths = ids_tensor([self.batch_size], self.num_frames)
-        padding_mask_videos = (
-            torch.ones([self.batch_size, self.num_frames], device=torch_device) < valid_lengths[:, None]
-        )
+        # Generate valid_lengths in range [1, num_frames] to ensure at least one valid frame
+        valid_lengths = ids_tensor([self.batch_size], self.num_frames - 1) + 1
+        padding_mask_videos = torch.arange(self.num_frames, device=torch_device).unsqueeze(0) < valid_lengths[:, None]
         padding_mask_videos = padding_mask_videos.int()
         config = self.get_config()
 
@@ -187,6 +186,10 @@ def test_feed_forward_chunking(self):
     def test_save_load(self):
         pass
 
+    @unittest.skip(reason="TimmWrapperModel does not support model parallelism")
+    def test_model_parallelism(self):
+        pass
+
     @unittest.skip(reason="@eustlb this is not really expected")
     def test_batching_equivalence(self):
         pass