单卡A100通过TensorRT-LLM量化推理Qwen3-8B简单教程

1
sudo docker run --rm -it \
2
    --ipc=host \
3
    --ulimit memlock=-1 \
4
    --ulimit stack=67108864 \
5
    --gpus=all \
6
    -v /home/ros/tensorrt-llm/models:/workspace/models \
7
    -p 8000:8000 \
8
    nvcr.io/nvidia/tensorrt-llm/release:1.0.0rc3

1
sudo apt install git-lfs
2
git clone https://www.modelscope.cn/Qwen/Qwen3-8B.git

1
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
2
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
3
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
4
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
5

6
sudo apt-get update
7

8
export NVIDIA_CONTAINER_TOOLKIT_VERSION=1.17.8-1
9
  sudo apt-get install -y \
10
      nvidia-container-toolkit=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
11
      nvidia-container-toolkit-base=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
12
      libnvidia-container-tools=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
13
      libnvidia-container1=${NVIDIA_CONTAINER_TOOLKIT_VERSION}

1
python /app/tensorrt_llm/examples/models/core/qwen/convert_checkpoint.py \
2
    --output_dir /workspace/models/Qwen3-8B-checkpoint \
3
    --model_dir /workspace/models/Qwen3-8B \
4
    --dtype float16 \
5
    --tp_size 1 \
6
    --pp_size 1 \
7
    --workers 32 \
8
    --use_parallel_embedding

1
{
2
  "tensor_parallel": 1,
3
  "pipeline_parallel": 1,
4
  "max_batch_size": 1,
5
  "max_input_len": 8192,
6
  "max_output_len": 8192,
7
  "max_seq_len": 16384,
8
  "max_num_tokens": 8192,
9
  "max_beam_width": 1,
10
  "opt_num_tokens":8192,
11
  "gemm_plugin": "float16",
12
  "gpt_attention_plugin": "float16",
13
  "remove_input_padding": "enable",
14
  "context_fmha": "enable",
15
  "log_level": "info"
16
}

1
trtllm-build \
2
    --output_dir /workspace/models/Qwen3-8B-engine \
3
    --checkpoint_dir /workspace/models/Qwen3-8B-checkpoint \
4
    --build_config /workspace/models/build_config.json \
5
    --workers 32 \
6
    --cluster_key A100-PCIe-40GB

1
python /app/tensorrt_llm/examples/run.py \
2
    --engine_dir /workspace/models/Qwen3-8B-engine \
3
    --tokenizer_dir /workspace/models/Qwen3-8B \
4
    --input_text "Born in north-east France, Soyer trained as a" \
5
    --max_output_len 128 \
6
    --log_level info

1
trtllm-serve serve /workspace/models/Qwen3-8B-engine --tokenizer /workspace/models/Qwen3-8B --host 0.0.0.0 --port 8000 --tp_size 1 --pp_size 1 --max_batch_size 1 --max_seq_len 2048 --max_beam_width 1 --max_num_tokens 1024 --log_level info --num_postprocess_workers 16

1
from openai import OpenAI
2

3
client = OpenAI(
4
    base_url="http://localhost:8000/v1",
5
    api_key="tensorrt_llm",
6
)
7

8
response = client.chat.completions.create(
9
    model="TinyLlama-1.1B-Chat-v1.0",
10
    messages=[{
11
        "role": "system",
12
        "content": "you are a helpful assistant"
13
    }, {
14
        "role": "user",
15
        "content": "Where is New York?"
16
    }],
17
    max_tokens=20,
18
)
19
print(response)