Size the elastic v0.2 micro batch from the data-parallel world

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
 

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])])