Fix BLOOM DeepSpeed inference issue

[younesbelkada]

What does this PR do?

This PR tries to address a strange behaviour observed when inferring bloom-176 model using DeepSpeed!
My intuitions are:

• In the previous code we used -10000 for the attention mask filling value whereas we should use fp32.min as it is written in the original cuda kernel of FusedScaledSoftmax. This might lead to inconsistent result between the old version and the new version, but the new version should be considered as the correct one
• @RezaYazdaniAminabadi discovered that attention scores should not be multiplied by the attention mask after the softmax, which makes sense and could fix the issue

cc @RezaYazdaniAminabadi @stas00 @thomasw21

[younesbelkada]

@RezaYazdaniAminabadi did you tried to infer by removing the elementwise multiplication after the softmax as proposed in the PR?
When trying to infer on 8xA100 80GB I did obtained the same generations using the old code vs the new one with batch_size=1

[RezaYazdaniAminabadi]

@RezaYazdaniAminabadi did you tried to infer by removing the elementwise multiplication after the softmax as proposed in the PR? When trying to infer on 8xA100 80GB I did obtained the same generations using the old code vs the new one with batch_size=1
Hi @younesbelkada,

I did try this on 16 A100-40GB previously and it was not giving similar results. I will try with this one and let you know. Anyhow, I think that multiply is not needed since the scores are already masked.
Thanks

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[younesbelkada]

Thank you very much @RezaYazdaniAminabadi !!

[younesbelkada]

Finally after doing some tests it appears that we need the multiplication with the attention mask because of the following:
in some cases we have an attention mask like the one below

0 0 0 0 0 
0 1 0 0 0
0 1 1 0 0 
0 1 1 1 0 
0 1 1 1 1

After replacing all zeros by torch.finfo(dtype.min) , the softmax will return the following on the first row:
0.2 0.2 0.2 0.2 0.2 because we have the same values on the first row. To avoid using these wrong values on the calculation later I had to multiply the attention scores by the original mask.

cc @NouamaneTazi

[stas00]

@younesbelkada, ok, so we have the first row of 0.2 0.2 0.2 0.2 0.2 let's follow through to the end - where does that manifest an issue?

Let's perhaps use a small concrete example and use it to document why things are done the way they are - otherwise everybody will keep on questioning why this is done this way.

[thomasw21]

Is this because of padding, we should not care about the padding row, ie when the padding is the query. The wrong values don't matter when they are in the padding no?

[younesbelkada]

My guess was this will impact the computation of the context_layer tensor here in the case we have padded inputs as mentioned by @thomasw21
So at the end you are right ! Indeed it impacts the computation of this tensor but I think that it does not matter at all. At the end we get a token-to-token correspondance for the computed hidden states - ie the context layer will have a shape batch_size x seq_len x hidden_dim and the hidden states corresponding to the padding tokens will not impact anyway the prediction of the next token. Do you think that this explanation makes sense?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.