Flash Attention v2 for Roberta-large about optimum HOT 5 CLOSED

@jamesharrisivi When using BetterTransformer, you will already get a FA2 style attention if your compute capability is high enough. Excited to see the the F.sdpa getting merged in attention - finally a hassle free implementation, potentially compatible for masking in training.

from optimum.

@jamesharrisivi SDPA (that can dispatch to FA2) support for BERT is being merged natively in Transformers: huggingface/transformers#28802

This would likely make it easy to use SDPA as well for Roberta natively there.

Going forward, we are privileging direct support of SDPA/torch compile in Transformers and will not be adding new features to BetterTransformer. You can read more here: https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one

from optimum.

Ok thank you.

I see the current implementation uses https://github.com/pytorch/pytorch/blob/0f5e24bda9450a89ba56d2fdd471f56d97fe4546/aten/src/ATen/native/transformers/transformer.cpp#L75

But this doesn't use flash attention or does it?

from optimum.

It does, it uses https://github.com/pytorch/pytorch/blob/0f5e24bda9450a89ba56d2fdd471f56d97fe4546/aten/src/ATen/native/transformers/cuda/attention.cu#L467, see https://github.com/pytorch/pytorch/blob/232f09e0ea7b7d9ac9752bd366aa2f4c372fd9d9/aten/src/ATen/native/native_functions.yaml#L14505

This one is able to dispatch to flash attention if the attn_mask argument is None (mask in the C++ code)

You could check that with pytorch profiler or for example the decorator https://pytorch.org/docs/master/backends.html#torch.backends.cuda.sdp_kernel

from optimum.

@michaelfeil Yes, masking works for training. We simply drop the mask when the is_causal argument of SDPA can be used: https://github.com/huggingface/transformers/blob/83e366bfd49708796e2c6461d3988d23d008502a/src/transformers/modeling_attn_mask_utils.py#L375-L380

from optimum.