add longbench and static_attention_dtype (#2393)

* add longbench and static_attention_dtype

Signed-off-by: He, Xin3 <xin3.he@intel.com>

Author: xin3he

examples/pytorch/nlp/huggingface_models/language-modeling/quantization/auto_round/llama3/README.md

@@ -6,9 +6,9 @@ In this example, you can verify the accuracy on HPU/CUDA device with emulation o
 
 ```bash
 # neural-compressor-pt
-pip install neural-compressor-pt==3.7
+pip install neural-compressor-pt
 # auto-round
-pip install auto-round==0.9.3
+pip install auto-round
 # other requirements
 pip install -r requirements.txt
 ```
@@ -79,6 +79,8 @@ Notes:
 Here we provide several recipes for Llama3 models. The relative accuracy loss of quantized model should be less than 1%.
 
 > Note: You can also enable static quantization for KV cache by adding `--static_kv_dtype fp8` argument to `quantize.py`， or `--static_kv_dtype=fp8` argument to `run_quant.sh` and `run_benchmark.sh`.
+>
+> You can also enable static quantization for attention by adding `--static_attention_dtype fp8` argument to `quantize.py`， or `--static_attention_dtype=fp8` argument to `run_quant.sh` and `run_benchmark.sh`. When enabled, it automatically sets KV cache dtype to fp8 as well.
 
 #### Llama 3.1 8B MXFP8
 
@@ -210,8 +212,10 @@ CUDA_VISIBLE_DEVICES=0,1 bash run_benchmark.sh --model_path=Llama-3.1-70B-MXFP8
 
 The script automatically:
 - Detects available GPUs from `CUDA_VISIBLE_DEVICES` and sets `tensor_parallel_size` accordingly
-- Runs default tasks: `piqa,hellaswag,mmlu_llama,gsm8k_llama` with batch size 8
+- Runs default tasks: `piqa,hellaswag,mmlu_llama,gsm8k_llama` with batch size 64
 - Supports custom task selection and batch size adjustment
+- Handles special tasks like `mmlu_llama`, `gsm8k_llama` (with chat template) and `longbench` (with extended context length) automatically
+- For longbench dataset evaluation, use the `--tasks=longbench` parameter
 
 
 ### NVFP4

examples/pytorch/nlp/huggingface_models/language-modeling/quantization/auto_round/llama3/quantize.py

@@ -169,6 +169,13 @@ def get_accuracy(model_name_or_path, tokenizer=None, eval_tasks="mmlu", limit=No
         choices=["fp8", "float8_e4m3fn"],
         help="Data type for static quantize key and value.",
     )
+    parser.add_argument(
+        "--static_attention_dtype",
+        default=None,
+        type=str,
+        choices=["fp8", "float8_e4m3fn"],
+        help="Data type for static quantize key and value.",
+    )
     parser.add_argument("--use_recipe", action="store_true", help="whether to use recipe to quantize model")
     parser.add_argument("--recipe_file", type=str, default="recipes/Meta-Llama-3.1-8B-Instruct_6bits.json", help="path of recipe file")
     parser.add_argument("--iters", default=200, type=int, help="iters for autoround.")
@@ -214,9 +221,23 @@ def get_accuracy(model_name_or_path, tokenizer=None, eval_tasks="mmlu", limit=No
     model, tokenizer = initialize_model_and_tokenizer(args.model_name_or_path)
 
     if args.quantize:
-        if args.dtype in ["uNVFP4", "NVFP4+"]:
-            from auto_round.schemes import QuantizationScheme
+        from auto_round.schemes import PRESET_SCHEMES, QuantizationScheme
 
+        # Check if RCEIL versions are available and use them instead
+        use_rceil = "MXFP4_RCEIL" in PRESET_SCHEMES and "MXFP8_RCEIL" in PRESET_SCHEMES
+        if use_rceil:
+            # Replace dtype if it's MXFP4 or MXFP8
+            if args.dtype == "MXFP4":
+                args.dtype = "MXFP4_RCEIL"
+            elif args.dtype == "MXFP8":
+                args.dtype = "MXFP8_RCEIL"
+            # Replace options list entries
+            args.options = [
+                "MXFP4_RCEIL" if opt == "MXFP4" else ("MXFP8_RCEIL" if opt == "MXFP8" else opt)
+                for opt in args.options
+            ]
+
+        if args.dtype in ["uNVFP4", "NVFP4+"]:
             uNVFP4 = QuantizationScheme.from_dict(
                 {
                     "bits": 4,
@@ -256,6 +277,7 @@ def load_recipe_results(file_path):
             options=args.options,
             shared_layers=args.shared_layers,
             static_kv_dtype=args.static_kv_dtype,
+            static_attention_dtype=args.static_attention_dtype,
             enable_torch_compile=args.enable_torch_compile,
             low_gpu_mem_usage=args.low_gpu_mem_usage,
             export_format=args.export_format,

examples/pytorch/nlp/huggingface_models/language-modeling/quantization/auto_round/llama3/requirements.txt

@@ -1,7 +1,11 @@
 transformers==4.57.3
 torch==2.9.0
 torchvision==0.24.0
-lm_eval==0.4.9.2
+lm_eval==0.4.10
 datasets==4.4.2
-auto-round==0.9.3
+auto-round>=0.9.3
 neural-compressor-pt>=3.7
+jieba
+fuzzywuzzy
+rouge
+hf_transfer

examples/pytorch/nlp/huggingface_models/language-modeling/quantization/auto_round/llama3/run_benchmark.sh

@@ -7,6 +7,7 @@ TASKS="piqa,hellaswag,mmlu_llama,gsm8k_llama"
 BATCH_SIZE=64
 GPU_MEMORY_UTILIZATION=0.8
 KV_CACHE_DTYPE="auto"
+ATTN_DTYPE="auto"
 
 while [[ $# -gt 0 ]]; do
     case $1 in
@@ -30,6 +31,10 @@ while [[ $# -gt 0 ]]; do
             KV_CACHE_DTYPE="${1#*=}"
             shift
             ;;
+        --static_attention_dtype=*)
+            ATTN_DTYPE="${1#*=}"
+            shift
+            ;;
         *)
             echo "Unknown parameter: $1"
             exit 1
@@ -44,6 +49,14 @@ if [[ "$KV_CACHE_DTYPE" == "fp8" ]]; then
     echo "Using FP8 for KV cache"
 fi
 
+# for fp8 attention cache
+if [[ "$ATTN_DTYPE" == "fp8" ]]; then
+    export VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION=0
+    export VLLM_ATTENTION_BACKEND="FLASHINFER"
+    KV_CACHE_DTYPE="fp8"
+    echo "Using FP8 Attention"
+fi
+
 # Validate required parameters
 if [[ -z "$MODEL_PATH" ]]; then
     echo "Usage: bash run_benchmark.sh --model_path=<path_to_quantized_model> [--tasks=<tasks>] [--batch_size=<size>]"
@@ -103,10 +116,11 @@ run_evaluation() {
 }
 
 
-# Check if tasks contain gsm8k_llama or mmlu_llama
+# Check if tasks contain gsm8k_llama, mmlu_llama, or longbench
 NEED_SPLIT=false
 OTHER_TASKS="$TASKS"
 SPECIAL_TASKS=""
+LONGBENCH_TASK=""
 
 if [[ "$TASKS" == *"gsm8k_llama"* ]]; then
     SPECIAL_TASKS="gsm8k_llama"
@@ -122,26 +136,24 @@ if [[ "$TASKS" == *"mmlu_llama"* ]]; then
     OTHER_TASKS=$(echo "$OTHER_TASKS" | sed 's/,*mmlu_llama,*//' | sed 's/^,//' | sed 's/,$//')
     NEED_SPLIT=true
 fi
+if [[ "$TASKS" == *"longbench"* ]]; then
+    LONGBENCH_TASK="longbench"
+    OTHER_TASKS=$(echo "$OTHER_TASKS" | sed 's/,*longbench,*//' | sed 's/^,//' | sed 's/,$//')
+    NEED_SPLIT=true
+fi
 
 if [[ "$NEED_SPLIT" == true ]]; then
     if [[ -n "$OTHER_TASKS" ]]; then
         echo "Running general tasks"
         run_evaluation "$OTHER_TASKS" true ""
-        if [[ $? -eq 0 ]]; then
-            IFS=',' read -ra SPECIAL_ARRAY <<< "$SPECIAL_TASKS"
-            for special_task in "${SPECIAL_ARRAY[@]}"; do
-                echo "Running $special_task with chat template"
-                run_evaluation "$special_task" true "--apply_chat_template --fewshot_as_multiturn"
-                if [[ $? -ne 0 ]]; then
-                    echo "Benchmark failed on $special_task!"
-                    exit 1
-                fi
-            done
-        else
+        if [[ $? -ne 0 ]]; then
             echo "Skipping special tasks due to previous failure"
             exit 1
         fi
-    else
+    fi
+    
+    # Run special tasks (gsm8k_llama, mmlu_llama)
+    if [[ -n "$SPECIAL_TASKS" ]]; then
         IFS=',' read -ra SPECIAL_ARRAY <<< "$SPECIAL_TASKS"
         for special_task in "${SPECIAL_ARRAY[@]}"; do
             echo "Running $special_task with chat template"
@@ -152,6 +164,26 @@ if [[ "$NEED_SPLIT" == true ]]; then
             fi
         done
     fi
+    
+    # Run longbench task with special configuration
+    if [[ -n "$LONGBENCH_TASK" ]]; then
+        echo "Running longbench with special configuration"
+        local longbench_cmd="lm_eval --model vllm --model_args pretrained=\"$MODEL_PATH\",trust_remote_code=True,dtype=bfloat16,max_model_len=66000,tensor_parallel_size=$TENSOR_PARALLEL_SIZE,gpu_memory_utilization=$GPU_MEMORY_UTILIZATION,enable_prefix_caching=False --tasks longbench --seed 42 --batch_size $BATCH_SIZE --apply_chat_template --gen_kwargs '{\"temperature\":0.0}'"
+        echo "Executing command: $longbench_cmd"
+        
+        lm_eval --model vllm \
+            --model_args pretrained="$MODEL_PATH",trust_remote_code=True,dtype=bfloat16,max_model_len=66000,tensor_parallel_size=$TENSOR_PARALLEL_SIZE,gpu_memory_utilization=$GPU_MEMORY_UTILIZATION,enable_prefix_caching=False \
+            --tasks longbench \
+            --seed 42 \
+            --batch_size $BATCH_SIZE \
+            --apply_chat_template \
+            --gen_kwargs '{"temperature":0.0}'
+        
+        if [[ $? -ne 0 ]]; then
+            echo "Benchmark failed on longbench!"
+            exit 1
+        fi
+    fi
 else
     run_evaluation "$TASKS" true ""
 fi

examples/pytorch/nlp/huggingface_models/language-modeling/quantization/auto_round/llama3/run_quant.sh

@@ -4,6 +4,7 @@
 
 # Parse command line arguments
 KV_CACHE_DTYPE="auto"
+STATIC_ATTENTION_DTYPE="auto"
 while [[ $# -gt 0 ]]; do
     case $1 in
         --topology=*)
@@ -26,6 +27,10 @@ while [[ $# -gt 0 ]]; do
             KV_CACHE_DTYPE="${1#*=}"
             shift
             ;;
+        --static_attention_dtype=*)
+            STATIC_ATTENTION_DTYPE="${1#*=}"
+            shift
+            ;;
         *)
             echo "Unknown parameter: $1"
             exit 1
@@ -48,10 +53,12 @@ echo "  Input Model: $INPUT_MODEL"
 echo "  Output Model: $OUTPUT_MODEL"
 
 # Set common parameters
-if [ "$KV_CACHE_DTYPE" = "auto" ]; then
-    COMMON_ARGS="--quantize --enable_torch_compile --low_gpu_mem_usage --export_format auto_round"
-else
-    COMMON_ARGS="--quantize --enable_torch_compile --low_gpu_mem_usage --export_format auto_round --static_kv_dtype $KV_CACHE_DTYPE"
+COMMON_ARGS="--quantize --enable_torch_compile --low_gpu_mem_usage --export_format auto_round"
+if [ "$KV_CACHE_DTYPE" != "auto" ]; then
+    COMMON_ARGS="$COMMON_ARGS --static_kv_dtype $KV_CACHE_DTYPE"
+fi
+if [ "$STATIC_ATTENTION_DTYPE" != "auto" ]; then
+    COMMON_ARGS="$COMMON_ARGS --static_attention_dtype $STATIC_ATTENTION_DTYPE"
 fi
 
 case "$TOPOLOGY" in
@@ -81,14 +88,15 @@ case "$TOPOLOGY" in
                 ;;
             "mxfp4_mixed")
                 echo "Running Llama 3.1 8B MXFP4 (Mixed with MXFP8) quantization..."
-                CMD="python quantize.py --model_name_or_path \"$INPUT_MODEL\" $COMMON_ARGS --target_bits 7.8 --options \"MXFP4\" \"MXFP8\" --shared_layers \"k_proj\" \"v_proj\" \"q_proj\" --shared_layers \"gate_proj\" \"up_proj\" --export_path \"$OUTPUT_MODEL\""
+                CMD="python quantize.py --model_name_or_path \"$INPUT_MODEL\" $COMMON_ARGS --target_bits 7.8 --iters 0 --options \"MXFP4\" \"MXFP8\" --shared_layers \"k_proj\" \"v_proj\" \"q_proj\" --shared_layers \"gate_proj\" \"up_proj\" --export_path \"$OUTPUT_MODEL\""
                 echo "Executing command: $CMD"
                 python quantize.py \
                     --model_name_or_path "$INPUT_MODEL" \
                     $COMMON_ARGS \
                     --target_bits 7.8 \
                     --options "MXFP4" "MXFP8" \
                     --shared_layers "k_proj" "v_proj" "q_proj" \
+                    --iters 0 \
                     --shared_layers "gate_proj" "up_proj" \
                     --export_path "$OUTPUT_MODEL"
                 ;;