Skip to content

Enable Flash Attention in recompute and causal modes (#21)#861

Closed
wszczurekhabana wants to merge 1 commit into
huggingface:mainfrom
wszczurekhabana:flash_attention_causal_recompute
Closed

Enable Flash Attention in recompute and causal modes (#21)#861
wszczurekhabana wants to merge 1 commit into
huggingface:mainfrom
wszczurekhabana:flash_attention_causal_recompute

Conversation

@wszczurekhabana
Copy link
Copy Markdown
Contributor

Cherry-pick from: HabanaAI#21

Original description:
This is a follow-up on: #623

where main issue is that when running with Flash Attention in causal mode (required for performance and memory optimizations on 1st token) it will generate a triangular attention mask - same on each batch of the input data.
If we have more than one batch of sequences that have different lengths, tokenizer will include a padding for those sentences that have smaller sequence lengths than 'max input tokens'. In the case of Flash Attention in causal mode, where triangular attention mask is applied to all the inputs, this will effectively mean that padding tokens are also attended to, which will in turn result in junk output being generated

This PR propagates the modes of running to run_generation.py so that optimization can still be used in case of the same length inputs in a batch. This is controlled through: --flash_attention_recompute and --flash_attention_causal_mask.

Additionally this PR provides a way to pass real input data to the model from Project Gutenberg Books for easier testing of large sequence lengths.

Below are the throughput measurements for different ratios of prompt to max seq length:

Ratio Max input tokens Max new tokens Batch size Throughput [tokens/s]
97% 31744 1042 12 85.54
75% 24576 8192 16 336.29
50% 16384 16384 24 521.39
25% 8192 24576 36 708.42

example of the command tested:
python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ --num_beams 1 --attn_softmax_bf16 --model_name_or_path meta-llama/Llama-2-70b-hf \ --warmup 2 --n_iterations 3 --use_hpu_graphs --use_kv_cache --max_input_tokens 31744 --max_new_tokens 1042 --bf16 --batch_size 12 --reuse_cache --trim_logits --limit_hpu_graphs --use_flash_attention --flash_attention_recompute --flash_attention_causal_mask --book_source

Results on finetuning:

No Flash Attention:
'train_runtime': 2499.5658, 'train_samples_per_second': 2.626

Flash Attention:
'train_runtime': 2487.0323, 'train_samples_per_second': 2.636

Flash Attention Causal:
'train_runtime': 2449.3563, 'train_samples_per_second': 2.686

* Enable Flash Attention in recompute and causal modes

* Add flash_attention_causal_mask to generation utils

* Propagate Flash Attention causal_mask to finetuning example

* Modify README example and provide additional description

* Add flash_attention_causal_mask to FT README
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant