From c177606fb47fd15109df0025ba15bdf6fe1f8bea Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Mon, 7 Aug 2023 17:45:41 +0200 Subject: [PATCH] Fix more offload edge cases (#25342) * fix * fix * fix --------- Co-authored-by: ydshieh --- tests/models/longformer/test_modeling_longformer.py | 3 +++ tests/test_modeling_common.py | 8 +++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/models/longformer/test_modeling_longformer.py b/tests/models/longformer/test_modeling_longformer.py index b40e464e60..7edcd206ab 100644 --- a/tests/models/longformer/test_modeling_longformer.py +++ b/tests/models/longformer/test_modeling_longformer.py @@ -326,6 +326,9 @@ class LongformerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCa else {} ) + # Need to use `0.6` instead of `0.5` for `test_disk_offload` + model_split_percents = [0.6, 0.7, 0.9] + # TODO: Fix the failed tests def is_pipeline_test_to_skip( self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index fe2e2bc819..1850ae8e5b 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -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)