From 8909ab4b192500bcc3c17d839ae101cc669e9d8e Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 1 May 2024 18:45:47 +0530 Subject: [PATCH] [Tests] fix: device map tests for models (#7825) * fix: device module tests * remove patch file * Empty-Commit --- tests/models/test_modeling_common.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index d8a93d40c8..d9e70c6dd7 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -691,6 +691,9 @@ class ModelTesterMixin: def test_cpu_offload(self): config, inputs_dict = self.prepare_init_args_and_inputs_for_common() model = self.model_class(**config).eval() + if model._no_split_modules is None: + return + model = model.to(torch_device) torch.manual_seed(0) @@ -718,6 +721,9 @@ class ModelTesterMixin: def test_disk_offload_without_safetensors(self): config, inputs_dict = self.prepare_init_args_and_inputs_for_common() model = self.model_class(**config).eval() + if model._no_split_modules is None: + return + model = model.to(torch_device) torch.manual_seed(0) @@ -728,12 +734,12 @@ class ModelTesterMixin: model.cpu().save_pretrained(tmp_dir, safe_serialization=False) with self.assertRaises(ValueError): - max_size = int(self.model_split_percents[1] * model_size) + max_size = int(self.model_split_percents[0] * model_size) max_memory = {0: max_size, "cpu": max_size} # This errors out because it's missing an offload folder new_model = self.model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory) - max_size = int(self.model_split_percents[1] * model_size) + max_size = int(self.model_split_percents[0] * model_size) max_memory = {0: max_size, "cpu": max_size} new_model = self.model_class.from_pretrained( tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir @@ -749,6 +755,9 @@ class ModelTesterMixin: def test_disk_offload_with_safetensors(self): config, inputs_dict = self.prepare_init_args_and_inputs_for_common() model = self.model_class(**config).eval() + if model._no_split_modules is None: + return + model = model.to(torch_device) torch.manual_seed(0) @@ -758,7 +767,7 @@ class ModelTesterMixin: with tempfile.TemporaryDirectory() as tmp_dir: model.cpu().save_pretrained(tmp_dir) - max_size = int(self.model_split_percents[1] * model_size) + max_size = int(self.model_split_percents[0] * model_size) max_memory = {0: max_size, "cpu": max_size} new_model = self.model_class.from_pretrained( tmp_dir, device_map="auto", offload_folder=tmp_dir, max_memory=max_memory @@ -774,6 +783,9 @@ class ModelTesterMixin: def test_model_parallelism(self): config, inputs_dict = self.prepare_init_args_and_inputs_for_common() model = self.model_class(**config).eval() + if model._no_split_modules is None: + return + model = model.to(torch_device) torch.manual_seed(0)