DeepSpeed

9148ab22 - Sum GPTBigCode's shared key/value gradient across tensor-parallel ranks (#8928)

Sum GPTBigCode's shared key/value gradient across tensor-parallel ranks (#8928) Fixes #8927. With AutoTP training, multi-query GPTBigCode's fused `c_attn` gives 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_LinearLayer` now registers a gradient hook on `c_attn`'s weight and bias for the `bigcodetype` layout. The hook all-reduces (sums) the key/value rows across the tensor-parallel group. - Query rows stay per rank. - The input gradient is unchanged: the column-parallel backward already sums every rank's contribution to it. - The hook is registered only when the TP group has more than one rank and the parameter requires grad. ## Test `TestGPTBigCodeMultiQueryTraining` in `tests/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: 1. After the first backward, the key/value-row gradients of `c_attn` (weight and bias) on every rank equal the single-device ones. 2. Over 4 SGD steps, the losses match single-device training (`atol=1e-5`). 3. The key/value rows stay bit-identical across ranks. 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.0` to 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.0` and SGD (lr 0.5), 6 steps, compared with single-device training: | | TP | ZeRO | Max loss gap | Max cross-rank diff of the K/V rows | |---|---|---|---|---| | `master` | 2 | 0, 1, 2 | 1.2e-4 | 9.1e-3 | | `master` | 4 | 0, 1, 2 | 1.9e-4 | 5.9e-3 | | this PR | 2 | 0, 1, 2 | 0 | 0 | | this PR | 4 | 0, 1, 2 | 4.8e-7 | 0 | ## Validation macOS 26.5 (arm64, CPU), Python 3.9.6, torch 2.8.0, transformers 4.44.2, `DS_ACCELERATOR=cpu LOCAL_SIZE=4`. - The new test passes at world sizes 2 and 4, and fails on `master`. - `test_autotp_training.py`, `test_autotp_custom_patterns.py`, `test_affine_shard_map.py` and `test_autotp_uc_checkpoint.py`: 146 passed, 42 failed, 2 skipped. - `master`: 145 passed and the same 42 failures. - All 42 are `KeyError: 'flags'` from building the CPU Adam op on macOS arm64. - `yapf` 0.43.0 is clean. `flake8` is clean on the changed lines; the 3 `F824` warnings in `layers.py` are on `master` too and come from a newer pyflakes than the pinned 5.0.4. - Not run on GPU. The change is in which gradients are summed, not in the communication backend. ## Not covered here - **Gradient clipping.** With clipping active (DeepSpeed's default is 1.0), the global gradient norm counts the replicated K/V rows once per TP rank, because the TP norm decides per parameter whether to count it, and `c_attn` is 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 on `master`. Counting replicated rows once needs the norm code in `runtime/utils.py`, both ZeRO stages, `bf16_optimizer.py` and `fp16/fused_optimizer.py` to know about them, which I'd rather do as a follow-up. - `multi_query=False` GPTBigCode still fails in the forward pass under AutoTP (the "Related" note in #8927). That is a separate change. - The new test does not exercise ZeRO-3 or `torch.compile`. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>