Fix MultiHeadAttention GQA with input normalization

[adhocyyf]

When use gqa with input normalization，both key_layer and value_layer inputs are not applying normalization.

[adhocyyf]

I recommend hide whitespace to review this pr: https://github.com/NVIDIA/TransformerEngine/pull/378/files?diff=unified&w=1

[cyanguwa]

Hi @adhocyyf,

Thanks for submitting this PR. If I understand correctly, the only component missing in our main is the self.norm for hidden_states when self.attention_type == "self" and self.num_gqa_groups != self.num_attention_heads. In that case, we're currently going through the same path as self.attention_type == "cross", which does not apply normalization for hidden_states (hence affecting the calculation for query_layer, key_layer and value_layer).

The code changes in this PR seem a little repetitive. If my understanding above is correct, would it be sufficient to just add the self.norm here:

TransformerEngine/transformer_engine/pytorch/attention.py

Line 1652 in 3b7b7c6

else:

 if self.attention_type == "cross":
input_tensor = encoder_output
else:
hidden_states = self.norm(hidden_states)
input_tensor = hidden_states

Thanks,
Charlene

[adhocyyf]

Hi @cyanguwa

thanks for your suggestion，I simplified my PR

[cyanguwa]

@adhocyyf thanks for the changes. Please see the comment regarding fusing layernorm and q/kv linears in #394 .

[cyanguwa]

Closing the PR as relevant changes have been made in #434 . Thanks for your contribution!!