Size the elastic v0.2 micro batch from the data-parallel world

[vineethsaivs]

Elasticity v0.2 exists to support model parallelism, but its micro-batch pick divides the global batch by every GPU instead of by the number of data-parallel groups. With model_parallel_size > 1 it is off by exactly that factor, so it returns a needlessly small micro batch, or None when only the correct one is on offer.

Cause: get_microbatch uses final_batch_size // current_num_gpus. A DP rank's batch is final_batch_size // (current_num_gpus // model_parallel_size). The surrounding code already knows this: the branch above it compares current_num_gpus // model_parallel_size against valid_dp_world_size.

Fix: pass the DP world size in and divide by that.

Test: pytest tests/unit/elasticity/test_elastic.py -> 4 failed / 24 passed before, 28 passed / 3 skipped after. New test_model_parallel_v2_microbatch_is_per_dp_rank, with model_parallel_size=1 as the unchanged control.

16 GPUs, 8 per node, max_train_batch_size=64 (batch resolves to 64):

model_parallel_size	micro_batch_sizes	DP size	batch per DP rank	before	after
1	[4, 8]	16	4	4	4
2	[4, 8]	8	8	4	8
4	[4, 8]	4	16	4	8
2	[8]	8	8	None	8
4	[8]	4	16	None	8

[ebarkhordar]

I ran this at c5601dbe and at the merge base 56de5705 in a clean python:3.11-slim container, CPU torch, DS_ACCELERATOR=cpu. Your table reproduces. But the only caller in the tree does not reach the value you fixed.

deepspeed/runtime/config.py:734 calls compute_elastic_config(..., world_size=self.world_size) and never passes return_microbatch, so world_size > 0 holds, the function returns at elasticity.py:368, and the micro batch comes from final_batch_size // world_size % mbsz at line 362. That is the divisor you are correcting in get_microbatch, and candidate_microbatch_size is dropped on that path. Your test omits world_size and sets return_microbatch=True, so it measures the other branch.

16 GPUs, 8 per node, batch 64, micro [4, 8]:

                  world_size omitted    world_size=16
                  return_microbatch     (as config.py:734 calls it)
mp=1  56de5705    (64, [16], 4)         (64, [16], 4)
      c5601dbe    (64, [16], 4)         (64, [16], 4)
mp=2  56de5705    (64, [8], 4)          ElasticityIncompatibleWorldSize
      c5601dbe    (64, [8], 8)          ElasticityIncompatibleWorldSize

World size (16) is not valid with the current list of valid GPU counts: [8]. It raises at the merge base too, so this PR neither causes it nor depends on it. For v0.2 valid_gpus holds DP world sizes and line 355 compares the raw world_size against them. Passing 8 instead does not help: line 175 divides by model_parallel_size again, and the valid list becomes [4].

So world_size is total GPUs at line 324 and a DP world size at line 355. Is line 362 in scope here, or would you rather keep this to get_microbatch? As it stands the runtime path keeps the old divisor.

I only called compute_elastic_config directly, no multi-node run.

[vineethsaivs]

In scope, and folded in as f7e6c25. You are right that the fix did not reach the runtime, and I reproduced your table before changing anything.

Two things go wrong in that block, both the same mix-up:

• valid_gpus for version 0.2 is valid_dp_world_size from _get_compatible_gpus_v02, and the log line one above the check says so ("Valid World Size (GPUs / Model Parallel Size)"). So world_size not in valid_gpus compares total GPUs against DP world sizes and rejects valid configurations outright. That is the ElasticityIncompatibleWorldSize you hit: 16 // 2 == 8 is in [4, 8].
• final_batch_size // world_size % mbsz is the divisor from get_microbatch, unfixed.

Both now go through dp_world_size = world_size // model_parallel_size. Version 0.1 raises on model_parallel_size > 1 a few lines up, so nothing there moves.

Same setup as yours, 16 GPUs, 8 per node, batch 64, micro [4, 8], called the way config.py:734 calls it:

        c5601dbe                                    f7e6c25
mp=1    (64, [8, 16], 4)                            (64, [8, 16], 4)
mp=2    ElasticityIncompatibleWorldSize             (64, [4, 8], 8)
mp=4    ElasticityIncompatibleWorldSize             (64, [2, 4], 8)

New test_model_parallel_v2_runtime_path_uses_dp_world_size covers that branch with world_size=16 and no return_microbatch, keeping mp=1 as the control: pytest unit/elasticity/test_elastic.py is 2 failed / 29 passed before, 31 passed after. yapf --style .style.yapf clean.

[ebarkhordar]

That is the branch, thanks. Both the membership check and the divisor now use dp_world_size, which is what config.py:734 reaches, and the new test calls it the way the runtime does. I read f7e6c25e rather than rerunning it.

[vineethsaivs]

Thanks for reading it through. f7e6c25e is the head; nothing outstanding from my side.