query_states = query_states * self.norm_factor
  attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)).float()
  htcore.mark_step()

deepspeed --num_nodes 1 --num_gpus 4 --master_addr 127.0.0.1 \
        --master_port 60008 run_generation.py \
        --model_name_or_path Qwen/Qwen2-57B-A14B-Instruct \
        --trust_remote_code \
        --n_iterations 1 \
        --warmup 1 \
        --bf16 \
        --trim_logits \
        --batch_size 1 \
        --max_input_tokens 128 \
        --max_new_tokens 128 \
        --use_kv_cache \
        --reuse_cache \
        --bucket_size 128 \
        --bucket_internal \
        --use_hpu_graphs \
        --use_flash_attention

Input/outputs:
input 1: ('DeepSpeed is a machine learning framework',)
output 1: ('DeepSpeed is a machine learning framework that provides thenumber one hundred and one dalm thenumber one then\nDeepSpeed is a deep thenumber one thenumber one hundred and thenumber one�\nDeepSpeed is a machine learning framework developed by the MoE (Mixture of Experts thenumber thenumber one hundred]Human thenumber one hundred and one d thenumber one hundred and thenumber one hundred thenumber then thenumber then thenumber one hundred]�\nDeepSpeed is a machine learning framework that is used to optimize and accelerate the thenumber one hundred and thenumber one hundred] training of large language�f thenumber one then thenumber one hundred',)

deepspeed --num_nodes 1 --num_gpus 4 --master_addr 127.0.0.1 \
        --master_port 60008 run_generation.py \
        --model_name_or_path Qwen/Qwen2-57B-A14B-Instruct \
        --trust_remote_code \
        --n_iterations 1 \
        --warmup 1 \
        --bf16 \
        --trim_logits \
        --batch_size 1 \
        --max_input_tokens 128 \
        --max_new_tokens 128 \
        --use_kv_cache \
        --reuse_cache \
        --bucket_size 128 \
        --bucket_internal \
        --use_hpu_graphs

Input/outputs:
input 1: ('DeepSpeed is a machine learning framework',)
output 1: ('DeepSpeed is a machine learning framework that provides optimization for deep learning models. It is designed to make deep learning more efficient and scalable by utilizing the power of GPUs, distributed computing, and other hardware accelerators. DeepSpeed can be used with popular deep learning frameworks such as PyTorch and TensorFlow.\n\nSome of the key features of DeepSpeed include:\n\n1. Zero Redundancy Optimizer (ZeRO): This feature allows for model parallelism without any redundancy in the optimizer state, which can significantly reduce memory usage.\n2. Pipeline Parallelism: This feature allows for the partitioning of a model across multiple GPUs, enabling faster training times on large models.\n3. Dynamic',)