From a906f7713fb818611386ea2c6c57fc6806b318e0 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 13 Jun 2025 15:49:53 +0530 Subject: [PATCH 1/2] device_map tests for all models. --- tests/models/test_modeling_common.py | 40 +++++++++++++++++++ .../unets/test_models_unet_2d_condition.py | 37 ----------------- 2 files changed, 40 insertions(+), 37 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 511fa4bfa9ea..b0dad0c522dd 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -1736,6 +1736,46 @@ def test_auto_model(self, expected_max_diff=5e-5): f"AutoModel forward pass diff: {max_diff} exceeds threshold {expected_max_diff}", ) + @parameterized.expand( + [ + (-1, "You can't pass device_map as a negative int"), + ("foo", "When passing device_map as a string, the value needs to be a device name"), + ] + ) + def test_wrong_device_map_raises_error(self, device_map, msg_substring): + init_dict, _ = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + with tempfile.TemporaryDirectory() as tmpdir: + model.save_pretrained(tmpdir) + with self.assertRaises(ValueError) as err_ctx: + _ = self.model_class.from_pretrained( + "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map + ) + + assert msg_substring in str(err_ctx.exception) + + @parameterized.expand([0, "cuda", torch.device("cuda"), torch.device("cuda:0")]) + @require_torch_gpu + def test_passing_non_dict_device_map_works(self, device_map): + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict).eval() + with tempfile.TemporaryDirectory() as tmpdir: + model.save_pretrained(tmpdir) + loaded_model = self.model_class.from_pretrained(tmpdir, device_map=device_map) + _ = loaded_model(**inputs_dict) + + @parameterized.expand([("", "cuda"), ("", torch.device("cuda"))]) + @require_torch_gpu + def test_passing_dict_device_map_works(self, name, device_map): + # There are other valid dict-based `device_map` values too. It's best to refer to + # the docs for those: https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference#the-devicemap. + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict).eval() + with tempfile.TemporaryDirectory() as tmpdir: + model.save_pretrained(tmpdir) + loaded_model = self.model_class.from_pretrained(tmpdir, device_map=device_map) + _ = loaded_model(**inputs_dict) + @is_staging_test class ModelPushToHubTester(unittest.TestCase): diff --git a/tests/models/unets/test_models_unet_2d_condition.py b/tests/models/unets/test_models_unet_2d_condition.py index e0331d15dd04..c8ed68c65b40 100644 --- a/tests/models/unets/test_models_unet_2d_condition.py +++ b/tests/models/unets/test_models_unet_2d_condition.py @@ -46,7 +46,6 @@ require_peft_backend, require_torch_accelerator, require_torch_accelerator_with_fp16, - require_torch_gpu, skip_mps, slow, torch_all_close, @@ -1084,42 +1083,6 @@ def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self): assert loaded_model assert new_output.sample.shape == (4, 4, 16, 16) - @parameterized.expand( - [ - (-1, "You can't pass device_map as a negative int"), - ("foo", "When passing device_map as a string, the value needs to be a device name"), - ] - ) - def test_wrong_device_map_raises_error(self, device_map, msg_substring): - with self.assertRaises(ValueError) as err_ctx: - _ = self.model_class.from_pretrained( - "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map - ) - - assert msg_substring in str(err_ctx.exception) - - @parameterized.expand([0, "cuda", torch.device("cuda"), torch.device("cuda:0")]) - @require_torch_gpu - def test_passing_non_dict_device_map_works(self, device_map): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - loaded_model = self.model_class.from_pretrained( - "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map - ) - output = loaded_model(**inputs_dict) - assert output.sample.shape == (4, 4, 16, 16) - - @parameterized.expand([("", "cuda"), ("", torch.device("cuda"))]) - @require_torch_gpu - def test_passing_dict_device_map_works(self, name, device_map): - # There are other valid dict-based `device_map` values too. It's best to refer to - # the docs for those: https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference#the-devicemap. - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - loaded_model = self.model_class.from_pretrained( - "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map={name: device_map} - ) - output = loaded_model(**inputs_dict) - assert output.sample.shape == (4, 4, 16, 16) - @require_peft_backend def test_load_attn_procs_raise_warning(self): init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() From fdb56cf0fb1e1a7e233fc2c0cfa33d8ebdcb0769 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 13 Jun 2025 15:54:42 +0530 Subject: [PATCH 2/2] updates --- tests/models/test_modeling_common.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index b0dad0c522dd..aab2c7d1ca3f 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -1748,9 +1748,7 @@ def test_wrong_device_map_raises_error(self, device_map, msg_substring): with tempfile.TemporaryDirectory() as tmpdir: model.save_pretrained(tmpdir) with self.assertRaises(ValueError) as err_ctx: - _ = self.model_class.from_pretrained( - "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map - ) + _ = self.model_class.from_pretrained(tmpdir, device_map=device_map) assert msg_substring in str(err_ctx.exception)