Skip to content

Count each expert once in the fp32 gradient-clipping norm - #8478

Open
alanhuangyoo wants to merge 8 commits into
deepspeedai:masterfrom
alanhuangyoo:fix/fp32-clip-moe-norm
Open

alanhuangyoo wants to merge 8 commits into
deepspeedai:masterfrom
alanhuangyoo:fix/fp32-clip-moe-norm

Conversation

@alanhuangyoo

@alanhuangyoo alanhuangyoo commented Sep 10, 2026 •

Copy link
Copy Markdown
Contributor

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:

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 describes a different set of experts. Four expert gradients across two ranks, measured on 2×H20:

layout global norm
reference — each expert counted once 10.954452 (sqrt(120))
replicated on both ranks 10.954452
one half per rank 7.236068

sqrt(20) and sqrt(100) averaged to 7.236 instead of sqrt(20 + 100). The same gradients, 34% apart, decided only by where the experts happen to live. clip_coef is max_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_layers already does it correctly — it takes the non-expert norm plus the expert gradients keyed by group, reduces each expert group over its own moe_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:

This implementation standardizes the grad_norm across ranks. A more precise implementation can be found in 'get_norm_with_moe_layers'.

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:

  • 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.

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.

TestClipFp32GradientsThroughEngine goes through deepspeed.initialize, backward and step in an fp32 run without ZeRO, with one expert on each of two ranks and a dense layer in front. The engine drops the norm clip_grad_norm_ returns, so the test reads it from its effect: with plain SGD the update is lr * 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 after step() against start - 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.

this branch                         16 passed
branch's test file on the code      TestClipFp32GradientsThroughEngine fails:
before this PR                          AssertionError: Tensor-likes are not close!
                                        Greatest absolute difference: 3.2e-05 (up to 1e-06 allowed)

It uses the plain MoE layer, not AutoEP: the code under test reads is_moe_param and group_name and 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.py compared by failure set on the same box:

master   7 failed, 29 passed, 14 skipped
branch   6 failed, 30 passed, 14 skipped

only failing on branch (regressions) -> none
only failing on master               -> test_sharded_experts_count_once   (the new one)
failing on both                      -> 6

The 6 were pre-existing there, ImportError: libcudart.so.13, unrelated to this change. I have not rerun test_moe.py since 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.

`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>
@alanhuangyoo

Copy link
Copy Markdown
Contributor Author

Do not review this as it stands — two defects found by @gss10282023 on #8469, both of which I reproduced.

  1. Collective deadlock when one rank holds no expert gradient. I branch on whether this rank has expert gradients and the two branches issue different collectives, so a rank whose experts got no tokens this step waits on a peer that took the other path. p.grad is not None filtering makes that ordinary in MoE training, not a corner case.

  2. Exception on expert-only norm_type=inf: norm_parameters is empty and torch.stack([]) raises.

Marking this as a draft while I fix both. The branch decision has to be rank-independent — _EXPERT_PARALLEL_GROUP is built identically everywhere, so iterating its group names and contributing a zero tensor where this rank holds nothing keeps the collectives symmetric no matter which experts got tokens.

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.

@alanhuangyoo
alanhuangyoo marked this pull request as draft September 11, 2026 05:00
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>
@alanhuangyoo

Copy link
Copy Markdown
Contributor Author

@gss10282023 Both fixed in 338ed51, and both reproduced on 2×H20 before I touched anything.

Deadlock. Your diagnosis was exactly right, and your own patch already had the answer in a comment I should have read more carefully:

Inspect ownership before filtering gradients: an unused expert on one rank must not make that rank choose a different collective from the others.

I was reading ownership after p.grad is not None filtering. Ownership now comes from the parameters.

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 _EXPERT_PARALLEL_GROUP (built identically on every rank). A rank holding 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.

Measured, same harness that hung before:

uneven  (rank1 owns no expert at all)     both ranks return 4.4721
nograd  (rank1's expert got no tokens)    both ranks return 4.4721
inf     (expert-only, norm_type=inf)      both ranks return 2.0000

sqrt(2**2 + 4**2) — the non-expert 2.0 replicated, plus one expert of norm 4.0 owned by rank 0. The original sharded-expert number is unchanged at 10.954452.

Expert-only inf norm was torch.stack([]) with nothing replicated left to take a max over; zero is the identity there.

Three tests added for your cases. On master, run individually, all four fail with AssertionError — no hangs, since master has no expert branch to diverge on. (Run together they cascade into TimeoutError: a failed assertion in one rank of a DistributedTest leaves the other waiting. Worth knowing if you see that in your own runs.)

On the split: having read #8476, I do not think either subsumes the other. Yours keys on ds_zero_placement_family == "autoep_expert" with its own ds_autoep_ep_size, norm_type == 2 and mpu is None; mine keys on is_moe_param/group_name and routes through get_norm_with_moe_layers for every norm type. Different families, different coverage.

What I would rather not do is leave clip_grad_norm_ with two special cases bolted on. If AutoEP experts can be given a group_name and an entry in _EXPERT_PARALLEL_GROUP, the helper route covers both and the ownership weighting falls out of the group size — but you know the AutoEP placement model far better than I do, so tell me whether that holds. If it does not, I will take yours as the shape and fold legacy MoE into it.

Un-drafting since the two defects are closed, but I am in no hurry for this to land ahead of that conversation.

Comment thread deepspeed/runtime/utils.py Outdated
],
device=get_accelerator().current_device_name(),
dtype=torch.float)
dist.all_reduce(vote, group=groups._get_data_parallel_group())

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:258 passes model.mpu() for a PipelineModule, which is the PipelineParallelGrid (runtime/pipe/module.py:566)
  • runtime/engine.py:2048 copies it into groups.mpu
  • utils/groups.py:707-711 then returns mpu.get_data_parallel_group()
  • runtime/pipe/topology.py:436 documents 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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, with mpu set and MoE-TP disabled, takes elif bwc_pipeline_parallel_world_size(mpu) > 1 and reduces over bwc_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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@alanhuangyoo

Copy link
Copy Markdown
Contributor Author

@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.

@0z5a 0z5a left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Collective ordering and empty-gradient ranks can still cause incorrect behavior or failure.

Review effort: Balanced
Findings: 2 High severity · 1 Low severity

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.

Comment thread deepspeed/runtime/utils.py Outdated
Comment on lines +432 to +442
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])
Comment thread deepspeed/runtime/utils.py Outdated
# 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>
@pengdurice

Copy link
Copy Markdown
Contributor

please address the codex comments

Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
@alanhuangyoo

Copy link
Copy Markdown
Contributor Author

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.

@pengdurice
pengdurice enabled auto-merge October 22, 2026 15:44
@pengdurice
pengdurice added this pull request to the merge queue Oct 22, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 22, 2026
@pengdurice
pengdurice added this pull request to the merge queue Oct 25, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 25, 2026
@pengdurice
pengdurice enabled auto-merge October 26, 2026 23:25
@pengdurice
pengdurice added this pull request to the merge queue Oct 26, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 27, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] AutoEP FP32 gradient clipping averages norms across ranks with different expert ownership

5 participants