Fix all offload and MP tests (#17533)
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user