Refactoring LLama Attention and mlp layers

[bgoldberg-habana]

Module for scope linearAllreduce
this change allows better memory consumption and better optimizations in synapse when running llama 70b on deepspeed

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

[MrGeva]

LGTM

[regisss]

I just left a last comment that will be addressed quickly.

Besides, do you have numbers to see the kind of memory that is saved doing this?

[bgoldberg-habana]

cmd line -
ENABLE_SYNAPSE_QUANTIZATION=false USE_DEFAULT_QUANT_PARAM=true UPDATE_GRAPH_OUTPUT_MME=false ENABLE_CALC_DYNAMIC_RANGE=false ENABLE_EXPERIMENTAL_FLAGS=true deepspeed --num_gpus 8 run_generation.py --model_name_or_path /mnt/weka/data/pytorch/llama2/Llama-2-70b-hf/ --use_hpu_graphs --use_kv_cache --kv_cache_fp8 --batch_size 50 --fp8 --reuse_cache --trim_logits --n_iterations 5 --attn_softmax_bf16 --limit_hpu_graphs --max_new_tokens 2048 --max_input_tokens 2048

pay attention i'm running already on 1.14 but i don't think the numbers changed much from 1.13

with change -
Throughput (including tokenization) = 1581.191910099665 tokens/second
Number of HPU graphs = 333
Memory allocated = 19.07 GB
Max memory allocated = 49.15 GB
Total memory available = 94.62 GB
Graph compilation duration = 524.4125659640013 seconds

reference
Throughput (including tokenization) = 1257.5571168775869 tokens/second
Number of HPU graphs = 333
Memory allocated = 27.33 GB
Max memory allocated = 87.02 GB
Total memory available = 94.62 GB
Graph compilation duration = 542.6321858290012 seconds