Repository: huggingface/transformers
Pull request: Fix test_all_tensors_are_parameter_or_buffer

# What does this PR do?

This PR fixes the test `test_all_tensors_are_parameter_or_buffer` across 5 models, effectively making the test pass for all models. 

## Failure cases and fixes

There are 2 reasons this test used to fail:
- the model has some triton code which crash when we try to do a forward pass on the CPU:
To address this, we add a keyword arg `non_meta_device` which is the device the model is transferred to. It defaults to `cpu`, so the default test behavior does not change. If set to `cuda`, the model is transferred to the GPU and the triton code can run.
- the model has inputs with nested tensors (list, dict, or objects with tensor attibutes):
 To address this, we add the `transfer_nested_tensors` flag, turned off by default so the default test behavior does not change. If the flag is passed, we go through inputs and try to transfer nested tensors whether they are part of a list, tuple, dict or plain attributes. 

## Models affected

No modeling code was changed, but the test was overloaded with kwargs in the following models:
- layoutlmv2 (the `images` input has a nested tensor attribute)
- modernbert (cannot run on cpu because of triton)
- pe_audio (cannot run on cpu because of triton)
- pe_video (cannot run on cpu because of triton)
- perceiver (the model has nested tensors in its inputs)

Changed files:
filename: tests/models/modernbert/test_modeling_modernbert.py
patch: @@ -27,6 +27,7 @@
     require_flash_attn,
     require_torch,
     require_torch_accelerator,
+    require_torch_gpu,
     slow,
     torch_device,
 )
@@ -506,6 +507,10 @@ def flash_attn_inference_equivalence(
                 f"Model architecture does not support {attn_implementation}, or setting its attention dynamically"
             )
 
+    @require_torch_gpu  # modernbert contains triton code which cannot run on CPU, so we only test on GPU
+    def test_all_tensors_are_parameter_or_buffer(self):
+        super().test_all_tensors_are_parameter_or_buffer()
+
 
 @require_torch
 class ModernBertModelIntegrationTest(unittest.TestCase):

filename: tests/models/pe_audio/test_modeling_pe_audio.py
patch: @@ -17,6 +17,7 @@
 from transformers.audio_utils import load_audio
 from transformers.testing_utils import (
     require_torch,
+    require_torch_gpu,
     slow,
     torch_device,
 )
@@ -330,6 +331,10 @@ def test_batching_equivalence(self):
     def test_can_init_all_missing_weights(self):
         pass
 
+    @require_torch_gpu  # pe-audio contains triton code which cannot run on CPU, so we only test on GPU
+    def test_all_tensors_are_parameter_or_buffer(self):
+        super().test_all_tensors_are_parameter_or_buffer()
+
 
 @require_torch
 class PeAudioIntegrationTest(unittest.TestCase):

filename: tests/models/pe_video/test_modeling_pe_video.py
patch: @@ -27,6 +27,7 @@
     floats_tensor,
     ids_tensor,
     random_attention_mask,
+    require_torch_gpu,
 )
 
 
@@ -362,6 +363,10 @@ def test_batching_equivalence(self):
     def test_can_init_all_missing_weights(self):
         pass
 
+    @require_torch_gpu  # pe-video contains triton code which cannot run on CPU, so we only test on GPU
+    def test_all_tensors_are_parameter_or_buffer(self):
+        super().test_all_tensors_are_parameter_or_buffer()
+
 
 @require_torch
 class PeVideoIntegrationTest(unittest.TestCase):

filename: tests/test_modeling_common.py
patch: @@ -1286,30 +1286,28 @@ def test_init_weights_can_init_buffers(self):
                 f"them correctly if the model is on meta device):\n{unique_bad_module_traceback}",
             )
 
-    def test_all_tensors_are_parameter_or_buffer(self):
+    def test_all_tensors_are_parameter_or_buffer(self) -> None:
         """Check that all tensors are registered as Parameter or Buffer, i.e. we don't have simple assignments such
-        as `self.x = torch.tensor(...)` in a Module (as we cannot correctly recover from meta device if it's not registered as
-        parameter/buffer)"""
+        as `self.x = torch.tensor(...)` in a Module (as we cannot correctly recover from meta device if it's not
+        registered as parameter/buffer). To test this, we initialize the model on a meta device and then move it onto
+        the torch_device and perform a forward pass."""
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
 
         for model_class in self.all_model_classes:
             # Apparently this model cannot correctly create its inputs and has to use another function....
             if "modeling_perceiver.py" in inspect.getfile(model_class):
                 _, inputs_dict = self.model_tester.prepare_config_and_inputs_for_model_class(model_class)
 
-            # Initialize the model fully on meta device, then move everything to cpu and run `init_weights`
+            # Initialize the model fully on meta device, then move everything to torch_device and run `init_weights`
             with torch.device("meta"):
                 model = model_class(copy.deepcopy(config)).eval()
-            # Move everything randomly to cpu
-            model.to_empty(device="cpu")
+            # Move everything randomly to torch_device
+            model.to_empty(device=torch_device)
             # Now, run all the inits
             model.init_weights()
 
             # Prepare inputs
             inputs = self._prepare_for_class(inputs_dict, model_class)
-            # Inputs may be on cuda -> move to cpu, we don't care about accelerator for this test
-            inputs = {k: v.to("cpu") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
-
             # Try running a forward, to see if a tensor stayed on meta somewhere
             try:
                 _ = model(**inputs)