DeepSpeed

2cfebbda - Keep the elastic batch size within max_train_batch_size (#8237)

Keep the elastic batch size within max_train_batch_size (#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, #8162, in `compute_elastic_config`'s `return_microbatch` tail. It does not overlap these lines. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>