Fix test_model_parallelism (#25359)

* fix

* fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-08-08 10:48:45 +02:00
committed by GitHub
parent d4bd33cc9f
commit 6ea3ee3cd2
25 changed files with 37 additions and 15 deletions

View File

@@ -2597,7 +2597,7 @@ class ModelTesterMixin:
model_size = compute_module_sizes(model)[""]
# We test several splits of sizes to make sure it works.
max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents[1:]]
with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir)