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

Commit
53 days ago
Warn when zero.Init silently falls back to a single rank (#8084) (#8089) ## 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__`: ```python 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](https://cursor.com) --------- Signed-off-by: FNU AKSHANSH <105249360+akshansh47@users.noreply.github.com> Signed-off-by: Akshansh <105249360+akshansh47@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Author
Parents
Loading