DeepSpeed

ece19cfe - Fix FPDT K/V gradient buffers under grouped-query attention (#8665)

Fix FPDT K/V gradient buffers under grouped-query attention (#8665) ## What breaks `_FPDTGPUAttentionImpl_` (the non-offloading FPDT attention, used when `cpu_offloading=False` or when there is a single attention chunk) allocates its K/V gradient accumulators with the **query** shape: ```python dk = [torch.zeros(global_q[0].shape, ...) for _ in range(num_chunks)] dv = [torch.zeros(global_q[0].shape, ...) for _ in range(num_chunks)] ``` With grouped-query attention (`num_key_value_heads < num_attention_heads`) the K/V chunks carry fewer heads than Q. The per-block `dk_this`/`dv_this` are already allocated from the K/V shapes, so they are accumulated into buffers with the wrong head count, and after the reverse all-to-all the projected K/V gradient no longer fits the K/V slice of the QKV weight. Backward fails, e.g. with 8 query heads, 2 KV heads, head_dim 32 and sp=2: ``` RuntimeError: The size of tensor a (64) must match the size of tensor b (256) at non-singleton dimension 0 ``` The offloading implementation (`_FPDTGPUOffloadingAttentionImpl_`) already allocates these buffers from `global_k`/`global_v`, so only the non-offloading path is affected. MHA is unaffected because Q and K/V shapes coincide. ## Fix Allocate `dk`/`dv` from the K/V chunk shapes, matching the offloading implementation. ## Tests Added `TestFPDTAttentionGQABackward` in `tests/unit/v1/sequence_parallelism/test_ulysses.py` (world size 2, `num_kv_heads` in {2, 4, 8} with 8 query heads). It runs the non-offloading implementation on load-balanced FPDT shards and compares the output and the input, QKV weight and QKV bias gradients against a full-sequence fp32 causal attention reference. Previously only the offloading implementation's forward output was tested; this is the first backward check for FPDT attention. Like the existing FPDT test, it requires flash-attn and is skipped otherwise. Local results (2x RTX 5090 D, torch 2.13 + CUDA 13.0, flash-attn 2.8.3.post1 built from source): | | kv=2 (GQA) | kv=4 (GQA) | kv=8 (MHA) | |---|---|---|---| | master | FAILED | FAILED | PASSED | | this PR | PASSED | PASSED | PASSED | On master the GQA cases fail in backward with the shape mismatch above. With the fix, in fp16, the largest gradient deviation from the fp32 reference is ~0.07 on values up to ~150 (QKV bias grad of V). Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>