Repository navigation
Count each expert once in the fp32 gradient-clipping norm - #8478
alanhuangyoo wants to merge 8 commits into
Conversation
`clip_grad_norm_` finishes by averaging the per-rank norms over the data
parallel group:
pg = groups._get_data_parallel_group()
scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg))
dist.all_reduce(scaled_norm_tensor, group=pg)
That reconstructs a global norm only while every rank holds the same
parameters. Expert parameters are not replicated, so each rank's norm covers a
different set of experts and the mean of those norms is not the norm of the
whole. Four expert gradients over two ranks, measured on 2xH20:
reference, each expert counted once 10.954452 sqrt(120)
replicated over both ranks 10.954452
one half per rank 7.236068 (sqrt(20) + sqrt(100)) / 2
The same gradients, 34% apart, decided by where the experts happen to live.
`clip_coef` is `max_norm / total_norm`, so whenever clipping is active this
changes the update.
The repository already has the correct combiner: `get_norm_with_moe_layers`
takes the non-expert norm and the expert gradients keyed by group, reduces each
expert group over its own `moe_ep_group`, and combines as
`(sum_i norm_i ** p) ** (1/p)`. Both the bf16 and the fp16 optimizer use it.
The fp32 path instead inlines the approximation that `get_norm_with_moe_layers_fast`
carries under its own name -- and that named copy has no callers at all, so the
comment on it explaining that a more precise implementation exists never reaches
the one place still doing it.
So: keep expert gradients out of the replicated norm, then fold them in through
that helper. Behaviour is unchanged when no expert parameter is present, which
is every non-MoE fp32 run -- that branch is the original code, untouched.
Two guards worth naming:
- An expert parameter carrying no `group_name` has no expert parallel group to
reduce over. `TestClipGradNorm::test_gather` builds one by hand exactly like
that, so such a call stays on the original path rather than guessing a group.
- `get_norm_with_moe_layers` reports a non-finite norm as `-1`. Used as-is that
would make `clip_coef` negative and flip the sign of every gradient, so it is
mapped to `inf`, which keeps what the averaging path already did: scale to zero.
Fixes deepspeedai#8469
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
|
Do not review this as it stands — two defects found by @gss10282023 on #8469, both of which I reproduced.
Marking this as a draft while I fix both. The branch decision has to be rank-independent — There is also #8476 from @gss10282023 covering the AutoEP FP32 L2 path. We are discussing on #8469 whether these should be one change; if a maintainer prefers a single PR, theirs came first. |
Two defects @gss10282023 found on the first version of this branch, both reproduced on 2xH20 before fixing. **Deadlock.** The branch read expert ownership after `p.grad is not None` filtering, so a rank whose experts received no tokens this step took the averaging path while its peer took the expert path. The two issue different collectives and they waited on each other: [rank0] entering clip_grad_norm_ (has expert grads) [rank1] entering clip_grad_norm_ (no expert grads) <neither returns; killed at 90s> Ownership now comes from the parameters, before that filter. A rank can also legitimately own no expert at all -- under pipeline parallelism -- so the choice is settled by one all-reduced flag rather than by what this rank happens to hold, and the group names come from `_EXPERT_PARALLEL_GROUP`, which is built identically everywhere. A rank with nothing for a group contributes a zero, which adds nothing to a sum of p-th powers and nothing to a max of absolute values. **Expert-only inf norm.** With every parameter an expert there is nothing replicated left to take a max over, and `torch.stack([])` raised. Zero is the identity for a max over absolute values. Three tests cover the cases: experts idle on one rank, a rank owning no expert, and an expert-only infinity norm. All three fail on master, as does the original sharded-expert case. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
|
@gss10282023 Both fixed in Deadlock. Your diagnosis was exactly right, and your own patch already had the answer in a comment I should have read more carefully:
I was reading ownership after That alone was not enough, though — your case was "no expert gradients", and there is a harder one: a rank owning no expert parameter at all, which happens under pipeline parallelism. Local inspection cannot settle that in kind, so the choice is now made by one all-reduced flag, with group names taken from Measured, same harness that hung before:
Expert-only inf norm was Three tests added for your cases. On On the split: having read #8476, I do not think either subsumes the other. Yours keys on What I would rather not do is leave Un-drafting since the two defects are closed, but I am in no hurry for this to land ahead of that conversation. |
| ], | ||
| device=get_accelerator().current_device_name(), | ||
| dtype=torch.float) | ||
| dist.all_reduce(vote, group=groups._get_data_parallel_group()) |
There was a problem hiding this comment.
I read this against 338ed51 rather than running it, so this is a question about the branch selector, not a repro.
Your comment above says the all-reduced flag is what settles the pipeline-parallel case, where a rank owns no expert parameter at all. The flag is reduced over groups._get_data_parallel_group(), and under pipeline parallelism that group does not span stages:
deepspeed/__init__.py:258passesmodel.mpu()for aPipelineModule, which is thePipelineParallelGrid(runtime/pipe/module.py:566)runtime/engine.py:2048copies it intogroups.mpuutils/groups.py:707-711then returnsmpu.get_data_parallel_group()runtime/pipe/topology.py:436documents that as "The group of ranks within the same stage of all pipelines"
So each stage votes among itself. If experts sit on some stages and not others, two stages can reach different expert_group_names and issue different collectives, which is the divergence the vote is there to prevent. Your harness is 2 ranks with no pipe dimension, so it would not show this either way. Is _clone_world_group() what you want here, or does something already keep MoE inside a single stage?
One smaller thing on the same lines. clip_grad_norm_ is a per-step call from _take_model_step, so the vote costs an all-reduce and two .item() syncs on every fp32 clip. With no MoE configured _EXPERT_PARALLEL_GROUP is empty and sorted(... or {}) gives [] whichever way the vote goes, so reading the registry locally first would skip it, on the same "built identically on every rank" argument you use two lines down.
There was a problem hiding this comment.
You are right on both, and the first one is worse than a hole in a corner: it fails in the exact case the comment named as its reason for existing. Fixed in 1a761e8.
I checked your chain before changing anything, since I could not reproduce it — my harness has no pipe dimension, so it would not show this either way, as you said.
PipelineParallelGrid.get_data_parallel_group— "The group of ranks within the same stage of all pipelines". One stage wide.get_global_norm_of_tensors, withmpuset and MoE-TP disabled, takeselif bwc_pipeline_parallel_world_size(mpu) > 1and reduces overbwc_pipeline_parallel_group(mpu)→get_pipe_parallel_group(), "The group of ranks within the same pipeline". That one spans stages.
So a stage that votes "experts" issues that pipeline-group all-reduce inside get_norm_with_moe_layers, and a stage that votes "no experts" never does — mismatched calls on one cross-stage communicator. I also checked nothing refuses MoE with a PipelineModule (no assert in pipe/engine.py or pipe/module.py), and the pipeline branch in the norm code only makes sense if the combination is meant to work, so I do not think anything keeps MoE inside one stage. The vote now reduces over _clone_world_group(). Its lazy new_group is itself a collective, but entry is decided by the registry, which every rank holds identically, so all ranks reach it together.
This is verified by reading, not by running. If you have a PP box, the case worth a real run is experts on some stages and not others.
On the cost, you were right and it is now skipped without MoE. Counted per call on the non-MoE path, 2 ranks:
master all_reduce=1 item=0
338ed51 all_reduce=2 item=1
1a761e8 all_reduce=1 item=0
One correction to the count, for the record: without MoE it was one .item() rather than two, because vote[0].item() > 0 and ... short-circuits when the first slot is zero. Still a host sync on every fp32 clip, so the conclusion stands.
test_runtime_utils.py: 13 passed, and the three cases from @gss10282023 (idle experts, a rank owning no expert, expert-only inf) still return on both ranks.
There was a problem hiding this comment.
Both read right to me at 1a761e8.
The part I wanted to check is the new early exit, since entry to the collective now depends on a local read where it used to be unconditional. It holds. _EXPERT_PARALLEL_GROUP has exactly two writers, _create_expert_and_data_parallel (utils/groups.py:393,403,461) and _create_expert_data_and_model_parallel (:588), and every write sits behind an if rank in ranks. The enclosing loops partition the whole world though, in both functions and on all three topology paths, so a rank that calls either one always lands in exactly one expert group. On top of that dist.new_group has to be entered by every rank in the job, so a run where one rank built expert groups and another skipped them would have hung at construction rather than reaching a clip. Empty and non-empty cannot disagree across ranks at that point.
The dense stages the world vote now sends down the expert path are covered by the setdefault(group_name, [zero]) from 338ed51, and _get_expert_parallel_group(exp_name) resolves on those ranks for the same reason.
Read, not run. Still no PP box here, so the experts-on-some-stages case is as untested by me as by you.
Both from @ebarkhordar's review of 338ed51. **The vote was one stage wide.** It was reduced over `groups._get_data_parallel_group()`, and under pipeline parallelism that returns `PipelineParallelGrid`'s data parallel group, documented as "the group of ranks within the same stage of all pipelines". Each stage therefore voted among itself. The expert norm the vote chooses between then reduces over the pipeline group, which spans stages (`get_global_norm_of_tensors`, the "MoE-TP disabled, sum over pipeline parallel" branch), so a stage holding experts and a stage holding only dense layers issued different collectives on one communicator -- the divergence the vote exists to prevent, in the exact case its comment named. Reduce it over `_clone_world_group()` instead. Its lazy `new_group` is a collective too, but every rank reaches it together, since entry is decided by the registry below. **The vote cost a collective without MoE.** `clip_grad_norm_` runs every step from `_take_model_step`. `_EXPERT_PARALLEL_GROUP` is built identically on every rank, so an empty registry already settles the choice, and the vote now runs only when it is not. Counted per call on the non-MoE path, 2 ranks: master all_reduce=1 item=0 338ed51 all_reduce=2 item=1 this all_reduce=1 item=0 Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
|
@tohtana @sfc-gh-truwase Could one of you take a look when you have time? The branch merges cleanly with master. @ebarkhordar reviewed both revisions and ran the cases (thread above): the per-step cost is back to master's, and the pipeline-parallel vote now reduces over the world group. With MoE and fp32 gradient clipping, expert gradients were counted twice in the global norm, so clipping triggered early and scaled every update too small. |
1a761e8 to
f92eef4
Compare
0z5a
left a comment
There was a problem hiding this comment.
Independent two-GPU RTX 5090 validation at head f92eef4b63277c87eaf2a3aca72581e196c043e6: the upstream TestClipGradNorm passed 6/6 cases. I added an active-clipping check with one replicated dense gradient and four sharded expert gradients; it passed with the expected global L2 norm sqrt(124) and clip coefficient 2/(sqrt(124)+1e-6). Every dense and expert gradient matched its scaled reference after clipping.
The runs used the PR's exact changed source and test over installed DeepSpeed 0.19.7. This validates the two-rank path; a real PP=2 × EP=2 topology was not available in this run.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Collective ordering and empty-gradient ranks can still cause incorrect behavior or failure.
Review effort: Balanced
Findings: 2
Open (3)
What changed in this PR
Corrects FP32 gradient clipping so sharded expert gradients contribute once to the global norm.
Changes:
- Routes expert norms through the precise MoE norm helper.
- Handles idle and expert-only ranks.
- Adds distributed norm regressions.
| File | Description |
|---|---|
deepspeed/runtime/utils.py |
Separates dense and expert norm computation. |
tests/unit/runtime/test_runtime_utils.py |
Tests sharded, idle, and infinity-norm cases. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| expert_tensors: Dict[str, List[torch.Tensor]] = {} | ||
| if expert_group_names: | ||
| for p in parameters: | ||
| if is_moe_param(p): | ||
| expert_tensors.setdefault(p.group_name, []).append(p.grad.data) | ||
| # Every rank has to reduce over every group, whether or not its own experts | ||
| # were used this step; a zero contributes nothing to the sum of p-th powers | ||
| # and nothing to a max over absolute values. | ||
| zero = torch.zeros(1, device=get_accelerator().current_device_name()) | ||
| for group_name in expert_group_names: | ||
| expert_tensors.setdefault(group_name, [zero]) |
| # behaviour the averaging path already had, which is to scale them to zero. | ||
| if moe_norm == -1: | ||
| moe_norm = float('inf') | ||
| total_norm = torch.tensor([float(moe_norm)], device=parameters[0].device, dtype=torch.float) |
|
|
||
| assert gathered_norm[0] == gathered_norm[1], "norm at rank 0 does not match the norm at rank 1" | ||
|
|
||
| def test_sharded_experts_count_once(self): |
… clip_grad_norm_ Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
|
please address the codex comments |
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
|
Thanks. I only see the three Copilot comments on this PR, no codex ones, so let me know if I missed something. The expert group order and the result device are fixed, and I added an engine-level test: backward and step on 2 GPUs, checked against a global norm computed separately. 16 passed on 2 H20s. |


Fixes #8475.
(Originally opened against #8469, which @gss10282023 withdrew and resubmitted as #8475 while this was being written. The discussion about how this and #8476 relate is on #8469 and continues on #8475.)
The measurement
clip_grad_norm_ends by averaging the per-rank norms over the data parallel group:That reconstructs a global norm only while every rank holds the same parameters. Expert parameters are not replicated, so each rank's norm describes a different set of experts. Four expert gradients across two ranks, measured on 2×H20:
sqrt(120))sqrt(20)andsqrt(100)averaged to 7.236 instead ofsqrt(20 + 100). The same gradients, 34% apart, decided only by where the experts happen to live.clip_coefismax_norm / total_norm, so wherever clipping is active this changes the update.@gss10282023 reported this (#8475) and measured a 0.2 shift in the clipping coefficient on their setup; the numbers above are an independent reproduction from a different direction.
Why the fix is the repo's own
get_norm_with_moe_layersalready does it correctly — it takes the non-expert norm plus the expert gradients keyed by group, reduces each expert group over its ownmoe_ep_group, and combines as(Σ norm_i ** p) ** (1/p). Both the bf16 and the fp16 optimizer already call it.The fp32 path instead inlines the approximation. That approximation also exists under its own name,
get_norm_with_moe_layers_fast, carrying this comment:and that named copy has no callers at all — so the warning never reaches the one place still doing it. This PR keeps expert gradients out of the replicated norm and folds them in through the precise helper, leaving fp32 the same shape as the other two optimizers.
Scope
Unchanged when no expert parameter is present — that branch is the original code, not touched. Every non-MoE fp32 run behaves exactly as before, which the existing suite confirms (below).
Two guards worth naming:
group_namehas no expert parallel group to reduce over.TestClipGradNorm::test_gatherbuilds one by hand exactly like that, so such a call stays on the original path rather than guessing a group.get_norm_with_moe_layersreports a non-finite norm as-1. Used as-is that would makeclip_coefnegative and flip the sign of every gradient, so it is mapped toinf— which keeps what the averaging path already did, scale to zero.Testing
tests/unit/runtime/test_runtime_utils.py, 2 GPUs (H20): 16 passed.The cases that call
clip_grad_norm_directly assert the norm value, not that a code path was reached, and cover experts spread over ranks, a rank whose experts got no tokens, a rank owning no expert, a rank with no gradient at all, the expert-only inf norm, and the order the expert groups are visited in.TestClipFp32GradientsThroughEnginegoes throughdeepspeed.initialize,backwardandstepin an fp32 run without ZeRO, with one expert on each of two ranks and a dense layer in front. The engine drops the normclip_grad_norm_returns, so the test reads it from its effect: with plain SGD the update islr * clip_coef * grad, so it computes the global norm independently from the gradients (each dense gradient once, each rank's expert once) and checks every parameter afterstep()againststart - lr * clip_coef * grad. It also checks that the clip is active and that the average of the ranks' own norms, the old answer, differs from the true one by more than 5%, so it cannot pass by accident. On this setup the true norm is 0.624 and the average is 0.542.It uses the plain
MoElayer, not AutoEP: the code under test readsis_moe_paramandgroup_nameand reduces over the expert parallel group, which both paths share.Earlier, on the first revision,
test_runtime_utils.py+tests/unit/v1/moe/test_moe.pycompared by failure set on the same box:The 6 were pre-existing there,
ImportError: libcudart.so.13, unrelated to this change. I have not reruntest_moe.pysince the later commits.Note
#8470 proposed a fix for the same issue and was closed by its author without discussion. This takes a different route: rather than repairing the averaging, it routes the fp32 path through the combiner the other two optimizers already use, so there is one implementation of the correct thing instead of two.
Environment
NVIDIA H20-3e, driver 580.105.08, Python 3.11.15, torch 2.9.1+cu128; the runs above used 2 GPUs.