fix: avoid ZeroDivisionError in compute_elastic_config return_microbatch path - #8162
Closed
Solaris-star wants to merge 1 commit into
Closed
Solaris-star wants to merge 1 commit into
Solaris-star wants to merge 1 commit into
Conversation
…tch path For elasticity versions other than 0.2, return_microbatch=True with the default world_size=0 divided by zero. Require a positive world_size (or WORLD_SIZE) and raise ElasticityConfigError instead. Fixes #8156 Signed-off-by: Solaris-star <820622658@qq.com>
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 12, 2026
…8237) ## Problem `_get_compatible_gpus_v01` validates every micro batch against `max_acceptable_batch_size`: ```python if not all(mb <= max_acceptable_batch_size for mb in micro_batches): raise ValueError(...) ``` but the first heuristic scales the **LCM** of the micro batches, and the LCM is never checked. It goes into `base_list` and reaches `get_candidate_batch_sizes`, where the `base >= max_acceptable_batch_size` branch appends it unscaled. So a batch size the caller already said was too large becomes a candidate, and since the LCM divides every micro batch it tends to win the most-valid-GPU-counts vote in `get_best_candidates`. The docstring says the heuristic produces "the largest batch size less than the max_acceptable batch size", and `config-json.md` documents `max_train_batch_size` as "Max acceptable batch size can be used in training", so the returned value is not supposed to exceed it. ## Repro ```python import deepspeed from deepspeed.git_version_info import version as ds_version ds_config = {"elasticity": {"enabled": True, "max_train_batch_size": 100, "micro_batch_sizes": [8, 10, 12], "min_gpus": 1, "max_gpus": 1500, "min_time": 20, "version": 0.1}} print(deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)) # (120, [...]) <- 120 against a declared max of 100 ``` `DeepSpeedConfig.__init__` writes that return value straight into `self._param_dict[TRAIN_BATCH_SIZE]`, so the job runs 20 percent over the limit the user set, with the matching effect on the LR schedule and step count. It is not an exotic corner. `[8, 12]` with a cap of 16 gives 24, and a brute-force sweep over micro batch sets of size 2 and 3 drawn from 1 to 32, against every cap up to 400, finds it in 946105 configurations. The clearest evidence is in this repo: `tests/unit/elasticity/test_elastic.py::test_proper_mbsz` sets `max_train_batch_size` to 32 with micro batches `[1, 2, 3, 7]`, whose LCM is 42, and gets 42 back today. ## Fix Skip a base larger than the cap. Scaling one can only make it bigger, so it can never yield a legal candidate, and every micro batch is already validated against the cap, so the candidate list cannot end up empty. ## What this changes for the existing tests `test_basic_10k` is unaffected: still 9792, still 23 valid GPU counts. `test_proper_mbsz` needed one number changed, and I want to be upfront about it rather than bury it. Its `world_size=7` was only reachable because the batch size came back as 42, over its own cap of 32; at any legal batch size for that config, 7 is not a valid GPU count. I changed it to 4, where the batch per GPU is 6, so 7 is still correctly ruled out and the assertion that 3 is chosen is unchanged. That keeps the test doing what it was written to do, which is check the micro batch picked for a given world size. If you would rather keep `world_size=7` working, then the LCM overshoot is load-bearing rather than a bug, and this PR is the wrong change; I would want to hear that before going further. I could not find a config for those micro batches that makes 7 valid without exceeding the cap. `test_batch_size_within_max` is new and pins the actual contract. ## Test ``` before after test_basic_10k PASS PASS test_proper_mbsz (world_size=7, the old value) PASS ElasticityIncompatibleWorldSize test_proper_mbsz (world_size=4, the new value) ElasticityIncompatibleWorldSize PASS test_batch_size_within_max (new) FAIL: 120 exceeds 100 PASS ``` Run on CPU by driving the test bodies against the real `compute_elastic_config`, once against `master` and once against this branch; this path is pure Python and needs no GPU. `yapf --style .style.yapf` and `flake8 --config .flake8` are clean on both changed files, and clean on the unmodified tree as a control. There is one other open PR touching this file, deepspeedai#8162, in `compute_elastic_config`'s `return_microbatch` tail. It does not overlap these lines. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 23, 2026
…non-0.2 elasticity (deepspeedai#8286) ## 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. deepspeedai#8162 proposed the same resolution in July, and its author closed it unmerged on 2026-08-12 without a review. Fixes deepspeedai#8156 Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
compute_elastic_config(..., return_microbatch=True)on elasticity version0.1(and any non-0.2version) divided by the defaultworld_size=0and raised a bareZeroDivisionError. Version0.2already has a dedicated microbatch path that does not need this division.Changes
world_sizeargument or numericWORLD_SIZEenv before selecting a microbatch.ElasticityConfigErrorwith a clear message instead ofZeroDivisionError.world_sizeagainstvalid_gpusbefore dividing.world_size=64.Related issue
Fixes #8156
Tests
test_return_microbatch_requires_world_size_v01test_return_microbatch_with_world_size_v01