Skip to main content
`Ctrl`+`K`

# Language Model QAT Using Quark and Trainer

This document provides examples of Quantization-Aware Training (QAT) for language models using Quark.

Note

For information on accessing Quark PyTorch examples, refer to Accessing PyTorch Examples. This example and the relevant files are available at `/torch/language_modeling/llm_qat`.

## Supported Models

Model Name | WEIGHT-ONLY (INT4.g128)
--- | ---
microsoft/Phi-3-mini-4k-instruct | ✓
THUDM/chatglm3-6b | ✓

## Preparation

Please install the required packages before running QAT by executing `pip install -r requirements.txt`. To evaluate the model, install the necessary dependencies by running `pip install -r ../llm_eval/requirements.txt`. If an NCCL timeout error occurs while saving the model during the program’s execution, you can try installing the accelerate==1.4.0 version to resolve it.

(Optional) For LLM models, download the Hugging Face checkpoint.

## QAT Scripts

You can run the following Python scripts in the `examples/torch/language_modeling/llm_qat` path. Here, Phi-3-mini-4k-instruct is used as an example.

### Recipe 1: QAT Finetuning ChatGLM and Export to Safetensors using FSDP (The configuration variable ONLY_TRAIN_SCALING_FACTOR determines whether to only make scaling factor trainable, while keeping other parameters frozen)

    SECONDS=0
    log_file=${log_dir}/llm_qat_${model_name}_finetune.log
    output_dir="./quantized_model/chatglm_6b"
    NUM_GPUS=4
    BATCH_SIZE_PER_GPU=2
    TOTAL_BATCH_SIZE=32
    GRADIENT_ACC_STEPS=$(($TOTAL_BATCH_SIZE/$NUM_GPUS/$BATCH_SIZE_PER_GPU))

    FSDP_CONFIG=./fsdp_configs/chatglm_fsdp_config.json
    NUM_EPOCHS=5
    LR=2e-5
    MAX_SEQ_LEN=512

    CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --num_processes=${NUM_GPUS} main.py \
                        --fsdp "full_shard auto_wrap" \
                        --fsdp_config ${FSDP_CONFIG} \
                        --model ${MODEL_DIR} \
                        --model_trust_remote_code \
                        --quant_scheme w_uint4_asym \
                        --group_size 128 \
                        --finetune_dataset wikitext \
                        --num_train_epochs ${NUM_EPOCHS} \
                        --learning_rate ${LR} \
                        --finetune_seqlen ${MAX_SEQ_LEN} \
                        --per_device_train_batch_size ${BATCH_SIZE_PER_GPU} \
                        --per_device_eval_batch_size ${BATCH_SIZE_PER_GPU} \
                        --model_export hf_format \
                        --output_dir $finetune_checkpoint \
                        --model_export_dir ${output_dir} \
                        --gradient_accumulation_steps ${GRADIENT_ACC_STEPS} \
                        --skip_evaluation 2>&1| tee $log_file
    date -ud "@$SECONDS" "+Time elapsed: %H:%M:%S" |tee -a ${log_file}
    TOTAL_TIME=$((TOTAL_TIME+SECONDS))

### Recipe 2: Reload and Evaluate QAT Finetuned Model

    SECONDS=0
    log_file=${log_dir}/llm_qat_${model_name}_test_finetuned.log
    EVAL_BATCH=4
    export CUDA_VISIBLE_DEVICES=5
    EVAL_TASK=wikitext,winogrande,mmlu
    EVAL_OUTPUT_PATH=./${model_name}_${EVAL_TASK//,/_}_quantized_eval_results
    python main.py \
            --model ${MODEL_DIR} \
            --output_dir $finetune_checkpoint \
            --model_trust_remote_code \
            --skip_finetune \
            --model_reload \
            --import_model_dir $output_dir \
            --eval_result_output_path ${EVAL_OUTPUT_PATH} \
            --per_device_eval_batch_size ${EVAL_BATCH} \
            --eval_task ${EVAL_TASK} 2>&1| tee $log_file
    date -ud "@$SECONDS" "+Time elapsed: %H:%M:%S" | tee -a ${log_file}
    TOTAL_TIME=$((TOTAL_TIME+SECONDS))

### Recipe 3: Evaluate Original Non-Quantized Model

    EVAL_TASK=wikitext,winogrande,mmlu
    EVAL_OUTPUT_PATH=./${model_name}_${EVAL_TASK//,/_}_non_quantized_eval_results
    SECONDS=0
    EVAL_BATCH=4
    log_file=${log_dir}/llm_qat_${model_name}_test_bf16.log
    export CUDA_VISIBLE_DEVICES=4
    python main.py \
            --model ${MODEL_DIR} \
            --output_dir $finetune_checkpoint \
            --model_trust_remote_code \
            --skip_quantization \
            --skip_finetune \
            --eval_result_output_path ${EVAL_OUTPUT_PATH} \
            --per_device_eval_batch_size ${EVAL_BATCH} \
            --eval_task ${EVAL_TASK} 2>&1| tee ${log_file}
    date -ud "@$SECONDS" "+Time elapsed: %H:%M:%S" |tee -a ${log_file}
    TOTAL_TIME=$((TOTAL_TIME+SECONDS))

### Results on Phi-3-mini-4k-instruct

Model Name | Wikitext PPL (Quark) | Wikitext PPL (LLM harness) | MMLU | Winogrande
--- | --- | --- | --- | ---
BF16 | 6.19 | 10.32 | 68.59 | 74.42
QAT Trainer | 6.21 | 11.51 | 65.97 | 73.24

Model Name | Wikitext PPL (LLM harness) | MMLU | Winogrande
--- | --- | --- | ---
BF16 | 10.32 | 68.59 | 74.42
QAT Trainer | 11.51 | 65.97 | 73.24

### Results on ChatGLM3-6B

Model Name | Wikitext PPL (Quark) | Wikitext PPL (LLM harness) | MMLU | Winogrande
--- | --- | --- | --- | ---
BF16 | 29.93 | 51.30 | 50.45 | 62.35
QAT Trainer | 9.84 | 29.97 | 49.36 | 65.50

Model Name | Wikitext PPL (LLM harness) | MMLU | Winogrande
--- | --- | --- | ---
BF16 | 51.30 | 50.45 | 62.35
QAT Trainer | 29.97 | 49.36 | 65.50