Skip to content

Commit 8909ab4

Browse files
authored
[Tests] fix: device map tests for models (#7825)
* fix: device module tests * remove patch file * Empty-Commit
1 parent c1edb03 commit 8909ab4

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

tests/models/test_modeling_common.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,9 @@ def test_deprecated_kwargs(self):
691691
def test_cpu_offload(self):
692692
config, inputs_dict = self.prepare_init_args_and_inputs_for_common()
693693
model = self.model_class(**config).eval()
694+
if model._no_split_modules is None:
695+
return
696+
694697
model = model.to(torch_device)
695698

696699
torch.manual_seed(0)
@@ -718,6 +721,9 @@ def test_cpu_offload(self):
718721
def test_disk_offload_without_safetensors(self):
719722
config, inputs_dict = self.prepare_init_args_and_inputs_for_common()
720723
model = self.model_class(**config).eval()
724+
if model._no_split_modules is None:
725+
return
726+
721727
model = model.to(torch_device)
722728

723729
torch.manual_seed(0)
@@ -728,12 +734,12 @@ def test_disk_offload_without_safetensors(self):
728734
model.cpu().save_pretrained(tmp_dir, safe_serialization=False)
729735

730736
with self.assertRaises(ValueError):
731-
max_size = int(self.model_split_percents[1] * model_size)
737+
max_size = int(self.model_split_percents[0] * model_size)
732738
max_memory = {0: max_size, "cpu": max_size}
733739
# This errors out because it's missing an offload folder
734740
new_model = self.model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory)
735741

736-
max_size = int(self.model_split_percents[1] * model_size)
742+
max_size = int(self.model_split_percents[0] * model_size)
737743
max_memory = {0: max_size, "cpu": max_size}
738744
new_model = self.model_class.from_pretrained(
739745
tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir
@@ -749,6 +755,9 @@ def test_disk_offload_without_safetensors(self):
749755
def test_disk_offload_with_safetensors(self):
750756
config, inputs_dict = self.prepare_init_args_and_inputs_for_common()
751757
model = self.model_class(**config).eval()
758+
if model._no_split_modules is None:
759+
return
760+
752761
model = model.to(torch_device)
753762

754763
torch.manual_seed(0)
@@ -758,7 +767,7 @@ def test_disk_offload_with_safetensors(self):
758767
with tempfile.TemporaryDirectory() as tmp_dir:
759768
model.cpu().save_pretrained(tmp_dir)
760769

761-
max_size = int(self.model_split_percents[1] * model_size)
770+
max_size = int(self.model_split_percents[0] * model_size)
762771
max_memory = {0: max_size, "cpu": max_size}
763772
new_model = self.model_class.from_pretrained(
764773
tmp_dir, device_map="auto", offload_folder=tmp_dir, max_memory=max_memory
@@ -774,6 +783,9 @@ def test_disk_offload_with_safetensors(self):
774783
def test_model_parallelism(self):
775784
config, inputs_dict = self.prepare_init_args_and_inputs_for_common()
776785
model = self.model_class(**config).eval()
786+
if model._no_split_modules is None:
787+
return
788+
777789
model = model.to(torch_device)
778790

779791
torch.manual_seed(0)

0 commit comments

Comments
 (0)