[FIX] offload_weight() takes from 3 to 4 positional arguments but 5…

… were given (#29457) * use require_torch_gpu * enable on XPU * fix
huggingface · Mar 6, 2024 · 00bf442 · 00bf442
1 parent 7b01579
commit 00bf442
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -796,7 +796,7 @@ def _load_state_dict_into_meta_model(
             if not is_safetensors:
                 offload_index = offload_weight(param, param_name, offload_folder, offload_index)
         elif param_device == "cpu" and state_dict_index is not None:
-            state_dict_index = offload_weight(param, param_name, model, state_dict_folder, state_dict_index)
+            state_dict_index = offload_weight(param, param_name, state_dict_folder, state_dict_index)
         elif (
             hf_quantizer is None
             or (not hf_quantizer.requires_parameters_quantization)

diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
@@ -765,7 +765,7 @@ def test_model_parallelism_gpt2(self):
 
     @require_accelerate
     @mark.accelerate_tests
-    @require_torch_accelerator
+    @require_torch_gpu
     def test_from_pretrained_disk_offload_task_model(self):
         model = AutoModel.from_pretrained("hf-internal-testing/tiny-random-gpt2")
         device_map = {
@@ -808,7 +808,7 @@ def test_from_pretrained_disk_offload_task_model(self):
 
     @require_accelerate
     @mark.accelerate_tests
-    @require_torch_accelerator
+    @require_torch_gpu
     def test_from_pretrained_disk_offload_derived_to_base_model(self):
         derived_model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")