fix: avoid ZeroDivisionError in compute_elastic_config return_microbatch path

deepspeed/elasticity/elasticity.py

@@ -364,15 +364,32 @@ def compute_elastic_config(ds_config: dict, target_deepspeed_version: str, world
         # Pick a valid micro batch size
         if float(elastic_config.version) == 0.2:
             return final_batch_size, valid_gpus, candidate_microbatch_size
-        else:
-            micro_batch_size = None
-            for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True):
-                if final_batch_size // 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" micro_batches={elastic_config.micro_batches}."
-            return final_batch_size, valid_gpus, micro_batch_size
+        # Non-0.2 versions need an explicit world_size (or WORLD_SIZE env) to
+        # choose a microbatch; without it the old path divided by world_size=0.
+        if world_size <= 0:
+            env_world_size = os.getenv("WORLD_SIZE")
+            if env_world_size is not None and env_world_size.isnumeric():
+                world_size = int(env_world_size)
+            else:
+                raise ElasticityConfigError(
+                    "return_microbatch=True requires a positive world_size "
+                    "(or WORLD_SIZE env) for elasticity versions other than 0.2. "
+                    f"Got world_size={world_size}, WORLD_SIZE={env_world_size!r}, "
+                    f"version={elastic_config.version}."
+                )
+        if world_size not in valid_gpus:
+            raise ElasticityIncompatibleWorldSize(
+                f"World size ({world_size}) is not valid with the current list "
+                f"of valid GPU counts: {valid_gpus}"
+            )
+        micro_batch_size = None
+        for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True):
+            if final_batch_size // 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" micro_batches={elastic_config.micro_batches}."
+        return final_batch_size, valid_gpus, micro_batch_size
 
     return final_batch_size, valid_gpus

tests/unit/elasticity/test_elastic.py

@@ -251,3 +251,26 @@ def test(self):
 
         with pytest.raises(deepspeed.elasticity.config.ElasticityError):
             model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+
+def test_return_microbatch_requires_world_size_v01(ds_config):
+    """return_microbatch on elasticity v0.1 needs world_size; do not ZeroDivisionError."""
+    # Ensure WORLD_SIZE is not set from other tests
+    os.environ.pop("WORLD_SIZE", None)
+    with pytest.raises(deepspeed.elasticity.config.ElasticityConfigError):
+        deepspeed.elasticity.compute_elastic_config(
+            ds_config=ds_config,
+            target_deepspeed_version=ds_version,
+            return_microbatch=True,
+        )
+
+
+def test_return_microbatch_with_world_size_v01(ds_config):
+    final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config(
+        ds_config=ds_config,
+        target_deepspeed_version=ds_version,
+        world_size=64,
+        return_microbatch=True,
+    )
+    assert mbsize == 17
+    assert final_batch_size % 64 == 0