Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 4 additions & 5 deletions tests/models/test_modeling_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -1169,17 +1169,16 @@ def test_disk_offload_without_safetensors(self):
base_output = model(**inputs_dict)

model_size = compute_module_sizes(model)[""]
max_size = int(self.model_split_percents[0] * model_size)
# Force disk offload by setting very small CPU memory
max_memory = {0: max_size, "cpu": int(0.1 * max_size)}

with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir, safe_serialization=False)

with self.assertRaises(ValueError):
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[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
)
Expand Down
5 changes: 3 additions & 2 deletions tests/models/transformers/test_models_transformer_omnigen.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ class OmniGenTransformerTests(ModelTesterMixin, unittest.TestCase):
model_class = OmniGenTransformer2DModel
main_input_name = "hidden_states"
uses_custom_attn_processor = True
model_split_percents = [0.1, 0.1, 0.1]

@property
def dummy_input(self):
Expand Down Expand Up @@ -73,9 +74,9 @@ def prepare_init_args_and_inputs_for_common(self):
"num_attention_heads": 4,
"num_key_value_heads": 4,
"intermediate_size": 32,
"num_layers": 1,
"num_layers": 20,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn’t it a bit too much for CPUs?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall model size is still very small (2.6 MB). Adds 2 seconds to the current CPU tests.

"pad_token_id": 0,
"vocab_size": 100,
"vocab_size": 1000,
"in_channels": 4,
"time_step_dim": 4,
"rope_scaling": {"long_factor": list(range(1, 3)), "short_factor": list(range(1, 3))},
Expand Down
6 changes: 4 additions & 2 deletions tests/models/transformers/test_models_transformer_sd3.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
class SD3TransformerTests(ModelTesterMixin, unittest.TestCase):
model_class = SD3Transformer2DModel
main_input_name = "hidden_states"
model_split_percents = [0.8, 0.8, 0.9]

@property
def dummy_input(self):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.8, 0.8, 0.9]

@property
def dummy_input(self):
Expand Down Expand Up @@ -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": 4,
"attention_head_dim": 8,
"num_attention_heads": 4,
"caption_projection_dim": 32,
Expand Down