Keep the elastic batch size within max_train_batch_size

[vineethsaivs]

Problem

_get_compatible_gpus_v01 validates every micro batch against max_acceptable_batch_size:

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

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, #8162, in compute_elastic_config's return_microbatch tail. It does not overlap these lines.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a66f83037c

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[tohtana]

Thank you @vineethsaivs for your contribution! Looks good to me.