Repository: deepspeedai/DeepSpeed
Pull request: Fix uneven head sequence parallelism bug (#6774)
Here `gather_idx < 2` represents `is_first_all2all`. During the first all2all, `uneven_head_all2all` will be called if either `num_heads % seq_world_size != 0` or `get_num_kv_heads() is None`. 

During the second all2all, it'll return return `uneven_head_all2all` if and only if `get_num_kv_heads() is None` which is always set during the first uneven all2all. This means that there will no longer be issue where `uneven_head_all2all ` is returned for the second all2all because of `num_heads % seq_world_size != 0`.

Fixes: #6774 

Problem statement:
[BUG] [Fix Suggestion] Uneven head sequence parallelism
**Describe the bug**

deepspeed 0.15.4 will think you are using unevenhead SP even though you aren't and raise the following assert:
`assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})"`

This happens because during the second all2all, the head count is already parallelized; hence, `num_heads % seq_world_size != 0` returns true.
Second all2all input: `[B, s, hc/sp, hs]`. However, not always `hc/sp % sp == 0`.
```
def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, async_op=False, handle=None, type=None):
    seq_world_size = dist.get_world_size(group)
    # we only need num_heads once
    num_heads = input.shape[2]

    if get_num_kv_heads() is not None or num_heads % seq_world_size != 0:
        # Assuming here that the number of heads for q is consistent with kv
        # If not, additional logic is required for cases like GQA
        if get_num_kv_heads() is None:
            assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})"
            # set heads at first call by num_total_heads.
            # then use ``get_num_kv_heads() is not None`` to re-entry uneven path.
            set_num_kv_heads(num_heads)
        assert async_op == False, "uneven head sp does not support async op"
        return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group)
```


**To Reproduce**

To reproduce the error, one can set the SP=head_count.
**Fix Suggestion:**
Only rely on `get_num_kv_heads()` during the second all2all:
```
    if get_num_kv_heads() is not None or num_heads % seq_world_size != 0:
```
to
```
    if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and gather_idx < 2):
```

Hints from issue discussion:
Thank you @Eugene29 , yes, your analysis is correct. Would you mind sending a patch to fix it?
Commit explanation: 
Here `gather_idx < 2` represents `is_first_all2all`. During the first all2all, `uneven_head_all2all` will be called if either `num_heads % seq_world_size != 0` or `get_num_kv_heads() is None`. 

During the second all2all, it'll return return `uneven_head_all2all` if and only if `get_num_kv_heads() is None` which is always set during the first uneven all2all. This means that there will no longer be issue where `uneven_head_all2all ` is returned for the second all2all because of `num_heads % seq_world_size != 0`.

Commit message: Fix uneven head sequence parallelism bug (#6774)

Patch to review:
diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py
index 8fd3ca918433..7a16592cbf30 100644
--- a/deepspeed/sequence/layer.py
+++ b/deepspeed/sequence/layer.py
@@ -155,7 +155,7 @@ def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, asyn
     # we only need num_heads once
     num_heads = input.shape[2]
 
-    if get_num_kv_heads() is not None or num_heads % seq_world_size != 0:
+    if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and gather_idx < 2):
         # Assuming here that the number of heads for q is consistent with kv
         # If not, additional logic is required for cases like GQA
         if get_num_kv_heads() is None:


Code review comments:
diff_hunk: @@ -155,7 +155,7 @@ def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, asyn
     # we only need num_heads once
     num_heads = input.shape[2]
 
-    if get_num_kv_heads() is not None or num_heads % seq_world_size != 0:
+    if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and gather_idx < 2):
path: deepspeed/sequence/layer.py
text: @user1:
By the way, can we use scatter_idx for the judgment here? This way, we can stay consistent with other parts by uniformly using ```scatter_idx``` to determine whether the current state is ```total_head``` or ```total_seq_len```

@user2:
Ditto @user1! Future refactor will have a better naming to differentiate between N/p (pre/post-attention) and h/p (attention block)

@author:
Hi I agree, scatter_idx booleans are confusing. Maybe I can share my opinion on the refactoring on Wed if thats ok.

@author:
Hey, I just made a quick change to use `scatter_idx`.