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

[wszczurekhabana]

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

[regisss]

Why adding use_fused_rope in the changes of this PR? You need it for recompute and causal mode?

[regisss]

@wszczurekhabana I changed the target of this PR to main because we first need to merge this change there and then I'll cherry-pick it in the release branch.
Can you make sure your main branch is up to date, run git merge main in your working branch and then solve the merge conflicts please?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.