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

Fix all offload and MP tests (#17533)

Browse Source
This commit is contained in:
Sylvain Gugger
2022-06-03 09:59:13 -04:00
committed by GitHub
parent 1c57242d7b
commit 8343901263
4 changed files with 9 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
tests/test_modeling_common.py
View File

@@ -153,6 +153,7 @@ class ModelTesterMixin:
test_model_parallel = False
is_encoder_decoder = False
has_attentions = True
model_split_percents = [0.5, 0.7, 0.9]
def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
inputs_dict = copy.deepcopy(inputs_dict)
@@ -2217,12 +2218,7 @@ class ModelTesterMixin:
@require_accelerate
@require_torch_gpu
def test_disk_offload(self):
if all([model_class._no_split_modules is None for model_class in self.all_model_classes]):
return
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
if isinstance(getattr(config, "num_hidden_layers", None), int) and config.num_hidden_layers < 4:
config.num_hidden_layers = 4
for model_class in self.all_model_classes:
if model_class._no_split_modules is None:
@@ -2234,8 +2230,7 @@ class ModelTesterMixin:
base_output = model(**inputs_dict)
model_size = compute_module_sizes(model)[""]
# We test several splits of sizes to make sure it works.
max_size = int(0.4 * model_size)
max_size = int(self.model_split_percents[0] * model_size)
with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir)
@@ -2256,12 +2251,7 @@ class ModelTesterMixin:
@require_accelerate
@require_torch_gpu
def test_cpu_offload(self):
if all([model_class._no_split_modules is None for model_class in self.all_model_classes]):
return
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
if isinstance(getattr(config, "num_hidden_layers", None), int) and config.num_hidden_layers < 4:
config.num_hidden_layers = 4
for model_class in self.all_model_classes:
if model_class._no_split_modules is None:
@@ -2274,7 +2264,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 [0.5, 0.7, 0.9]]
max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir)
@@ -2292,12 +2282,7 @@ class ModelTesterMixin:
@require_accelerate
@require_torch_multi_gpu
def test_model_parallelism(self):
if all([model_class._no_split_modules is None for model_class in self.all_model_classes]):
return
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
if isinstance(getattr(config, "num_hidden_layers", None), int) and config.num_hidden_layers < 4:
config.num_hidden_layers = 4
for model_class in self.all_model_classes:
if model_class._no_split_modules is None:
@@ -2310,7 +2295,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 [0.5, 0.7, 0.9]]
max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir)