Skip to content

Latest commit

 

History

History

trtllm

README.md

TensorRT-LLMによるoffline推論

[ English | 日本語 ]

TensorRT-LLMを用いてプロンプトの推論結果を出力します。

推論実行前の準備や評価の実行方法については 推論および評価実行方法 を参照してください。

動作検証環境

本ディレクトリの実行環境はDockerの利用を前提としています。

  • NGC Container
    • nvcr.io/nvidia/tritonserver:24.08-trtllm-python-py3
  • TensorRT-LLM 0.12.0.dev2024080600
    • tritonserver:24.08に組み込まれているバージョンと同じコミットをcheckoutして使用
  • Transformers v4.44.2
    • Llama-3.1を使用するために必要
    • pip install時にエラーが表示されるが問題ない

インストール手順

HOMEディレクトリ直下にllm-jp-evalがcloneされている前提で、次の手順を実行します。 （llm-jp-evalのrequirements.txtのインストールは不要です）

$ docker build -t llm-jp-trt-llm ../../ -f Dockerfile
$ docker run -it --rm --ipc=host --gpus all --shm-size 64g \
  -v $(pwd)../../../../:/llm-jp-eval \
  -w /llm-jp-eval/llm-jp-eval-inference/inference-modules/trtllm/ \
  llm-jp-trt-llm

モデル変換・推論・評価の一括実行

build_and_eval.shを使用して簡単にモデルを評価することができます。 引数等の詳細はスクリプトを参照して下さい。 （モデル変換・推論・評価の各ステップを個別に実行する方法については以降の章に詳細な説明があります）

./build_and_eval.sh llm-jp/llm-jp-3-3.7b-instruct bfloat16 origin 1 1 2048 enable A100-SXM-40GB

モデル変換

(以下のモデル変換処理はbuild_and_eval.sh)に組み込まれています)

Llama2ベースのモデルの変換例を以降に示します。大まかな流れは、

  • Hugging Face Transformers形式のモデルをローカルにダウンロード
  • Checkpoint形式に変換
  • Engine形式に変換

となります。

詳細はTensorRT-LLM/examples/llama/を参照してください。

モデルダウンロード

apt install git-lfs
git lfs install

MODEL_ORG=llm-jp
MODEL_DIR=llm-jp-3-3.7b-instruct
git clone https://huggingface.co/${MODEL_ORG}/${MODEL_DIR}
# TensorRT-LLMは古いLLaMA表記のモデル形式に対応していないため config.json の書き換えを行う
cd ${MODEL_DIR} && cp config.json config.json.old && sed "s/LLaMAForCausalLM/LlamaForCausalLM/g" config.json.old > config.json && cd ..

Checkpoint形式への変換

Hugging Face Hubで配布されるLLMのモデルの保存精度は、モデルフォルダのconfig.jsonのtorch_dtypeフィールドで確認できます。 以下のサンプルのDTYPEには、config.jsonのtorch_dtypeフィールドの値をセットしてください。

量子化を行わない場合

DTYPE=bfloat16 # or float16, float32
TP=1
PP=1
MAX_SEQ_LEN=2048
SETTINGS=${DTYPE}_tp${TP}_pp${PP}_$((MAX_SEQ_LEN / 1024))k
CKPT_DIR=./${MODEL_DIR}.${SETTINGS}.ckpt
ENGINE_DIR=./${MODEL_DIR}.${SETTINGS}
python3 scripts_tensorrt/convert_checkpoint.py \
  --model_dir ${MODEL_DIR} --dtype ${DTYPE} \
  --tp_size ${TP} --pp_size ${PP} \
  --output_dir ${CKPT_DIR}

SmoothQuant による量子化

DTYPE=bfloat16 # or float16
TP=1
PP=1
MAX_SEQ_LEN=2048
SETTINGS=sq-0.5_tp${TP}_pp${PP}_$((MAX_SEQ_LEN / 1024))k
CKPT_DIR=./${MODEL_DIR}.${SETTINGS}.ckpt
ENGINE_DIR=./${MODEL_DIR}.${SETTINGS}
python3 TensorRT-LLM/examples/llama/convert_checkpoint.py \
  --model_dir ${MODEL_DIR} --dtype ${DTYPE} \
  --tp_size ${TP} --pp_size ${PP} \
  --smoothquant 0.5 --per_token --per_channel \
  --output_dir ${CKPT_DIR}

FP8 による量子化

FP8を使用するにはAda Lovelace/Hopper以降（sm89以上）のGPUが必要です。

DTYPE=bfloat16 # or float16
TP=1
PP=1
MAX_SEQ_LEN=2048
SETTINGS=fp8_tp${TP}_pp${PP}_$((MAX_SEQ_LEN / 1024))k
CKPT_DIR=./${MODEL_DIR}.${SETTINGS}.ckpt
ENGINE_DIR=./${MODEL_DIR}.${SETTINGS}
python3 TensorRT-LLM/examples/quantization/quantize.py \
  --model_dir ${MODEL_DIR} --dtype ${DTYPE} \
  --tp_size ${TP} --pp_size ${PP} \
  --qformat fp8 --kv_cache_dtype fp8 --calib_size 512 \
  --output_dir ${CKPT_DIR}

INT8 KV cache + AWQ による4bit量子化

DTYPE=bfloat16 # or float16
TP=1
PP=1
MAX_SEQ_LEN=2048
SETTINGS=awq-int4_tp${TP}_tp${TP}_pp${PP}_$((MAX_SEQ_LEN / 1024))k
CKPT_DIR=./${MODEL_DIR}.${SETTINGS}.ckpt
ENGINE_DIR=./${MODEL_DIR}.${SETTINGS}
python3 TensorRT-LLM/examples/quantization/quantize.py \
  --model_dir ${MODEL_DIR} --dtype ${DTYPE} \
  --tp_size ${TP} --pp_size ${PP} \
  --qformat int4_awq --awq_block_size 128 --kv_cache_dtype int8 --calib_size 32 \
  --output_dir ${CKPT_DIR}

Engine形式への変換

Checkpoint形式からEngine形式への変換はtrtllm-buildコマンドで使用して次のように実行します。 （ただしint4_awqで量子化した場合は--use_paged_context_fmha enableの指定を削除する必要があります）

trtllm-build --checkpoint_dir ${CKPT_DIR} --output_dir ${ENGINE_DIR} --max_seq_len ${MAX_SEQ_LEN} --gemm_plugin auto --use_paged_context_fmha enable
rm -r ${CKPT_DIR}

configの編集

config_offline_inference_trtllm.yamlを編集して使用してください。

推論の実行

config_offline_inference_trtllm.yamlの設定にすべて従う場合はこちら。

# for TP=1 PP=1
CUDA_VISIBLE_DEVICES=0 python3 offline_inference_trtllm.py

# use mpirun for TP*PP>1
mpirun -n 2 --allow-run-as-root -x CUDA_VISIBLE_DEVICES=0,1 python3 offline_inference_trtllm.py

各種設定をコマンドラインでオーバーライドする場合はこちら。 （推論結果の保存先ディレクトリは $OUTPUT_BASE_DIR/$RUN_NAME/ となる）

$ CUDA_VISIBLE_DEVICES=0 uv run inference.py \
  --config ./config_sample.yaml \
  --runner.endinge_dir=llm-jp-3-13b.bfloat16_tp1_2k \
  --runner.tokenizer_dir=llm-jp-3-13b \
  --run_name="$RUN_NAME" \
  --output_base_dir="$OUTPUT_BASE_DIR" \
  --prompt_json_path=['../../../local_files/datasets/2.0.0/evaluation/path/to/prompts/*.eval-prompt.json']

主要パラメータを指定して簡易に推論を実行する eval.sh の使用例。

  • 第1引数: ENGINE_DIR
  • 第2引数: TOKENIZER_DIR
  • 第3引数: TP
  • 第4引数: PP
  • 第5引数: MAX_SEQ_LEN
  • 第6引数: PREFIX_CACHING (enable or disable)
  • 第7引数以降: OPTIONS
$ CUDA_VISIBLE_DEVICES=0,1 ./eval.sh llm-jp-3-13b.bfloat16_tp1_2k llm-jp-3-13b 2 1 2048 enable run_name="13b-trtllm-test1"