Warn when zero.Init silently falls back to a single rank (#8084)

[akshansh47]

Problem

Fixes #8084.

deepspeed.zero.Init resolves its partition group from dist.get_world_group(). If the distributed process group has not been initialized before zero.Init runs (the classic case: AutoModel.from_pretrained(...) under an active HfDeepSpeedConfig executes before deepspeed.init_distributed()), the call at the top of Init.__init__:

if not dist.is_initialized():
    init_distributed()

ends up with a process group that only sees the local rank, so self.dp_world_size == 1. zero.Init then materializes every parameter whole on every rank instead of partitioning it. Under deepspeed --num_gpus N every rank loads the full (unsharded) model and OOMs. The failure is silent and indistinguishable from an honest "model too big" OOM, which makes it very hard to diagnose.

Fix

Detect the case and emit a loud, actionable warning: the launcher reports WORLD_SIZE > 1 but the resolved group collapsed to a single rank. The warning names the likely cause and the fix (call deepspeed.init_distributed() before building the model under zero.Init).

This is intentionally a warning, not an error or an auto-fix:

• An explicitly supplied size-1 data_parallel_group is treated as intentional and never warns.
• A genuine single-process run (WORLD_SIZE unset or 1) never warns.
• It does not change any partitioning behavior, so it cannot break a working setup.

Detection is factored into a small pure helper _unsharded_single_rank_warning(...) so it is unit-testable without a GPU or a live process group.

Tests

tests/unit/runtime/zero/test_zero_init_unsharded_warning.py covers:

• warns when launcher WORLD_SIZE>1 but group is single-rank,
• no warning for a genuine single-process run (WORLD_SIZE unset / 1),
• no warning when the group actually shards (dp_world_size>1),
• no warning when an explicit data_parallel_group is supplied,
• malformed WORLD_SIZE values do not raise.

Made with Cursor

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[sfc-gh-truwase]

@akshansh47 thanks for the PR. Do you know why init_distributed() fails? I think the better solution is to make it work.
https://github.com/akshansh47/DeepSpeed/blob/642a9d2827218bd5465d553933f9d75afcc3ce31/deepspeed/runtime/zero/partition_parameters.py#L1046-L1048

[akshansh47]

Good question — I dug into this, and in this scenario init_distributed() doesn't actually fail. zero.Init already auto-initializes when no process group exists:

# partition_parameters.py L1017-1019
if not dist.is_initialized():
    init_distributed()
    assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm"

and when the launcher env (RANK/WORLD_SIZE/MASTER_ADDR/MASTER_PORT) is visible, that path initializes at the correct world size. So the silent world_size=1 outcome reported in #8084 arises when the gate is skipped: a process group already exists at zero.Init time but is single-rank — e.g. user code or another library (accelerate/transformers init paths, per-GPU Ray actors, notebook helpers) called init_process_group earlier with a size-1 world, before from_pretrained ran under the active HfDeepSpeedConfig.

In that case I don't think DeepSpeed can safely "make it work": the existing PG is user/library state, and respecting it is correct — tearing it down and re-initializing from env inside zero.Init would be surprising and unsafe (in-flight collectives, custom groups). What DeepSpeed can own is detecting the contradiction (launcher says WORLD_SIZE>1, resolved group is size 1, no explicit data_parallel_group passed) and saying loudly that nothing will be partitioned — which is what this PR does.

Happy to adjust if you'd prefer a different behavior:

1. keep it a warning (current PR),
2. escalate to a hard error (strictest, but could break exotic-but-intentional setups), or
3. attempt re-init from env behind an opt-in flag (I'd advise against default-on for the reasons above).

I've also asked the reporter on #8084 to share their exact init order to confirm which actor created the single-rank group in their case.

[sfc-gh-truwase]

@akshansh47 thanks for providing useful clarity on this issue. I think option 2 (fail loudly) is the right direction here. It is unclear how zero3 (or deepspeed) can work correctly with a pre-existing and conflicting PG. I would prefer to engage with the exotic-but-intentional setups as needed.

[sfc-gh-truwase]

@akshansh47 to be clear, my proposal is that zero.Init() should fail in the face of contradicting PGs.

[akshansh47]

Done — escalated to a hard failure per your proposal. zero.Init now raises RuntimeError when the resolved default process group is single-rank while the launcher environment reports WORLD_SIZE > 1 (a contradicting PG that ZeRO-3 cannot partition against). The error message names the likely cause (from_pretrained before deepspeed.init_distributed()) and the two escape hatches: initialize the PG first, or pass an explicit data_parallel_group if a single-rank group is intentional (that path never raises, so the exotic-but-intentional setups keep a supported route).

Helper renamed to _contradicting_single_rank_pg_error and the unit tests updated accordingly (tests/unit/runtime/zero/test_zero_init_pg_contradiction.py).

[sfc-gh-truwase]

Move this UT to tests/unit/v1/zero/

[tohtana]

Hi @akshansh47, @sfc-gh-truwase,

While I'm reviewing #8117, I noticed one edge case in this PR.
sequence_data_parallel_group is deprecated, but zero.Init still accepts it and assigns it to self.ds_process_group.
This will lead to passing unrelated (inconsistent) arguments to _contradicting_single_rank_pg_error.