huggingface_example.sh

#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -e

script_dir="$(dirname "$(readlink -f "$0")")"

source $script_dir/parser.sh
parse_options "$@"

set -x

# This will prevent the script from hanging on Selene/EOS due to the MPI support.
echo "********** unset all SLURM_, PMI_, PMIX_ Variables **********"
for i in $(env | grep ^SLURM_ | cut -d"=" -f 1); do unset -v $i; done
for i in $(env | grep ^PMI_ | cut -d"=" -f 1); do unset -v $i; done
for i in $(env | grep ^PMIX_ | cut -d"=" -f 1); do unset -v $i; done

# Fail on errors inside pipelines (e.g. `python eval.py | tee result.txt`), otherwise a crashing
# eval is masked by tee's exit code and the script passes silently.
set -o pipefail

if [ -z "$MODEL_PATH" ]; then
    echo "Unsupported model argument: Expected a huggingface model path or model name" >&2
    exit 1
fi

# Check if ENABLE_SPARSITY environment variable is set to "true"
if [ "$SPARSITY_FMT" = "dense" ]; then
    ENABLE_SPARSITY=false
else
    ENABLE_SPARSITY=true
fi

case $SPARSITY_FMT in
dense | sparsegpt) ;;
*)
    echo "Unknown sparsity argument: Expected one of: [dense, sparsegpt]" >&2
    exit 1
    ;;
esac

# Quant format / recipe validation is delegated to hf_ptq.py.

script_dir="$(dirname "$(readlink -f "$0")")"

pushd $script_dir/..

if [ -z "$ROOT_SAVE_PATH" ]; then
    ROOT_SAVE_PATH=$(pwd)
fi

QFORMAT_MODIFIED="${QFORMAT//,/_}"

# When using --recipe, build the model name from the recipe basename (without
# directory or .yaml suffix) so each recipe gets its own SAVE_PATH.
if [ -n "$RECIPE" ]; then
    RECIPE_TAG=$(basename "$RECIPE" .yaml | sed 's/[^0-9a-zA-Z\-]/_/g')
    MODEL_NAME=$(basename "$MODEL_PATH" | sed 's/[^0-9a-zA-Z\-]/_/g')_recipe_${RECIPE_TAG}
else
    MODEL_NAME=$(basename "$MODEL_PATH" | sed 's/[^0-9a-zA-Z\-]/_/g')_${QFORMAT_MODIFIED}${KV_CACHE_QUANT:+_kv_${KV_CACHE_QUANT}}
fi

SAVE_PATH=${ROOT_SAVE_PATH}/saved_models_${MODEL_NAME}

MODEL_CONFIG=${SAVE_PATH}/config.json

mkdir -p $SAVE_PATH

if [ "${REMOVE_EXISTING_MODEL_CONFIG,,}" = "true" ]; then
    rm -f $MODEL_CONFIG
fi

PTQ_ARGS=""

if $CALIB_WITH_IMAGES; then
    PTQ_ARGS+=" --calib_with_images "
fi

if [ "$LOW_MEMORY_MODE" = "true" ]; then
    PTQ_ARGS+=" --low_memory_mode "
fi

# AutoQuantize runs via an AutoQuantize --recipe or the deprecated --auto_quantize_bits CLI path.
# Auto-generate a checkpoint path (to save/restore the search state) when the user didn't supply one.
if [ -z "$AUTO_QUANTIZE_CHECKPOINT" ] && { [[ "$RECIPE" == *auto_quantize* ]] || [ -n "$AUTO_QUANTIZE_BITS" ]; }; then
    AUTO_QUANTIZE_CHECKPOINT="${ROOT_SAVE_PATH}/auto_quantize_checkpoints/${MODEL_NAME}.pth"
    mkdir -p "$(dirname "$AUTO_QUANTIZE_CHECKPOINT")"
    echo "Auto-generated auto_quantize checkpoint path: $AUTO_QUANTIZE_CHECKPOINT"
fi
if [ -n "$AUTO_QUANTIZE_CHECKPOINT" ]; then
    PTQ_ARGS+=" --auto_quantize_checkpoint=$AUTO_QUANTIZE_CHECKPOINT "
fi

# Deprecated AutoQuantize CLI flags: passed through to hf_ptq.py, which converts them into an
# AutoQuantizeConfig on the fly. Prefer an AutoQuantize --recipe.
if [ -n "$AUTO_QUANTIZE_BITS" ]; then
    PTQ_ARGS+=" --auto_quantize_bits=$AUTO_QUANTIZE_BITS "
    PTQ_ARGS+=" --auto_quantize_method=${AUTO_QUANTIZE_METHOD:-gradient} "
    PTQ_ARGS+=" --auto_quantize_score_size=${AUTO_QUANTIZE_SCORE_SIZE:-128} "
    PTQ_ARGS+=" --auto_quantize_cost_model=${AUTO_QUANTIZE_COST_MODEL:-weight} "
    if [ -n "$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO" ]; then
        PTQ_ARGS+=" --auto_quantize_active_moe_expert_ratio=$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO "
    fi
fi

if [ -n "$CALIB_DATASET" ]; then
    PTQ_ARGS+=" --dataset=$CALIB_DATASET "
fi

if [ -n "$KV_CACHE_QUANT" ]; then
    PTQ_ARGS+=" --kv_cache_qformat=$KV_CACHE_QUANT "
fi

if $TRUST_REMOTE_CODE; then
    PTQ_ARGS+=" --trust_remote_code "
fi

if $USE_SEQ_DEVICE_MAP; then
    PTQ_ARGS+=" --use_seq_device_map "
fi

if [ -n "$GPU_MAX_MEM_PERCENTAGE" ]; then
    PTQ_ARGS+=" --gpu_max_mem_percentage=$GPU_MAX_MEM_PERCENTAGE "
fi

if [ -n "$CALIB_SEQ" ]; then
    PTQ_ARGS+=" --calib_seq=$CALIB_SEQ "
fi

if [ -n "$MOE_CALIB_EXPERTS_RATIO" ]; then
    PTQ_ARGS+=" --moe_calib_experts_ratio=$MOE_CALIB_EXPERTS_RATIO "
fi

if $CAST_MXFP4_TO_NVFP4; then
    PTQ_ARGS+=" --cast_mxfp4_to_nvfp4 "
fi

if ! $VERBOSE; then
    PTQ_ARGS+=" --no-verbose "
fi

AWQ_ARGS=""
if [ -n "$AWQ_BLOCK_SIZE" ]; then
    AWQ_ARGS+="--awq_block_size=$AWQ_BLOCK_SIZE"
fi

if [[ -f $MODEL_CONFIG ]] || [[ -f "$SAVE_PATH/encoder/config.json" && -f "$SAVE_PATH/decoder/config.json" ]]; then
    MODEL_CONFIG_EXIST=true
else
    MODEL_CONFIG_EXIST=false
fi

if [[ $TASKS =~ "quant" ]] || [[ ! -d "$SAVE_PATH" ]] || [[ ! $(ls -A $SAVE_PATH) ]]; then

    if [[ "$MODEL_CONFIG_EXIST" == false ]]; then
        echo "Quantizing original model..."
        if [ -n "$RECIPE" ]; then
            QUANT_SPEC_ARGS="--recipe=$RECIPE"
        else
            QUANT_SPEC_ARGS="--qformat=${QFORMAT// /,}"
        fi
        python hf_ptq.py \
            --pyt_ckpt_path=$MODEL_PATH \
            --export_path=$SAVE_PATH \
            --sparsity_fmt=$SPARSITY_FMT \
            $QUANT_SPEC_ARGS \
            --calib_size=$CALIB_SIZE \
            --batch_size=$CALIB_BATCH_SIZE \
            --inference_tensor_parallel=$TP \
            --inference_pipeline_parallel=$PP \
            $PTQ_ARGS \
            $AWQ_ARGS
    else
        echo "Quantized model config $MODEL_CONFIG exists, skipping the quantization stage"
    fi

    # for enc-dec model, users need to refer TRT-LLM example for deployment
    if [[ -f "$SAVE_PATH/encoder/config.json" && -f "$SAVE_PATH/decoder/config.json" && ! -f $MODEL_CONFIG ]]; then
        echo "Please continue to deployment with the TRT-LLM enc_dec example, https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec. Checkpoint export_path: $SAVE_PATH"
        exit 0
    fi

    if [[ "$SPARSITY_FMT" != "dense" ]]; then
        echo "Sparse quantization detected (SPARSITY_FMT=$SPARSITY_FMT). Please deploy with the TRT-LLM using trtllm-build. Checkpoint export_path: $SAVE_PATH"
        exit 0
    fi

    if [[ "$QFORMAT" == *"nvfp4"* ]] || [[ "$KV_CACHE_QUANT" == *"nvfp4"* ]] || [[ "$RECIPE" == *"nvfp4"* ]]; then
        cuda_major=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0 | cut -d. -f1)

        if [ "$cuda_major" -lt 10 ]; then
            echo "Please deploy the NVFP4 checkpoint on a Blackwell GPU. Checkpoint export_path: $SAVE_PATH"
            exit 0
        fi
    fi

    if [ -n "$RECIPE" ]; then
        echo "Recipe $RECIPE used. Please deploy with TensorRT-LLM directly. Checkpoint export_path: $SAVE_PATH"
        exit 0
    fi

    if [[ ! " fp8 nvfp4 bf16 fp16 " =~ " ${QFORMAT} " ]]; then
        echo "Quant $QFORMAT specified. Please read TensorRT-LLM quantization support matrix https://nvidia.github.io/TensorRT-LLM/features/quantization.html#quantization-in-tensorrt-llm and use TensorRT-LLM for deployment. Checkpoint export_path: $SAVE_PATH"
        exit 0
    fi

    if $TRUST_REMOTE_CODE; then
        RUN_ARGS+=" --trust_remote_code "
    fi

    # Only run the deploy+generate smoke test when "quant" is explicitly requested. Eval tasks
    # (lm_eval/mmlu/simple_eval) deploy the checkpoint themselves, so it is redundant there.
    if [[ $TASKS =~ "quant" ]]; then
        if $VLM; then
            # VLMs use the TRT-LLM multimodal quickstart for the deploy smoke test.
            if [ -z "$TRT_LLM_CODE_PATH" ]; then
                TRT_LLM_CODE_PATH=/app/tensorrt_llm # default path for the TRT-LLM release docker image
                echo "Setting default TRT_LLM_CODE_PATH to $TRT_LLM_CODE_PATH."
            fi
            QUICK_START_MULTIMODAL=$TRT_LLM_CODE_PATH/examples/llm-api/quickstart_multimodal.py
            if [ -f "$QUICK_START_MULTIMODAL" ]; then
                python3 "$QUICK_START_MULTIMODAL" --model_dir "$SAVE_PATH" --modality image
            else
                echo "Warning: $QUICK_START_MULTIMODAL cannot be found. Please set TRT_LLM_CODE_PATH to the TRT-LLM code path or test the quantized checkpoint $SAVE_PATH with the TRT-LLM repo directly."
            fi
        else
            python run_tensorrt_llm.py --checkpoint_dir="$SAVE_PATH" $RUN_ARGS
        fi
    fi
fi

if [[ -d "${MODEL_PATH}" ]]; then
    MODEL_ABS_PATH=$(realpath ${MODEL_PATH})
else
    # model_path is a HF reference, not a local directory, no need to make the path absolute
    MODEL_ABS_PATH=${MODEL_PATH}
fi

if [[ $TASKS =~ "lm_eval" ]]; then

    if [ -z "$LM_EVAL_TASKS" ]; then
        echo "lm_eval_tasks not specified"
        exit 1
    fi

    lm_eval_flags=""
    if [[ "$LM_EVAL_TASKS" == *"llama"* ]]; then
        # Flags instructed by https://github.com/EleutherAI/lm-evaluation-harness/tree/main/lm_eval/tasks/llama3#paper
        lm_eval_flags+=" --fewshot_as_multiturn --apply_chat_template "
    fi

    if [ -n "$LM_EVAL_LIMIT" ]; then
        lm_eval_flags+=" --limit $LM_EVAL_LIMIT "
    fi

    if $TRUST_REMOTE_CODE; then
        lm_eval_flags+=" --trust_remote_code "
    fi

    LM_EVAL_RESULT=${SAVE_PATH}/lm_eval.txt
    echo "Evaluating lm_eval, result saved to $LM_EVAL_RESULT..."

    pushd ../llm_eval/

    pip install -r requirements.txt

    echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE"

    python lm_eval_tensorrt_llm.py \
        --model trt-llm \
        --model_args tokenizer=$MODEL_PATH,checkpoint_dir=$SAVE_PATH,max_gen_toks=$BUILD_MAX_OUTPUT_LEN \
        --tasks $LM_EVAL_TASKS \
        --batch_size $BUILD_MAX_BATCH_SIZE $lm_eval_flags | tee $LM_EVAL_RESULT

    popd

fi

if [[ $TASKS =~ "mmlu" ]]; then

    MMLU_RESULT=${SAVE_PATH}/mmlu.txt
    echo "Evaluating MMLU, result saved to $MMLU_RESULT..."

    pushd ../llm_eval/

    pip install -r requirements.txt

    if [ -z "$MMLU_DATA_PATH" ]; then
        MMLU_DATA_PATH=data/mmlu
    fi
    if [[ ! -d "$MMLU_DATA_PATH" ]] || [[ ! $(ls -A $MMLU_DATA_PATH) ]]; then
        echo "Preparing the MMLU test data"
        wget https://people.eecs.berkeley.edu/~hendrycks/data.tar -O /tmp/mmlu.tar
        mkdir -p data
        tar -xf /tmp/mmlu.tar -C data && mv data/data $MMLU_DATA_PATH
    fi

    mmlu_flags=""
    if [ -n "$MMLU_LIMIT" ]; then
        mmlu_flags+=" --limit $MMLU_LIMIT "
    fi

    python mmlu.py \
        --model_name causal \
        --model_path $MODEL_ABS_PATH \
        --checkpoint_dir $SAVE_PATH \
        --data_dir $MMLU_DATA_PATH $mmlu_flags | tee $MMLU_RESULT
    popd

fi


if [[ $TASKS =~ "livecodebench" || $TASKS =~ "simple_eval" ]]; then
    # Clean a previous session if exists
    pkill -f "trtllm-serve" && while pgrep -f "trtllm-serve" >/dev/null; do sleep 1; done
    HASH=$(echo -n "$SAVE_PATH" | md5sum | awk '{print $1}')
    PORT=$((10000 + (0x${HASH:0:4} % 50001)))
    echo "Starting trtllm-serve on $PORT"
    trtllm-serve $SAVE_PATH --host 0.0.0.0 --port $PORT >$SAVE_PATH/serve.txt 2>&1 &
    SERVE_PID=$!

    # Poll the log instead of `tail -f | while ... break`: under `set -o pipefail` (set above),
    # breaking out of that pipeline leaves tail to die by SIGPIPE, which would abort the script.
    while ! grep -q "Application startup complete" $SAVE_PATH/serve.txt 2>/dev/null; do
        if ! kill -0 $SERVE_PID 2>/dev/null; then
            echo "trtllm-serve has exited."
            exit 1
        fi
        sleep 2
    done
    echo "Application startup complete."

    pushd ../llm_eval/

    if [[ $TASKS =~ "livecodebench" ]]; then
        echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE"
        bash run_livecodebench.sh $MODEL_NAME $BUILD_MAX_BATCH_SIZE $BUILD_MAX_OUTPUT_LEN $PORT | tee $SAVE_PATH/livecodebench.txt
        mkdir -p $SAVE_PATH/livecodebench
        mv LiveCodeBench/output/$MODEL_NAME/* $SAVE_PATH/livecodebench
        echo "LiveCodeBench results are saved under $SAVE_PATH/livecodebench."

    fi

    if [[ $TASKS =~ "simple_eval" ]]; then
        echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE"
        bash run_simple_eval.sh $MODEL_NAME $SIMPLE_EVAL_TASKS $BUILD_MAX_OUTPUT_LEN $PORT $SIMPLE_EVAL_LIMIT | tee $SAVE_PATH/simple_eval.txt
        echo "Simple eval results are saved under $SAVE_PATH/simple_eval.txt."
    fi

    popd

    kill $SERVE_PID
fi


popd