[QUESTION] Does Megatron-LM supports Flash Attention for BERT and T5 Pretraining?

[Leo-T-Zang]

My question

Does Megatron-LM supports Flash Attention for BERT and T5 Pretraining? If so, where is the code specifically supports such feature?

Thanks!!!

[elliottnv]

@shanmugamr1992 Please help answering the question. Thank you!

[shanmugamr1992]

Megatron lm , when you use mcore models will support flash attention in the next couple of weeks.
The flash attention is provided by a library called Transformer engine that we use . Not sure what you mean by T5 Pretraining ?

[Leo-T-Zang]

Thanks for your response!

I mean, when we train a BERT or T5 model using Megatron LM, how do we enable flash attention? Is this supported? I ask this question because some other libraries, like HuggingFace, support flash attention for LLaMA/GPT but not for BERT or T5.