SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix more offload edge cases (#25342)

Browse Source 
* fix

* fix

* fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-08-07 17:45:41 +02:00
committed by GitHub
parent 7d65697da7
commit c177606fb4
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/test_modeling_common.py
View File

@@ -2521,15 +2521,17 @@ class ModelTesterMixin:
base_output = model(**inputs_dict_class)
model_size = compute_module_sizes(model)[""]
max_size = int(self.model_split_percents[1] * model_size)
with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir)
max_memory = {0: max_size, "cpu": max_size}
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 cause it's missing an offload folder
new_model = model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory)
max_size = int(self.model_split_percents[1] * model_size)
max_memory = {0: max_size, "cpu": max_size}
new_model = model_class.from_pretrained(
tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir
)
@@ -2559,7 +2561,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)