mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
476d7003f8
* update waive list Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> * update waives.txt Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> --------- Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com>
117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2022-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.
|
|
"""Module test_qwen test qwenvl examples."""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
import pytest
|
|
from defs.common import venv_check_call, venv_check_output
|
|
from defs.trt_test_alternative import check_call
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def qwenvl_example_root(llm_root, llm_venv):
|
|
"Get qwenvl example root"
|
|
example_root = os.path.join(llm_root, "examples", "models", "core",
|
|
"qwenvl")
|
|
llm_venv.run_cmd([
|
|
"-m", "pip", "install", "-r",
|
|
os.path.join(example_root, "requirements.txt")
|
|
])
|
|
|
|
return example_root
|
|
|
|
|
|
@pytest.mark.parametrize("llm_qwen_model_root", ["qwen-vl-chat"], indirect=True)
|
|
def test_llm_qwenvl_single_gpu_summary(qwenvl_example_root, llm_qwen_model_root,
|
|
llm_venv, engine_dir):
|
|
"Build & run qwenvl on 1 gpu."
|
|
workspace = llm_venv.get_working_directory()
|
|
|
|
print("Generate vit onnx file and engine...")
|
|
plan_file = f"{workspace}/plan/visual_encoder/visual_encoder_fp16.plan"
|
|
onnx_file = f"{workspace}/visual_encoder/visual_encoder.onnx"
|
|
image = f"{qwenvl_example_root}/pics/demo.jpeg"
|
|
vit_cmd = [
|
|
f"{qwenvl_example_root}/vit_onnx_trt.py",
|
|
f"--pretrained_model_path={llm_qwen_model_root}",
|
|
f"--planFile={plan_file}",
|
|
f"--onnxFile={onnx_file}",
|
|
f"--image_url={image}",
|
|
]
|
|
|
|
venv_check_call(llm_venv, vit_cmd)
|
|
|
|
print("Quantize weight...")
|
|
convert_cmd = [
|
|
f"{qwenvl_example_root}/../qwen/convert_checkpoint.py",
|
|
f"--model_dir={llm_qwen_model_root}",
|
|
f"--output_dir={workspace}/Qwen-VL-Chat",
|
|
f"--dtype=float16",
|
|
]
|
|
|
|
venv_check_call(llm_venv, convert_cmd)
|
|
|
|
print("Build TRT-LLM engines...")
|
|
build_cmd = [
|
|
"trtllm-build",
|
|
f"--checkpoint_dir={workspace}/Qwen-VL-Chat",
|
|
f"--gemm_plugin=float16",
|
|
f"--gpt_attention_plugin=float16",
|
|
f"--max_input_len=2048",
|
|
f"--max_seq_len=3072",
|
|
f"--max_prompt_embedding_table_size=2048",
|
|
f"--remove_input_padding=enable",
|
|
f"--max_beam_width=4",
|
|
f"--output_dir={engine_dir}",
|
|
f"--max_batch_size={8}",
|
|
]
|
|
|
|
check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env)
|
|
|
|
print("Run summary...")
|
|
image_path = json.dumps({"image": image})
|
|
|
|
run_cmd = [
|
|
f"{qwenvl_example_root}/run.py",
|
|
f"--tokenizer_dir={llm_qwen_model_root}",
|
|
f"--qwen_engine_dir={engine_dir}",
|
|
f"--vit_engine_path={workspace}/plan/visual_encoder/visual_encoder_fp16.plan",
|
|
f"--images_path={image_path}",
|
|
"--num_beams=4",
|
|
]
|
|
|
|
output = venv_check_output(llm_venv, run_cmd)
|
|
output = [line for line in output.split("\n") if "Output(beam:" in line]
|
|
print(output)
|
|
|
|
print("Verify the output...")
|
|
results = []
|
|
for item in output:
|
|
match = re.search(r"Output\(beam: \d+\): \"(.*)", item)
|
|
if match:
|
|
results.append(match.group(1))
|
|
|
|
for item in results:
|
|
# check the output if it contains key words
|
|
if ("dog" in item or "labrador" in item) and (
|
|
"sea" in item or "beach" in item) and ("woman" in item
|
|
or "girl" in item):
|
|
pass
|
|
else:
|
|
assert False, f"output is: {item}"
|