Repository navigation
Conversation
…tion group A ZeRO-3 engine whose data-parallel group is a subset of the world (two trainer layouts living in the same processes, or a trainer next to inference-only ranks) hung or failed its trace checks: the checks broadcast from global rank 0 over the world and the barriers waited for every rank. They now run over the parameters' partition group; with the default group nothing changes. Signed-off-by: RunguoLi <li19107254665@gmail.com>
|
Thanks @LiRunGuo. I read the diff against master: it covers every caller of the three check helpers (the two in Init._all_gather, the three trace checks in PartitionedParameterCoordinator, and the sanity check in stage 3) and all three world barriers in stage3.py, and group=None keeps today's behaviour. dist.get_global_rank(group, 0) exists in deepspeed.comm, so the first-rank lookup works for any group. Looks right to me. One related case this PR leaves out, which I hit on master while rebuilding ZeRO-3 engines in the same processes for the layout-change runs: ZeRO-3 numbers modules and parameters with process-wide counters, DeepSpeedZeRoOffload._register_deepspeed_module(self, module, count=[0]) (a mutable default) and the class attribute zero.Init.param_id, and neither is reset between engines. So after a phase in which only some ranks built an engine (DP4 → DP2 on two of four ranks), a new engine over all four ranks numbers its modules differently on the ranks that sat out, and the first trace check fails with "Detected a disagreement on list contents between rank0 and rank2". This PR doesn't touch those counters, so a subgroup engine followed by a new full-group engine should hit the same check; #9074 avoids it by keeping both engines alive. Resetting the counters per engine isn't safe in general, because parameter ids must stay unique across several zero.Init contexts for one model, so I'd treat it as a separate issue. I can open one with the reproduction if that's useful. |
TestZero3OverSubgroup trains ZeRO-3 on ranks 1-2 of a world of three through an mpu, with no engine on rank 0, and compares three SGD steps with plain torch SGD on the same model and batch. Without the previous commit the test does not finish: the world barriers and the broadcasts from global rank 0 wait for rank 0. Signed-off-by: RunguoLi <li19107254665@gmail.com>
|
Thanks @Achyuthan-S for reading the diff. The test is in. The second commit (54680fd) adds The counters. Both are on master as you describe (
Yes please open the issue with your reproduction; I agree it is separate from this PR. One direction, untested: the checks could compare each id's position among the engine's own ids instead of the process-wide number. Ids stay unique across I've marked the PR ready for review. |
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 54680fd9a0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights, | ||
| fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module, | ||
| log_trace_cache_warnings=self.log_trace_cache_warnings, | ||
| process_group=self.dp_process_group, |
There was a problem hiding this comment.
In supported zero.Init-before-AutoTP setups, the optimizer DP group is not necessarily the parameters' partition group: tests/unit/checkpoint/test_zero_affine.py:430-443 explicitly builds world-partitioned parameters before creating TP2/DP2. Wiring the coordinator to self.dp_process_group lets its submodule and parameter-order checks pass independently in each DP lane, while its later all_gather_coalesced calls still communicate those parameters over param.ds_process_group; if TP lanes record different orders, equal-sized tensors can be mixed or the all-gather can hang instead of the disagreement being reported. Validate the trace for every represented parameter partition group rather than only the optimizer DP group.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
| with get_accelerator().stream(self.reduce_and_partition_stream): | ||
| if self.enable_sanity_checks: | ||
| assert_ints_same_as_other_ranks([p.ds_id for p in params_in_bucket]) | ||
| assert_ints_same_as_other_ranks([p.ds_id for p in params_in_bucket], self.dp_process_group) |
There was a problem hiding this comment.
When a parameter's partition group differs from the optimizer DP group, as in the world-partitioned zero.Init-before-AutoTP path, this check compares only DP-lane peers even though the reduction immediately below uses bucket.process_group or _get_param_partition_group(params_in_bucket[0]) (stage3.py:1945-1960 and 1994-2013). With enable_sanity_checks, a different bucket order among the actual reduction peers can therefore go undetected and equal-shaped gradients can be reduced into the wrong parameters; pass the bucket's partition group here.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
gpu-test-labeler
Bot
added
the
skip-gpu-test
|
A correction to my comment above. I wrote that I had not reproduced the counter failure; I had met it before and forgot. In an early run of the DP4 ↔ DP2 switch, an engine built on a subset of the ranks first, and one over all ranks after it, failed the trace check with "disagreement on list contents". #9074 avoids it by building the engine over every rank first. That matches what the code suggests. |
ZeRO-3's rank-agreement checks and init barriers ran over the world. They now run over the parameters' partition group, so a ZeRO-3 engine can train on a subset of the ranks while the other ranks hold no engine. With the default group the same ranks take part as before.
This is the first commit of #9074, sent on its own because it does not depend on RFC #9026 and is small (31 lines added and 19 removed in
deepspeed/), plus a regression test as a second commit (85 lines intests/unit/v1/zero/test_zero.py).What breaks on master. Give
deepspeed.initializeanmpuwhose data-parallel group is a subset of the world, for example ranks 1-2 of three, and leave rank 0 without an engine (a trainer next to inference-only ranks, or two trainer layouts in the same processes):DeepSpeedZeroOptimizer_Stage3._setup_for_real_optimizercallsdist.barrier()over the world, twice, and_create_fp16_partitions_with_defragmentationonce more. The trainer ranks wait for rank 0, which never calls it.assert_ints_same_as_other_ranksbroadcasts from global rank 0 over the world (assert_lst_len_same_as_other_ranks,get_lst_from_rank0). It is called when a trace is recorded (PartitionedParameterCoordinator), in the coalesced all-gather (Init._all_gather) and, withenable_sanity_checks, in the gradient reduction. Rank 0 is outside the group and never sends.Root cause. These calls predate engines over a subgroup: the data path of ZeRO-3 already uses the partition group (
all_gather,reduce_scatterand the parameter broadcasts all pass it), and only the checks and barriers around it still assumed the world.The change
zero/utils.py:assert_ints_same_as_other_ranks,assert_lst_len_same_as_other_ranksandget_lst_from_rank0take an optionalgroupand broadcast from the group's first rank.group=Noneis the old behaviour, so other callers are unaffected.stage3.py: the three barriers and the sanity check passself.dp_process_group.partition_parameters.py: the two checks in_all_gatherpass the parameters'ds_process_group.partitioned_param_coordinator.py: takesprocess_group(defaultNone) and passes it to its three trace checks;parameter_offload.pyhands it the optimizer'sdp_process_group.Validation
One node with 4× RTX 3090 (24 GB), driver 595.58.03, torch 2.13.0+cu130, NCCL 2.29.7.
Regression test
TestZero3OverSubgroup(intests/unit/v1/zero/test_zero.py): a world of three, ZeRO-3 over ranks 1-2 through anmpu, rank 0 without an engine. Three SGD steps are compared with plain torch SGD on the same model and batch. SGD withgradient_clippingset to 0, so that the comparison sees the scale of the averaged gradient (an engine clips to norm 1.0 by default, and Adam hides a constant scale).python -m pytest -q -rfEs, fromtests/)unit/v1/zero/test_zero.py::TestZero3OverSubgroupDS_ACCELERATOR=cpu)unit/v1/zero/test_zero.py,unit/v1/zero/test_zero_init_pg_contradiction.py,unit/runtime/zero/test_zero_context.pyI did not look at where the run on master blocks; the first world barrier in
_setup_for_real_optimizeris the first call rank 0 never answers.Compatibility
dp_process_groupholds every rank, so the same ranks take part in every barrier and check as before. The ZeRO test files above run at the default group.enable_sanity_checks, several nodes.save_checkpointand_commit_decoupled_checkpointinengine.pystill calldist.barrier()over the world, so a subgroup engine cannot save a checkpoint while the other ranks stay out. That is a separate change.AI assistance: AI assistance was used. The code, the test and this description were drafted with an AI coding agent (Claude Code). The submitter reviewed every changed line and ran the tests above.
cc @Achyuthan-S @jinyouzhi