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.