Add mark step and inplace residual add in llama model code to reduce memory consumption

[puneeshkhanna]

Mark step helping in reducing workspace memory by
approx twice of (BS,seq len, hidden dim).

Inplace add helping in reducing persistent tensors by
approx twice of (BS, seq len, hidden dim).

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[puneeshkhanna]

@dvarshney-habana - please review.
@libinta - Can you please check finetuning once ?

[vivekgoe]

Please add mark_step calls under lazy mode flag. Same modeling file is used for torch compile mode also where mark_step is not relevant.

[puneeshkhanna]

@MrGeva - You may want to review this. Accuracy seems fine. However I need to address mark step comment from Vivek and also need to check finetuning script.

[puneeshkhanna]

@mandy-li - this PR is very important from memory usage perspective for llama inference.

As an example for the config of BS-172, seq len-2048, hidden dim-8191 (size is ~5.3 GB) for llama-70B on 8x.
Max memory usage (without flash attention) reduced from ~86 Gb to ~66Gb.
Max memory usage (with flash attention) reduced from ~70 GB to ~59GB.

[puneeshkhanna]

@schoi-habana - Can you check finetuning once with this PR ?

[puneeshkhanna]

@vivekgoe - lazy mode flag and check added.

[vivekgoe]

@vivekgoe - lazy mode flag and check added.

LGTM.

[puneeshkhanna]

In place add is having loss divergence issue while training so updated PR to perform the in place add operation only in inference.

Ran below command without any fixes of this PR :
PT_HPU_MAX_COMPOUND_OP_SIZE=10 DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED=1 python3 ../gaudi_spawn.py --use_deepspeed --world_size 8 run_lora_clm.py --model_name_or_path /mnt/weka/data/pytorch/llama2/Llama-2-70b-hf --deepspeed llama2_ds_zero3_config.json --dataset_name tatsu-lab/alpaca --bf16 True --output_dir ./lora_out --num_train_epochs 1 --max_seq_len 2048 --per_device_train_batch_size 10 --per_device_eval_batch_size 10 --gradient_checkpointing --evaluation_strategy epoch --eval_delay 2 --save_strategy no --learning_rate 0.0018 --warmup_ratio 0.03 --lr_scheduler_type "cosine" --logging_steps 1 --dataset_concatenation --attn_softmax_bf16 True --do_train --do_eval --use_habana --use_lazy_mode --pipelining_fwd_bwd --throughput_warmup_steps 3 --lora_rank 4 --lora_target_modules "q_proj" "v_proj" "k_proj" "o_proj" --validation_split_percentage 4 --use_flash_attention True
{
"epoch": 1.0,
"eval_accuracy": 0.791171470444553,
"eval_loss": 0.7647133469581604,
"eval_runtime": 27.0564,
"eval_samples": 125,
"eval_samples_per_second": 4.62,
"eval_steps_per_second": 0.074,
"max_memory_allocated (GB)": 81.61,
"memory_allocated (GB)": 26.91,
"perplexity": 2.148378447163791,
"total_memory_available (GB)": 94.62,
"train_loss": 0.8714751173288394,
"train_runtime": 1321.3399,
"train_samples_per_second": 2.628,
"train_steps_per_second": 0.033
}

Ran below command with the updated changes in this PR (only markstep fix will apply to finetuning) :
PT_HPU_MAX_COMPOUND_OP_SIZE=10 DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED=1 python3 ../gaudi_spawn.py --use_deepspeed --world_size 8 run_lora_clm.py --model_name_or_path /mnt/weka/data/pytorch/llama2/Llama-2-70b-hf --deepspeed llama2_ds_zero3_config.json --dataset_name tatsu-lab/alpaca --bf16 True --output_dir ./lora_out --num_train_epochs 1 --max_seq_len 2048 --per_device_train_batch_size 10 --per_device_eval_batch_size 10 --gradient_checkpointing --evaluation_strategy epoch --eval_delay 2 --save_strategy no --learning_rate 0.0018 --warmup_ratio 0.03 --lr_scheduler_type "cosine" --logging_steps 1 --dataset_concatenation --attn_softmax_bf16 True --do_train --do_eval --use_habana --use_lazy_mode --pipelining_fwd_bwd --throughput_warmup_steps 3 --lora_rank 4 --lora_target_modules "q_proj" "v_proj" "k_proj" "o_proj" --validation_split_percentage 4 --use_flash_attention True
{
"epoch": 1.0,
"eval_accuracy": 0.7912496336101612,
"eval_loss": 0.7647190690040588,
"eval_runtime": 26.8381,
"eval_samples": 125,
"eval_samples_per_second": 4.658,
"eval_steps_per_second": 0.075,
"max_memory_allocated (GB)": 81.58,
"memory_allocated (GB)": 26.91,
"perplexity": 2.148390740319044,
"total_memory_available (GB)": 94.62,
"train_loss": 0.8714751173288394,
"train_runtime": 1319.6586,
"train_samples_per_second": 2.672,
"train_steps_per_second": 0.034
}

@libinta , @schoi-habana - FYI.
log_lora_without_fixes.txt
log_lora_with_fixes.txt

[astachowiczhabana]

huggingface#833