Grouped Query Attention + Refactor Attn

[sashaDoubov]

Adding grouped query attention to LLM-foundry, and making the GQA class a superclass of MQA and MHA attention, as it is a generalization of those two variants of attention.
Things to note:

• we currently use repeat_interleave to make the grouped query tensor the same dimensions as multi-head attention, which does allocate new memory, compared to using expand. This can be updated in the future, but is the safer bet for now, given that we previously saw edge-cases with using expand vs repeat for particular head_dim settings causing NaNs
• Part of this change is also changing how we initialize the QKV matrix, which used to initialize (n_heads * head_dim, d_model), but now is changed to initialize each n_head (head_dim, d_model) matrix separately

[vchiley]

GeneralizedAttention is a valid name, but should be call it GroupedQueryAttention to not confuse ppl?

or alternatively create

class GroupedQueryAttention(GeneralizedAttention):
    def __init__(..., groups=G, ...):
        super().__init__(..., kv_n_heads=G, ...)

or something like that (following convention of other impl)

[sashaDoubov]

GeneralizedAttention is a valid name, but should be call it GroupedQueryAttention to not confuse ppl?

or alternatively create

class GroupedQueryAttention(GeneralizedAttention):
    def __init__(..., groups=G, ...):
        super().__init__(..., kv_n_heads=G, ...)

or something like that (following convention of other impl)

I've renamed it to GroupedQueryAttention, as per discussion offline

[dakinggg]

Could you please include a comparison run for a normal mha model before and after this pr? I'd like to make sure we don't have a perf regression (or at least know the quantity of it)

[dakinggg]

few nits, but lgtm pending the performance comparison

[sashaDoubov]

Could you please include a comparison run for a normal mha model before and after this pr? I'd like to make sure we don't have a perf regression (or at least know the quantity of it)

Here are some perf numbers, which look good outside of the loss spikes. However, I did see loss spikes with both "before and after" the change: