Repository navigation
Sum GPTBigCode's shared key/value gradient across tensor-parallel ranks - #8928
Achyuthan-S wants to merge 1 commit into
Conversation
With multi-query attention every rank holds the whole key/value block of the fused c_attn but attends with only its own query heads, so its gradient for those rows is partial. Nothing summed them, so the ranks applied different updates to rows that must stay identical and training departed from the unsharded model. All-reduce the gradient of those rows in c_attn's weight and bias across the tensor-parallel group. Fixes deepspeedai#8927 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.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. |
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
![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: 0bdd7b73ac
ℹ️ 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".
|
|
||
| def sum_kv_rows(grad): | ||
| grad = grad.clone() | ||
| dist.all_reduce(grad.narrow(0, kv_start, grad.shape[0] - kv_start), group=group) |
There was a problem hiding this comment.
With the default positive gradient_clipping (1.0), this new all-reduce leaves the full summed K/V gradient in every TP copy. config_tp_params marks the entire fused weight and bias as model-parallel, so clip_grad_norm_ and the ZeRO 1/2/3 norm routines include each local tensor and then sum norms across the model-parallel group; consequently the logical K/V squared norm is counted tp_world_size times while the sharded query rows are counted once. Whenever clipping activates, the coefficient is therefore too small and AutoTP updates no longer match the unsharded model; the added test explicitly sets clipping to 0.0, so it cannot detect this default-path numerical regression. The replicated segment needs to be excluded on all but one TP rank (without excluding each rank's query segment) when computing the norm.
AGENTS.md reference: AGENTS.md:L37-L38
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good point, verified. With the default clipping at TP 2, the global norm now counts the replicated K/V rows on every rank: 1.5416 against the true 1.5382 at step 0 on the repro model, so when clipping activates the coefficient is slightly too small. With clipping on, the loss gap against single-device training is 3.3e-5 with this PR versus 5.3e-4 on master, so the default path is much closer but not exact. On master the norm wasn't exact for these rows either: it summed each rank's partial-gradient norm instead of taking the norm of their sum.
Counting rows that are replicated inside a sharded parameter once needs the TP grad-norm code to know about them. That check (is_model_parallel_parameter(p) or mp_rank == 0) is decided per parameter in runtime/utils.py, both ZeRO stages, bf16_optimizer.py and fp16/fused_optimizer.py, so I'd rather do it as a follow-up than widen this fix into the core optimizer paths. I've noted the limitation in the description.
Fixes #8927.
With AutoTP training, multi-query GPTBigCode's fused
c_attngives every rank its slice of the query rows plus the whole key/value block (_bigcode_type_transpose). Each rank attends with only its own query heads, so the gradient it computes for the shared key/value rows covers those heads alone. Nothing summed those partial gradients. Each rank therefore applied a different update to rows that must stay identical, and training departed from the unsharded model.Change
fused_LinearLayernow registers a gradient hook onc_attn's weight and bias for thebigcodetypelayout. The hook all-reduces (sums) the key/value rows across the tensor-parallel group.Test
TestGPTBigCodeMultiQueryTrainingintests/unit/v1/autotp/test_autotp_training.py(world sizes 2 and 4) trains a tiny multi-query GPTBigCode with AutoTP and with plain PyTorch on the same data. It checks three things:c_attn(weight and bias) on every rank equal the single-device ones.atol=1e-5).On
master, the first check fails: 471 of 512 key/value-weight gradient elements differ, by up to 1.4e-2.The test uses plain SGD with
gradient_clipping: 0.0to keep the comparison clean. DeepSpeed clips gradients to norm 1.0 by default and the reference does not, and Adam turns floating-point noise on gradients that are zero in the unsharded model into visible updates.Results
The script from #8927 with
gradient_clipping: 0.0and SGD (lr 0.5), 6 steps, compared with single-device training:mastermasterValidation
macOS 26.5 (arm64, CPU), Python 3.9.6, torch 2.8.0, transformers 4.44.2,
DS_ACCELERATOR=cpu LOCAL_SIZE=4.master.test_autotp_training.py,test_autotp_custom_patterns.py,test_affine_shard_map.pyandtest_autotp_uc_checkpoint.py: 146 passed, 42 failed, 2 skipped.master: 145 passed and the same 42 failures.KeyError: 'flags'from building the CPU Adam op on macOS arm64.yapf0.43.0 is clean.flake8is clean on the changed lines; the 3F824warnings inlayers.pyare onmastertoo and come from a newer pyflakes than the pinned 5.0.4.Not covered here
c_attnis marked as sharded. The clipping coefficient is therefore slightly too small. On the repro model at TP 2 with clipping on, the loss gap against single-device training is 3.3e-5 with this PR and 5.3e-4 onmaster. Counting replicated rows once needs the norm code inruntime/utils.py, both ZeRO stages,bf16_optimizer.pyandfp16/fused_optimizer.pyto know about them, which I'd rather do as a follow-up.multi_query=FalseGPTBigCode still fails in the forward pass under AutoTP (the "Related" note in [BUG] AutoTP training of multi-query GPTBigCode diverges from single-device training at TP > 1 #8927). That is a separate change.torch.compile.