Repository navigation
Conversation
The fused linear layer used for this projection now stores a mark on both the weight and the bias: the first replicated row, how many rows are replicated, and the tensor-parallel world size. In the finite-norm accumulation, a model-parallel parameter carrying that mark contributes its leading rows at full power and its replicated rows at that power divided by the tensor-parallel size, so the later sum across ranks counts those rows once. Parameters without the mark still contribute their full gradient power, and each stored term is already raised to the norm power so the stacked terms are summed directly. Fixes deepspeedai#9024. Refs deepspeedai#9024 Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
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: 079df62743
ℹ️ 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".
The single-device comparison ran with clipping off, so it could not see a clip coefficient computed from a norm that counts the shared K/V rows once per rank. Run it again with gradient_clipping 0.5, which is below every step's gradient norm, and clip the reference with torch.nn.utils.clip_grad_norm_ at the same threshold. Refs deepspeedai#9024 Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Raising each parameter's norm to norm_type one at a time added a scalar kernel per parameter on every clipped FP32 step, even for models without replicated K/V rows. Unmarked parameters go back to one stacked pow, and only parameters marked with ds_tp_replicated_rows take the per-row path. Refs deepspeedai#9024 Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
|
Thanks for this fix! I’ve verified the norm accounting is correct for the path it touches.
Test environment details for reference:
A quick note on scopeThis fix specifically resolves the ZeRO-0 fp32 case. The Two possible paths forward
Small refinements to suggest
Per the guidelines in AGENTS.md, could you please run the new integration test on a GPU setup and share the hardware specs, dependency versions, and the exact run command? For context, my validation used the clipping-parity test from my branch, rather than the test case added in this PR. |
|
Thanks for running this on both A100 and 910B4. Option 2 fits best: this PR covers only the ZeRO-0 fp32 clip_grad_norm_ path, and the ZeRO 1/2/3 and fp16/bf16 norm paths need offset-aware K/V handling that belongs in a separate change, with #9024 left open for it. I don't have GPU hardware on this machine, so your A100 and 910B4 runs are the only device results for this so far, and you're right that utils.py still hardcodes 'ds_tp_replicated_rows' instead of importing DS_TP_REPLICATED_ROWS from layers.py. |
gpu-test-labeler
Bot
added
the
skip-gpu-test
Background
Tensor-parallel multi-query GPTBigCode training clips gradients more aggressively than the same step on one device. The key and value rows of the fused attention projection sit on every tensor-parallel rank, and the norm that sets the clip coefficient adds every copy, so updates are scaled down relative to single-device training. The gradient values already match that single-device run; the mismatch is the clip scale, and it grows with the key/value share of the model gradient and with the tensor-parallel degree.
Clipping totals a power of each local gradient and all-reduces that total across the tensor-parallel group, which is correct when every model-parallel row belongs to exactly one rank. The fused attention weight and bias are marked model-parallel, and each rank's shard holds its query rows plus the full key/value block that the backward hook has already reduced.
Query rows from different ranks are disjoint and combine into one copy, while the key/value block is identical on every rank, so the all-reduce counts it once per rank and the reported norm is too large.
Summary
The fused linear layer used for this projection now stores a mark on both the weight and the bias: the first replicated row, how many rows are replicated, and the tensor-parallel world size. In the finite-norm accumulation, a model-parallel parameter carrying that mark contributes its leading rows at full power and its replicated rows at that power divided by the tensor-parallel size, so the later sum across ranks counts those rows once.
Parameters without the mark still contribute their full gradient power, and each stored term is already raised to the norm power so the stacked terms are summed directly. Refs #9024.
Tests
make testlocally on macOS for the DeepSpeed fused_LinearLayer change; it fails the same way on the base branch, so the failure predates this change (it fails the same way on upstream/master).Refs #9024
AI was used for assistance.