compute_elastic_config raises bare ZeroDivisionError for return_microbatch=True on non-0.2 elasticity versions without explicit world_size

[ErenAta16]

Description

compute_elastic_config (deepspeed/elasticity/elasticity.py) raises a bare ZeroDivisionError instead of a clear error when called with return_microbatch=True, no explicit world_size, and an elasticity config version other than 0.2 (e.g. the still-documented 0.1).

Root cause

def compute_elastic_config(ds_config: dict, target_deepspeed_version: str, world_size=0, return_microbatch=False):

world_size defaults to 0. Further down, two separate blocks divide final_batch_size by world_size:

if world_size > 0:                                  # line 347 - correctly guarded
    ...
    if final_batch_size // world_size % mbsz == 0:   # line 355
    ...

if return_microbatch:                                # line 363 - NOT guarded on world_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:   # line 370 - world_size can be 0 here
                micro_batch_size = mbsz
                break
        ...

The first block (line 347) is correctly guarded by if world_size > 0:. The second block (line 363-376), reached when return_microbatch=True and world_size was never set, has no equivalent guard before reusing world_size as a divisor at line 370.

Version 0.2 configs don't hit this because _get_compatible_gpus_v02 (called earlier for version == 0.2) requires a concrete GPU count and raises a clear ElasticityConfigError ("...needs WORLD_SIZE to compute valid batch size...") long before reaching line 363 if world_size wasn't resolved. Version 0.1 (and anything else routed through the else branch) has no such requirement earlier in the function, so it computes final_batch_size/valid_gpus successfully and only fails at the very last step.

Reproduction

from deepspeed.elasticity import compute_elastic_config

ds_config = {
    "elasticity": {
        "enabled": True,
        "max_train_batch_size": 2000,
        "micro_batch_sizes": [2, 4, 6],
        "min_gpus": 1,
        "max_gpus": 10000,
        "min_time": 20,
        "version": 0.1,
    }
}

compute_elastic_config(ds_config, target_deepspeed_version="0.19.3", return_microbatch=True)

Actual:

ZeroDivisionError: integer division or modulo by zero

Expected: either a clear, actionable error (matching the v0.2 path's ElasticityConfigError), or a successful result if world_size isn't actually required for this branch.

Controls, both confirmed correct:

• Same call with an explicit world_size=4 succeeds: (1680, [1, 2, ..., 840], 6).
• The same call shape against a version: 0.2 config fails cleanly with ElasticityConfigError: Elasticity V 0.2 needs WORLD_SIZE to compute valid batch size... instead of crashing.

Why this is reachable

compute_elastic_config is a public entry point (re-exported from deepspeed/elasticity/__init__.py), documented as "Intended to be called both by scheduling infrastructure and deepspeed runtime" - callers invoking it before a world size is known (e.g. scheduling-time capacity checks) are exactly the documented use case. version: 0.1 is still shown as a supported value in the function's own docstring example, so this isn't a deprecated-path-only issue.

Suggested fix

Guard line 370 the same way line 347 is guarded (e.g. raise ElasticityConfigError when world_size <= 0 and return_microbatch=True for non-0.2 versions), or route the non-0.2 branch through the same WORLD_SIZE environment-variable resolution used in the 0.2 path (lines 316-329) before reaching this block.

Environment

Verified against deepspeedai/DeepSpeed master branch (current as of today), pure Python - the affected logic has no CUDA/GPU dependency, reproduced by loading the real, unmodified elasticity.py/config.py/constants.py directly (only the unrelated deepspeed.utils/deepspeed.git_version_info/deepspeed.elasticity.utils import chain was stubbed to avoid pulling in torch, none of the logic under test was modified or mocked).

[ebarkhordar]

Reproduced on master at 4c275f9e, in a clean container with the real import chain (pip install -e . against torch 2.13.0+cpu, nothing stubbed):

File "/work/deepspeed/elasticity/elasticity.py", line 370, in compute_elastic_config
    if final_batch_size // world_size % mbsz == 0:
       ~~~~~~~~~~~~~~~~~^^~~~~~~~~~~~
ZeroDivisionError: integer division or modulo by zero

Your structural read of the fault is correct. Two corrections to the scope, both measured on that run.

Only version 0.1 is affected, not "anything else routed through the else branch". That else is unreachable for any other version, because two earlier guards fire first:

version	result
0.3	ElasticityConfigError: Attempting to run elasticity version 0.3 but runtime only supports up to 0.2 (line 297)
0.15	NotImplementedError: Unable to find elastic logic for version: 0.15 (line 343)

So the surface is exactly version: 0.1 with return_microbatch=True and world_size left at its default, and there it is deterministic rather than occasional: the if world_size > 0: block at line 347 returns, so line 370 is reached only when world_size is 0.

The documented WORLD_SIZE environment variable does not rescue 0.1. This is the part I would flag as the real user-facing trap:

config	WORLD_SIZE=4 in env	result
0.2	yes	OK, (1680, [...], 6)
0.2	no	ElasticityConfigError naming WORLD_SIZE
0.1	yes	ZeroDivisionError
0.1	no	ZeroDivisionError

Only the 0.2 branch resolves a GPU count from the environment (lines 316-329); the 0.1 path never looks at it. A caller who follows the 0.2 error message's own advice, "set it as an environment variable", still crashes on a 0.1 config. That makes your second suggested fix, routing the non-0.2 branch through the same WORLD_SIZE resolution, the one that closes the asymmetry rather than only improving the message text.

On reachability: the only non-test in-tree caller is deepspeed/runtime/config.py:717, which always passes world_size, so the runtime path cannot hit this. It needs an external caller, which matches the scheduling-time use case you describe.

[ErenAta16]

Both corrections are right and sharpen this a lot.

The version scope was imprecise on my end, I said "anything else routed through the else branch" without checking that 0.3 and other future/unsupported versions never actually get there, they're caught by the two earlier guards you listed. Good catch, the surface really is just 0.1.

The WORLD_SIZE table is the more important addition though. I only checked that 0.2 resolves it from the environment and assumed the fix should route 0.1 through the same mechanism, I didn't actually verify what happens if a 0.1 caller follows the 0.2 error message's own advice and sets WORLD_SIZE. That it still crashes either way is the sharper bug report: it's not just a bad error message, the documented recovery path for the sibling version doesn't work here at all. That's a stronger case for wiring the WORLD_SIZE resolution into the 0.1 path rather than just adding a guard that raises a clearer error.