Reuse existing hpZeRO group in nested zero.Init contexts - #8361
Reuse existing hpZeRO group in nested zero.Init contexts#8361zupengwang wants to merge 1 commit into
Conversation
Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
![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: a491d71408
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if self.zero_param_process_group is None: | ||
| groups._create_zero_param_parallel_group(hpz_partition_size) | ||
| self.zero_param_process_group = groups._get_zero_param_intra_parallel_group() | ||
| elif dist.get_world_size(group=self.zero_param_process_group) != hpz_partition_size: |
There was a problem hiding this comment.
When an outer Init has created world HPZ groups and an inner Init supplies a differently arranged data_parallel_group with the same cardinality—for example, outer HPZ groups [0,1]/[2,3] and inner DP groups [0,2]/[1,3]—this condition accepts reuse solely because the sizes match. The reused HPZ group then crosses the inner context's DP replicas, so secondary parameter gathers can combine values from different DP/model-parallel replicas and silently corrupt results. Verify that the existing group has the required membership within self.ds_process_group, or reject reuse, rather than comparing only its size.
Useful? React with 👍 / 👎.
| self.zero_param_process_group = zero_param_parallel_group | ||
| if _ds_config is not None and _ds_config.zero_config.zero_hpz_partition_size > 1 and self.zero_param_process_group is None: | ||
| groups._create_zero_param_parallel_group(_ds_config.zero_config.zero_hpz_partition_size) | ||
| hpz_partition_size = min(_ds_config.zero_config.zero_hpz_partition_size, dist.get_world_size()) |
There was a problem hiding this comment.
This is a non-merge commit, but its message has no Signed-off-by: trailer. Recreate the commit with --signoff so it satisfies the repository's mandatory commit requirement.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
Codex Review SummaryThis 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. |
1 participant
Reuse existing hpZeRO groups in nested
zero.InitcontextsProblem
Nested
deepspeed.zero.Initcontexts fail whenzero_hpz_partition_size > 1because each context unconditionally calls_create_zero_param_parallel_group(). The outer context has already initialized the process-wide HPZ group, so the inner context raises:This occurs in real model composition flows where an outer model is constructed under ZeRO-3 and a nested
from_pretrained()call enters anotherzero.Initcontext. It is reported in #4901 and reproduced again in #7066.Fix
RuntimeErrorwhen nested contexts request incompatible HPZ sizes instead of silently attaching the wrong group.The change is local to
zero.Init; it does not change public APIs or the non-HPZ path.Test
Added a two-rank regression that nests two
zero.Initcontexts with ZeRO-3 + HPZ, constructs a parameterized module, and checks that the parameter is partitioned with the expected HPZ group.Validation on a single-node RTX 3090 host, based on
87d9ecd8e0a4fd7778a58ac0f69cc85951f78ea0:1 passed.tests/unit/runtime/zero/test_zero_context.py:12 passed.pre-commit run --files deepspeed/runtime/zero/partition_parameters.py tests/unit/runtime/zero/test_zero_context.py: passed all hooks.pip wheel --no-deps --no-build-isolation .: passed withDS_BUILD_OPS=0.outer HPZ=2,inner HPZ=4): all ranks failed with the intended actionable error.Related: #4901, #7066