Integrate Habana flash attention to Llama2-70B finetune

[mandy-li]

As title

[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.

[regisss]

LGTM!

I ran the command line from the example with and without --use_flash_attention:

• Without:

  max_memory_allocated (GB)   =      73.87
  memory_allocated (GB)       =       21.6
  

• With:

  max_memory_allocated (GB)   =      76.01
  memory_allocated (GB)       =      26.91
  

Memory consumption should be lower with flash attention no?

[mandy-li]

@regisss , that is a good observation. With recompute in flash attention disabled, we are supposed to have better performance, but memory consumption is more. I made another commit to add recompute arg to the script to allow user to configure that option.

Here are performance and memory usage comparisons:

Configuration	Performance	Memory
wo/ Habana Flash Attention	'train_runtime': 2669.7156, 'train_samples_per_second': 2.458,	max_memory_allocated (GB):  73.99 memory_allocated (GB): 21.6
w/ Habana Flash Attention + recompute=False	'train_runtime': 2625.176, 'train_samples_per_second': 2.499,	max_memory_allocated (GB): 76.39  memory_allocated (GB):  26.91
w/ Habana Flash Attention + recompute=True	'train_runtime': 2839.2379, 'train_samples_per_second': 2.358,	max_memory_allocated (GB): 67.42 memory_allocated (GB):  21.53
w/ Habana Flash Attention + recompute=True + PT_HPU_SDPA_BATCH_NUMHEADS_SLICE=0	'train_runtime': 2675.4069, 'train_samples_per_second': 2.45,	max_memory_allocated (GB):  77.83 memory_allocated (GB):  31.93

[regisss]

LGTM!