Enable Flash Attention in recompute and causal modes

[wszczurekhabana]

This is a follow-up on: huggingface#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

[puneeshkhanna]

This seems as a counter intuitive argument for inferencing.

[puneeshkhanna]

If there is no performance penalty and memory is also saved then we can internally pass it as True for 1st token when flash attention is used.

[wszczurekhabana]

Discussed offline. Here is summary:
We should left it controllable, cause it may be reasonable to turn it off for example on finetuning.
Additionnaly this parameter may cause some slight overhead even on inference, so maybe it's reasonable to turn it off if we don't need it.

[puneeshkhanna]

Can we forcefully set this to True when batch size is 1 and when use_flash_attention is passed. We can add it to help text that this will be taken care.

[wszczurekhabana]

Discussed offline. Here is summary:
The improvements in memory are to achieve higher batch size and thus increase throughput. Improvements from setting causal=True (and using triangular softmax underneath) will scale with batch size and seq-lenght. So causal on bigger batches is still usefull, but we need to be mindfull that without speciall handling of padding tokens it requires same lenght inputs in a batch.

[mandy-li]

@wszczurekhabana , please provide llama2-70b finetuning test results.

[wszczurekhabana]

@mandy-li please see the finetuning results updated in the description of the PR.

[mandy-li]

@mandy-li please see the finetuning results updated in the description of the PR.

Thanks @wszczurekhabana . Since FT perf with flash_attention_causal_mask is better, can you add this arg to the command example in README?

[wszczurekhabana]

Hi @mandy-li , I've added the argument to README in latest commit.

[puneeshkhanna]

Changes look good to me. @dvarshney-habana - we can merge.

[astachowiczhabana]

huggingface#862