Uh oh!
There was an error while loading. Please reload this page.
Add support for multi-query and grouped-query attention - #338
Conversation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
cyanguwa
commented
Jul 21, 2023
/te-ci |
cyanguwa
commented
Jul 21, 2023
/te-ci |
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ksivaman
commented
Jul 28, 2023
/te-ci |
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
cyanguwa
commented
Jul 28, 2023
/te-ci |
blahBlahhhJ
commented
Jul 28, 2023
Hi, overall LGTM with only one small thing. I see @ksivaman mentioned that this PR uses a PR in MegatronLM as a reference. The main difference I see is that in TE you use a Q projection and a KV projection while in megatron it's packed to be a single QKV projection matrix. Is there any particular reasons why you separate them? I think it would be better if the implementations are consistent so that people don't need to convert checkpoints back and forth. especially because I see megatron core uses some TE layers internally. |
cyanguwa
commented
Jul 29, 2023
Hi @blahBlahhhJ , With MQA and GQA, the kv projection size might be different from the query projection size. That itself isn't a problem for the projection gemm, but when we do ColumnParallel and the split of parameters (into q, k, v, in TE), the |
* add support for multi-query/grouped-query attention Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert to flash-attn 1.0.6 and build 2.0.0.post1 manually in CI Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add keyword name for DPA input Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fused attn tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix skipif for pytest Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Update transformer_engine/pytorch/attention.py Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Update tests/pytorch/test_fused_attn.py Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix TP and SP case Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * add skipifs for pytest Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove higher limit for flash-attn version Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
This PR adds support for multi-query attention (MQA) and grouped-query attention (GQA). Specifically, it enables the use of MQA and GQA in FlashAttention() and UnfusedDotProductAttention. MQA and MHA are special cases of GQA, i.e. when
num_gqa_groups=1andnum_heads. GQA only applies to keys and values, not queries.