Likely bug in Alibi positional embeddings together with Tensor Parallelism.

[DanielHesslow]

The way alibi is implemented we first generate the alibi attention bias tensor

Megatron-DeepSpeed/megatron/model/transformer.py

Line 591 in 7ab5c05

def _build_alibi_tensor(max_seq_len, num_attention_heads, batch_size):

and then extract the part of it:

Megatron-DeepSpeed/megatron/model/transformer.py

Line 291 in 7ab5c05

matmul_result = alibi[:output_size[0]*output_size[1], :, :output_size[3]]

When we run with tensor parallelism, the attention heads will be distributed. We then need to extract the correct part of alibi tensor corresponding to our tensor parallel rank.

In order to fix this we need to:

1. Create a test that makes sure that we get the same output with and without tensor parallelism.
2. Extract the correct part of the alibi_tensor corresponding to our tensor_parallel rank.

[thomasw21]

@stas00 for visibility in case we're going for Alibi and TP > 1

[stas00]

After yesterday's call we are definitely going for TP=8.

And it does sounds like we are going for Alibi - at least that's what @ibeltagy said in the channel.

[ibeltagy]

Yes, we certainly want to use ALiBi in the final model, and we need to fix the ALiBi/TP bug.

[DanielHesslow]

Closing this since the fix has been unit tested and is now merged into main.