update

Author: DN6

tests/models/test_modeling_common.py

@@ -1179,10 +1179,11 @@ def test_disk_offload_without_safetensors(self):
                 new_model = self.model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory)
 
             max_size = int(self.model_split_percents[0] * model_size)
-            max_memory = {0: max_size, "cpu": max_size}
+            max_memory = {0: max_size, "disk": max_size}
             new_model = self.model_class.from_pretrained(
                 tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir
             )
+            __import__("ipdb").set_trace()
 
             self.check_device_map_is_respected(new_model, new_model.hf_device_map)
             torch.manual_seed(0)

tests/models/transformers/test_models_transformer_sd3.py

@@ -33,6 +33,7 @@
 class SD3TransformerTests(ModelTesterMixin, unittest.TestCase):
     model_class = SD3Transformer2DModel
     main_input_name = "hidden_states"
+    model_split_percents = [0.4, 0.8, 0.9]
 
     @property
     def dummy_input(self):
@@ -67,7 +68,7 @@ def prepare_init_args_and_inputs_for_common(self):
             "sample_size": 32,
             "patch_size": 1,
             "in_channels": 4,
-            "num_layers": 1,
+            "num_layers": 4,
             "attention_head_dim": 8,
             "num_attention_heads": 4,
             "caption_projection_dim": 32,
@@ -107,6 +108,7 @@ def test_gradient_checkpointing_is_applied(self):
 class SD35TransformerTests(ModelTesterMixin, unittest.TestCase):
     model_class = SD3Transformer2DModel
     main_input_name = "hidden_states"
+    model_split_percents = [0.4, 0.8, 0.9]
 
     @property
     def dummy_input(self):
@@ -141,7 +143,7 @@ def prepare_init_args_and_inputs_for_common(self):
             "sample_size": 32,
             "patch_size": 1,
             "in_channels": 4,
-            "num_layers": 2,
+            "num_layers": 5,
             "attention_head_dim": 8,
             "num_attention_heads": 4,
             "caption_projection_dim": 32,