BetterTransformer optimization / flash_attn{_2}

[thomaschhh]

BetterTransformer optimization / flash_attn{_2}

[guynich]

I see: I've only used "flash_attn" and needed this PR.

[sanchit-gandhi]

Closing as resolved! Feel free to re-open if you continue to encounter issues

[guynich]

This might help: #76

[thomaschhh]

Looks good for --attn_type "flash_attn" but not for --attn_type "flash_attn_2". In that case I still get the above-mentioned error.

[sanchit-gandhi]

Fixed in #101! You can now set --attn_implementation to either {"eager", "sdpa", "flash_attn_2"}:

distil-whisper/training/run_pseudo_labelling.py

Lines 136 to 139 in b948d02

	"Which attention implementation to use in the encoder and decoder attention layers. Can be one of:\n"
	"1. `eager` or `None`: default Transformers attention implementation.\n"
	"2. `sdpa`: Flash Attention through PyTorch SDPA. Requires `torch>=2.1`. Recommended for hardware where Flash Attention 2 is not supported, e.g. Turing GPUs, (T4, RTX 2080).\n"
	"3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."

The README has been updated to reflect this change: https://github.com/huggingface/distil-whisper/tree/main/training#1-pseudo-labelling

[thomaschhh]

I just looked into it again and it seems like there is a mismatch between the help string

distil-whisper/training/run_distillation.py

Line 136 in b6400a3

"3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."

and the value that is expected in L142, which is flash_attention_2 and not flash_attn_2.