@@ -691,6 +691,9 @@ def test_deprecated_kwargs(self):
691
691
def test_cpu_offload (self ):
692
692
config , inputs_dict = self .prepare_init_args_and_inputs_for_common ()
693
693
model = self .model_class (** config ).eval ()
694
+ if model ._no_split_modules is None :
695
+ return
696
+
694
697
model = model .to (torch_device )
695
698
696
699
torch .manual_seed (0 )
@@ -718,6 +721,9 @@ def test_cpu_offload(self):
718
721
def test_disk_offload_without_safetensors (self ):
719
722
config , inputs_dict = self .prepare_init_args_and_inputs_for_common ()
720
723
model = self .model_class (** config ).eval ()
724
+ if model ._no_split_modules is None :
725
+ return
726
+
721
727
model = model .to (torch_device )
722
728
723
729
torch .manual_seed (0 )
@@ -728,12 +734,12 @@ def test_disk_offload_without_safetensors(self):
728
734
model .cpu ().save_pretrained (tmp_dir , safe_serialization = False )
729
735
730
736
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 )
732
738
max_memory = {0 : max_size , "cpu" : max_size }
733
739
# This errors out because it's missing an offload folder
734
740
new_model = self .model_class .from_pretrained (tmp_dir , device_map = "auto" , max_memory = max_memory )
735
741
736
- max_size = int (self .model_split_percents [1 ] * model_size )
742
+ max_size = int (self .model_split_percents [0 ] * model_size )
737
743
max_memory = {0 : max_size , "cpu" : max_size }
738
744
new_model = self .model_class .from_pretrained (
739
745
tmp_dir , device_map = "auto" , max_memory = max_memory , offload_folder = tmp_dir
@@ -749,6 +755,9 @@ def test_disk_offload_without_safetensors(self):
749
755
def test_disk_offload_with_safetensors (self ):
750
756
config , inputs_dict = self .prepare_init_args_and_inputs_for_common ()
751
757
model = self .model_class (** config ).eval ()
758
+ if model ._no_split_modules is None :
759
+ return
760
+
752
761
model = model .to (torch_device )
753
762
754
763
torch .manual_seed (0 )
@@ -758,7 +767,7 @@ def test_disk_offload_with_safetensors(self):
758
767
with tempfile .TemporaryDirectory () as tmp_dir :
759
768
model .cpu ().save_pretrained (tmp_dir )
760
769
761
- max_size = int (self .model_split_percents [1 ] * model_size )
770
+ max_size = int (self .model_split_percents [0 ] * model_size )
762
771
max_memory = {0 : max_size , "cpu" : max_size }
763
772
new_model = self .model_class .from_pretrained (
764
773
tmp_dir , device_map = "auto" , offload_folder = tmp_dir , max_memory = max_memory
@@ -774,6 +783,9 @@ def test_disk_offload_with_safetensors(self):
774
783
def test_model_parallelism (self ):
775
784
config , inputs_dict = self .prepare_init_args_and_inputs_for_common ()
776
785
model = self .model_class (** config ).eval ()
786
+ if model ._no_split_modules is None :
787
+ return
788
+
777
789
model = model .to (torch_device )
778
790
779
791
torch .manual_seed (0 )
0 commit comments