Repository navigation
Conversation
Many modules guard all_reduce with 'group is not None and world_size > 1' so that code written for an optional TP/SP group runs without it. Provide this as a comm helper where group=None explicitly means an absent parallel dimension (not the world group), and use it in sequence/cross_entropy.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
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. |
![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: 2d64bd8725
ℹ️ 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".
| ``tensor`` unchanged. Only synchronous reductions are supported, since a skipped call has no work | ||
| handle to return. | ||
| """ | ||
| if group is None or get_world_size(group) == 1: |
There was a problem hiding this comment.
When a caller passes a process group that excludes the current rank, get_world_size(group) returns -1 per this module's contract, so this equality check falls through even though there is no local work. This regresses the replaced > 1 guards and is not harmless on the custom CCL backend: CCLBackend.get_all_ranks_from_group derives an empty rank list and calls _new_group, which indexes ranks[0]; reduce only when the reported size is greater than one so non-members skip the collective.
AGENTS.md reference: AGENTS.md:L40-L40
Useful? React with 👍 / 👎.
| return cdb.all_reduce(tensor, op, group, async_op) | ||
|
|
||
|
|
||
| def all_reduce_if_parallel(tensor, op=ReduceOp.SUM, group=None): |
There was a problem hiding this comment.
all_reduce_if_parallel is exported through deepspeed.comm and gives group=None semantics that intentionally differ from all_reduce, but this commit adds no entry under docs/code-docs/source, and a repository-wide rg finds no documentation of the name outside source and tests. Without a public reference, downstream callers can easily apply torch.distributed's world-group convention to it; add the documentation required for this new public feature.
AGENTS.md reference: AGENTS.md:L71-L71
Useful? React with 👍 / 👎.
…ll sites Keep deepspeed/comm/comm.py for plain collectives, where group=None means the world group, and put helpers for optional parallel groups, where group=None means the dimension is absent, in deepspeed/comm/pg_utils.py. pg_utils is not star-exported from deepspeed.comm, so the two meanings of None stay apart. - pg_utils.is_parallel(group): group is not None and has more than one rank. - pg_utils.all_reduce_if_parallel(tensor, op, group): synchronous all-reduce guarded by is_parallel. Replace the hand-written "group is not None and world_size > 1" guards in sequence (vocab_parallel_ops, cross_entropy), MoE (sharded_moe), inference v1 (ds_mlp, ds_attention, triton mlp/attention, moe_inference), ZeRO (stage_1_and_2, stage3), module_inject/layers.py and runtime/utils.py. Guards that compare a separately supplied world size are left unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
changed the title
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
changed the title
Drop all_reduce_if_parallel: a single predicate works for every collective (sync, async, all_gather, ...) without one wrapper per primitive. Also guard the lm_head inference all-reduce and the MoE inference expert-MP gather. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Replace the per-call 'mp_group is None or tp_world_size == 1' guards in the AutoTP layers with a value computed once from pg_utils.is_parallel at construction, so forward/backward pay no extra cost. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…parallel Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
The AutoEP expert group is now reduced only when it spans more than one rank, so the test's stand-in group needs a world size like the neighbouring tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
|
@delock I reproed the CI failure cases, I fixed |
Head branch was pushed to by a user without write access
This branch has not been deployed
2 participants
Part of #8814, following @delock's suggestion in #8791 and the review discussion on where the helper should live.
Change
deepspeed/comm/pg_utils.pywithis_parallel(group):groupis notNoneand spans more than one rank.comm.pykeeps the plain collectives, wheregroup=Nonemeans the world group. Inpg_utils,group=Nonemeans the parallel dimension is absent.pg_utilsis not star-exported fromdeepspeed.comm; callers writefrom deepspeed.comm import pg_utils, so the two meanings ofNonestay apart.all_reduce_if_parallel), so one helper guards every primitive: sync or async all-reduce, all-gather, inference all-reduce, ...:group is not None and dist.get_world_size(group) > 1guards, and variants that only checkNone:vocab_parallel_ops.py(5; local_tp_active/_all_reduceremoved),cross_entropy.py(4,_global_sp_sum, empty-shard check)sharded_moe.pycapacity MAX reductions (3)ds_mlp.py,triton/mlp.py,ds_attention.py,triton/attention.py,moe_inference.py(async all-reduce, expert-MP all-gather and split)stage_1_and_2.py/stage3.py_model_parallel_all_reduce,stage3.pyAutoEP norm reductions (3)module_inject/layers.py:RowParallelforward,ScatterToTensorParallelRegion/ColumnParallelbackward,LmHeadLinearAllreduceforward;TensorParallel_Layercachesis_tp_parallel = is_parallel(mp_group)once at construction, used by the 6mp_group is None or tp_world_size == 1guards inLinearAllreduce/LinearLayer/SubParam*/ fused-QKV layers (no per-call cost)runtime/utils.pyMoE EP norm reductions (2)Behavior is unchanged. Sites that only checked
Nonenow also skip single-rank groups, where every guarded collective is a no-op.Not migrated:
len(partition_sizes), atp_world_sizeargument) in gather helpers: the size is not derived from a group.Nonemeans the world group:kv_cachewarm-up,moe_inference._alltoall.Tests
tests/unit/comm/test_dist.py:pg_utilscontract on 4 ranks.Noneand single-rank groups are not.group=Noneworks without a process group.is_parallelis not exported asdeepspeed.comm.is_parallel.tests/unit/module_inject/test_tp_autograd_reduce.py, with 2-rank, single-rank andNonegroups:RowParallelforward.ColumnParallel/ScatterToTensorParallelRegionbackward.LmHeadLinearAllreduceoutput matches the unsharded linear.tests/unit/runtime/test_runtime_utils.py::TestGlobalNormMoEGroup: expert-group 2-norm and inf-norm with the same three group kinds.is_parallel(None)return True (a world reduction) fails theNonecases of the new tests.test_runtime_utils.pytest_moe.py::TestMoESingletontest_vocab_parallel_ops.pytest_vocab_parallel_cross_entropy.pytest_liger_vocab_parallel_cross_entropy.pytest_invalid_grad_norm.pyRun on 8× Intel Arc Pro B70 (XPU, XCCL), torch 2.14.0.dev20260804+xpu.
Contribution checklist
Related work. This is part of RFC #8814. A search for
is_parallel,all_reduce_if_paralleland "optional parallel group" found no other PR that adds an optional-group predicate.Commands and results. Hardware and software: 8× Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu.
tests/unit/v1/moe/test_autoep_unit.py: 192 passed. This covers the ZeRO-3 AutoEP norm paths. Its fake EP group now mocksdist.get_world_size, becauseis_parallelqueries the group size.Numerics, performance, memory. Nothing changes. Every site keeps the same collective; sites that only checked for
Nonenow also skip single-rank groups, where the collective is a no-op.Compatibility matrix.
RowParallel/ColumnParallel/ScatterToTensorParallelRegion/LmHeadLinearAllreducewith 2-rank, single-rank andNonegroups_model_parallel_all_reduce; stage-3 AutoEP norm paths in unit testsTestMoESingletonds_mlp,ds_attention,triton/*,moe_inference: mechanical guard replacement onlyAI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.