TransformerEngine icon indicating copy to clipboard operation
TransformerEngine copied to clipboard
verified publisher icon NVIDIA

Transformer Engine using FlashAttention V3

Open heavyrain-lzy opened this issue 1 year ago • 1 comments

I find that TE don't support FA-V3. There some error when I use flash_attn==2.6.3 and transformer_engine=1.9.0 enable context parallel in Megatron-LM. Do you have the plan to support it?

heavyrain-lzy avatar Aug 21 '24 06:08 heavyrain-lzy

FA3 support is added in https://github.com/NVIDIA/TransformerEngine/pull/1019.

yaox12 avatar Aug 27 '24 02:08 yaox12

The context parallel support with FA3 is added in #1232. Please give it a try and let us know if there's any problems.

cyanguwa avatar Oct 21 '24 20:10 cyanguwa