Repository: microsoft/DeepSpeed
File: deepspeed/sequence/layer.py (Python)

Code before:
        seq_len, h, batch_size = input.shape
        num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size)
        local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()]
        h_dim = h // local_heads
        local_seq_len = seq_len // seq_world_size

        input = input.view(seq_len * h, batch_size)
        local_seq_len_with_heads = int(input.shape[0] / seq_world_size)  # dim size of local_seq_len*local_heads*hdim
        input_splits = [local_seq_len_with_heads] * seq_world_size
        coeff = local_seq_len_with_heads // local_heads  #per head: dim size of local_seq_len*hdim

        #uneven seq_world_size coeff,  total_heads/local_heads.
        heads_scale_coeff = get_num_kv_heads() / local_heads

        output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list]
        output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads)
        total_h = int(inp_shape[gather_idx] * heads_scale_coeff)
        output = torch.empty(output_buff_d1_size, input.shape[1], device=input.device, dtype=input.dtype)
        dist.all_to_all_single(output,input,output_split_sizes=output_splits, \
            input_split_sizes=input_splits,group=group)
        ##################
        #suppose 7 heads divide into 4 ranks [2,2,2,1]
        #chunk_num_heads_small=floor(7/4)=1
        #chunk_num_heads_large=ceil(7/4)=2
        #num_chunk_heads_large=len([2,2,2])=3, all2all_buffer_counts
        #num_chunk_heads_large=len([1])=1, all2all_buffer_counts
        #total_num_large_heads=sum([2,2,2])=7
        #total_num_small_heads=sum([1])=1

        chunk_num_heads_small = get_num_kv_heads() // seq_world_size  # even heads compatible
        chunk_num_heads_large = chunk_num_heads_small + 1
        num_chunk_heads_large = get_num_kv_heads() % seq_world_size
        num_chunk_heads_small = seq_world_size - num_chunk_heads_large
        total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large
        total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small

        heads_large_combine_size = coeff * total_num_large_heads
        heads_small_combine_size = coeff * total_num_small_heads
        heads_large_chunk, heads_small_chunk = output.split([heads_large_combine_size, heads_small_combine_size],
                                                            dim=0)
        heads_large_chunk = heads_large_chunk.view(num_chunk_heads_large, local_seq_len, chunk_num_heads_large, h_dim,
                                                   batch_size)
        heads_small_chunk = heads_small_chunk.view(num_chunk_heads_small, local_seq_len, chunk_num_heads_small, h_dim,
                                                   batch_size)
        if batch_dim_idx == 0:
            #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[batch,local_seq_len,all2all_buffer_counts*n_heads,dim]
            order = [4, 1, 0, 2, 3]
            heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(batch_size, local_seq_len,
                                                                                   total_num_large_heads, h_dim)
            heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(batch_size, local_seq_len,
                                                                                   total_num_small_heads, h_dim)

Review comment (suggestion):
```suggestion
        #num_chunk_heads_small=len([1])=1, all2all_buffer_counts
```

Code after:
        seq_len, h, batch_size = input.shape
        num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size)
        local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()]
        h_dim = h // local_heads
        local_seq_len = seq_len // seq_world_size

        input = input.view(seq_len * h, batch_size)
        local_seq_len_with_heads = int(input.shape[0] / seq_world_size)  # dim size of local_seq_len*local_heads*hdim
        input_splits = [local_seq_len_with_heads] * seq_world_size
        coeff = local_seq_len_with_heads // local_heads  #per head: dim size of local_seq_len*hdim

        #uneven seq_world_size coeff,  total_heads/local_heads.
        heads_scale_coeff = get_num_kv_heads() / local_heads

        output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list]
        output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads)
        total_h = int(inp_shape[gather_idx] * heads_scale_coeff)
        output = torch.empty(output_buff_d1_size, input.shape[1], device=input.device, dtype=input.dtype)
        dist.all_to_all_single(output,input,output_split_sizes=output_splits, \
            input_split_sizes=input_splits,group=group)
        ##################
        #suppose 7 heads divide into 4 ranks [2,2,2,1]
        #chunk_num_heads_small=floor(7/4)=1
        #chunk_num_heads_large=ceil(7/4)=2
        #num_chunk_heads_large=len([2,2,2])=3, all2all_buffer_counts
        #num_chunk_heads_small=len([1])=1, all2all_buffer_counts
        #total_num_large_heads=sum([2,2,2])=7
        #total_num_small_heads=sum([1])=1

        chunk_num_heads_small = get_num_kv_heads() // seq_world_size  # even heads compatible
        chunk_num_heads_large = chunk_num_heads_small + 1
        num_chunk_heads_large = get_num_kv_heads() % seq_world_size
        num_chunk_heads_small = seq_world_size - num_chunk_heads_large
        total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large
        total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small

        heads_large_combine_size = coeff * total_num_large_heads
        heads_small_combine_size = coeff * total_num_small_heads
        heads_large_chunk, heads_small_chunk = output.split([heads_large_combine_size, heads_small_combine_size],
                                                            dim=0)
        heads_large_chunk = heads_large_chunk.view(num_chunk_heads_large, local_seq_len, chunk_num_heads_large, h_dim,
                                                   batch_size)
        heads_small_chunk = heads_small_chunk.view(num_chunk_heads_small, local_seq_len, chunk_num_heads_small, h_dim,
                                                   batch_size)
        if batch_dim_idx == 0:
            #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[batch,local_seq_len,all2all_buffer_counts*n_heads,dim]
            order = [4, 1, 0, 2, 3]
            heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(batch_size, local_seq_len,
                                                                                   total_num_large_heads, h_dim)
            heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(batch_size, local_seq_len,
                                                                                   total_num_small_heads, h_dim)