From 00bf44270f9def905af70ee994c290adc12ef2cb Mon Sep 17 00:00:00 2001 From: Fanli Lin Date: Wed, 6 Mar 2024 10:58:42 +0800 Subject: [PATCH] [FIX] `offload_weight()` takes from 3 to 4 positional arguments but 5 were given (#29457) * use require_torch_gpu * enable on XPU * fix --- src/transformers/modeling_utils.py | 2 +- tests/test_modeling_utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index b542307794168..5aa9d0a770cfa 100644 --- 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 index a334cb0f2853b..57f0f11dbb8a0 100755 --- 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")