Fix ZeroDivisionError in compute_elastic_config return_microbatch on non-0.2 elasticity

[ebarkhordar]

Root cause

With return_microbatch=True, compute_elastic_config takes the else at elasticity.py:372 for any elasticity version other than 0.2, and that branch divides final_batch_size by world_size. In practice that means version 0.1: 0.3 is rejected at line 301 and any other value raises NotImplementedError at line 348, both before this point. The branch is only reachable when world_size is unset, because the if world_size > 0: block above it returns for every positive value. So on master the division is always by zero and the caller gets a bare ZeroDivisionError at line 375 rather than a configuration error.

Version 0.2 avoids this by resolving world_size from the WORLD_SIZE environment variable at lines 321-334, and raising ElasticityConfigError naming that variable when it cannot. Version 0.1 never reads the environment, so a caller who follows the 0.2 message's own advice, "set it as an environment variable", still crashes:

config	WORLD_SIZE in env	master
0.2	yes	returns (9792, [...], 17)
0.2	no	ElasticityConfigError naming WORLD_SIZE
0.1	yes	ZeroDivisionError
0.1	no	ZeroDivisionError

Fix

Resolve world_size from WORLD_SIZE in the non-0.2 branch the way 0.2 already does, and raise ElasticityConfigError with the same guidance when it cannot be resolved. Then check the resolved value against valid_gpus before dividing, matching the sibling block at lines 352-355; without that check an out-of-range WORLD_SIZE would reach the loop and fail on the micro_batch_size is not None assertion instead of ElasticityIncompatibleWorldSize.

Both divisions by world_size in this function now run only on a value that is positive and a member of valid_gpus. Nothing that works today changes: on master this branch raised ZeroDivisionError for every input, and the 0.2 path is untouched.

Verification

• Three new tests in tests/unit/elasticity/test_elastic.py cover the unset case, resolution from WORLD_SIZE, and an out-of-range WORLD_SIZE. All three fail on master with ZeroDivisionError at elasticity.py:375 and pass here.
• pytest unit/elasticity/ gives 23 passed, 3 skipped, on Python 3.12 with torch==2.10.0+cpu to match the cpu-torch-latest leg. The 3 skips are the DistributedTest classes that need FusedLambBuilder, and they skip on master too.
• pre-commit run --files passes on both changed files, yapf, flake8, codespell and check-torchdist included.
• Not checked: the GPU legs, and the 0.2 return_microbatch return at line 371, which no test in the repo reaches either before or after this change.

#8162 proposed the same resolution in July, and its author closed it unmerged on 2026-08-12 without a review.

Fixes #8156

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e37bb0fe06

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required sign-off trailer

This is a non-merge commit, but its message has no Signed-off-by trailer, so it does not satisfy the repository's commit requirements and may fail the DCO/CI gate. Recreate the commit with --signoff using the configured author name and email.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[tohtana]

Hi @ebarkhordar,
Thank you for your contribution! Looks good to me.

[ebarkhordar]

Thanks for the review and the merge, Masahiro. Glad this one was useful. I will keep an eye on the elasticity config paths and send anything else I hit there your way.