diff --git a/deepspeed/elasticity/elasticity.py b/deepspeed/elasticity/elasticity.py index 17f5aaca6edf..ee1b71bef2bd 100644 --- a/deepspeed/elasticity/elasticity.py +++ b/deepspeed/elasticity/elasticity.py @@ -148,11 +148,13 @@ def _get_compatible_gpus_v02(micro_batches, f"{num_gpus_per_node} should be divisible by " \ f"model parallel size {model_parallel_size}") - def get_microbatch(final_batch_size): + def get_microbatch(final_batch_size, dp_world_size): candidate_microbatch = None for micro_batch in micro_batches: - if final_batch_size // current_num_gpus % micro_batch == 0: + # A data-parallel rank sees the global batch split across the DP groups, not + # across every GPU: dividing by current_num_gpus is off by model_parallel_size. + if final_batch_size // dp_world_size % micro_batch == 0: if candidate_microbatch is None: candidate_microbatch = micro_batch if prefer_larger and candidate_microbatch < micro_batch: @@ -171,7 +173,7 @@ def get_microbatch(final_batch_size): final_batch_size = int(final_batch_size) * dp_size_per_node valid_dp_world_size = [i * dp_size_per_node for i in valid_world_size] if current_num_gpus // model_parallel_size in valid_dp_world_size: - candidate_microbatch = get_microbatch(final_batch_size) + candidate_microbatch = get_microbatch(final_batch_size, current_num_gpus // model_parallel_size) return final_batch_size, valid_dp_world_size, candidate_microbatch current_dp_size = (current_num_gpus / num_gpus_per_node) * dp_size_per_node @@ -188,7 +190,7 @@ def get_microbatch(final_batch_size): else: candidate_batch_size = min(candidate_batch_sizes) - candidate_microbatch = get_microbatch(candidate_batch_size) + candidate_microbatch = get_microbatch(candidate_batch_size, int(current_dp_size)) return candidate_batch_size, [int(current_dp_size)], candidate_microbatch @@ -350,18 +352,24 @@ def compute_elastic_config(ds_config: dict, target_deepspeed_version: str, world logger.info(f"Valid World Size (GPUs / Model Parallel Size): {valid_gpus}") if world_size > 0: - if world_size not in valid_gpus: + # `valid_gpus` holds DP world sizes, which is what the log line above says, so a world + # size that spans model-parallel groups has to be reduced the same way before it is + # compared against the list or used as the divisor. Version 0.1 rejects + # model_parallel_size > 1 above, so there this is the arithmetic it already does. + dp_world_size = world_size // model_parallel_size + if dp_world_size not in valid_gpus: raise ElasticityIncompatibleWorldSize(f"World size ({world_size}) is not valid " \ f"with the current list of valid GPU counts: {valid_gpus}") # Pick largest valid micro batch size micro_batch_size = None for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True): - if final_batch_size // world_size % mbsz == 0: + if final_batch_size // dp_world_size % mbsz == 0: micro_batch_size = mbsz break assert micro_batch_size is not None, "Unable to find divisible micro batch size" \ - f" world_size={world_size}, final_batch_size={final_batch_size}, and " \ + f" world_size={world_size}, dp_world_size={dp_world_size}," \ + f" final_batch_size={final_batch_size}, and " \ f" micro_batches={elastic_config.micro_batches}." return final_batch_size, valid_gpus, micro_batch_size diff --git a/tests/unit/elasticity/test_elastic.py b/tests/unit/elasticity/test_elastic.py index 79f3ee206d18..e984935f1837 100644 --- a/tests/unit/elasticity/test_elastic.py +++ b/tests/unit/elasticity/test_elastic.py @@ -153,6 +153,65 @@ def test_model_parallel_v2_valid(ds_config): os.environ.pop("WORLD_SIZE") +@pytest.mark.parametrize("model_parallel_size, micro_batch_sizes, expected_mbsize", [ + (1, [4, 8], 4), + (2, [4, 8], 8), + (4, [4, 8], 8), + (2, [8], 8), + (4, [8], 8), +]) +def test_model_parallel_v2_microbatch_is_per_dp_rank(ds_config, model_parallel_size, micro_batch_sizes, + expected_mbsize): + # A data-parallel rank gets the global batch split across the DP groups, so the micro batch + # has to divide batch / (gpus // model_parallel_size). Dividing by every GPU instead is off + # by model_parallel_size: it picked a needlessly small micro batch, and returned None when + # only the correct one was on offer. model_parallel_size 1 is the control, where the two + # divisors are equal. + ds_config["elasticity"]["version"] = 0.2 + ds_config["elasticity"]["model_parallel_size"] = model_parallel_size + ds_config["elasticity"]["num_gpus_per_node"] = 8 + ds_config["elasticity"]["max_train_batch_size"] = 64 + ds_config["elasticity"]["micro_batch_sizes"] = micro_batch_sizes + ds_config["elasticity"]["min_gpus"] = 16 + ds_config["elasticity"]["max_gpus"] = 16 + + os.environ["WORLD_SIZE"] = str(16) + try: + final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config( + ds_config=ds_config, target_deepspeed_version=ds_version, return_microbatch=True) + finally: + os.environ.pop("WORLD_SIZE") + + dp_world_size = 16 // model_parallel_size + assert mbsize == expected_mbsize + assert final_batch_size // dp_world_size % mbsize == 0 + + +@pytest.mark.parametrize('model_parallel_size, expected_mbsize', [(1, 4), (2, 8), (4, 8)]) +def test_model_parallel_v2_runtime_path_uses_dp_world_size(ds_config, model_parallel_size, expected_mbsize): + # deepspeed/runtime/config.py calls compute_elastic_config with world_size set and without + # return_microbatch, so the world_size > 0 block is the branch every real run takes. For + # version 0.2 valid_gpus holds DP world sizes, which is what the log line above it says, so + # both the membership check and the micro-batch divisor need world_size reduced the same way. + # Comparing the raw world size rejected valid configurations outright. model_parallel_size 1 + # is the control, where the two are equal. + ds_config["elasticity"]["version"] = 0.2 + ds_config["elasticity"]["model_parallel_size"] = model_parallel_size + ds_config["elasticity"]["num_gpus_per_node"] = 8 + ds_config["elasticity"]["max_train_batch_size"] = 64 + ds_config["elasticity"]["micro_batch_sizes"] = [4, 8] + ds_config["elasticity"]["min_gpus"] = 8 + ds_config["elasticity"]["max_gpus"] = 16 + + final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config( + ds_config=ds_config, target_deepspeed_version=ds_version, world_size=16) + + dp_world_size = 16 // model_parallel_size + assert dp_world_size in valid_gpus + assert mbsize == expected_mbsize + assert final_batch_size // dp_world_size % mbsize == 0 + + @pytest.mark.parametrize('key, value', [('micro_batch_sizes', [1, 4, -1, 2, -10]), ('min_gpus', -1), ('max_gpus', -1), ('micro_batch_sizes', 5), ('micro_batch_sizes', ['a', None, 0.5]), ('micro_batch_sizes', [2, 0.5, 4])])