Fix all offload and MP tests (#17533)

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

View File

@@ -63,7 +63,7 @@ class OPTModelTester:
use_labels=False,
vocab_size=99,
hidden_size=16,
num_hidden_layers=2,
num_hidden_layers=5,
num_attention_heads=4,
intermediate_size=4,
hidden_act="gelu",

View File

@@ -515,6 +515,8 @@ class T5ModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
test_resize_embeddings = True
test_model_parallel = True
is_encoder_decoder = True
# The small T5 model needs higher percentages for CPU/MP tests
model_split_percents = [0.8, 0.9]
def setUp(self):
self.model_tester = T5ModelTester(self)

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)