mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-06 11:11:36 +08:00
502 lines
21 KiB
Python
502 lines
21 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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.
|
|
import argparse
|
|
import json
|
|
import os
|
|
from typing import List, Tuple
|
|
|
|
import tensorrt as trt
|
|
import torch
|
|
from transformers import AutoConfig, AutoTokenizer
|
|
|
|
import tensorrt_llm
|
|
import tensorrt_llm.profiler as profiler
|
|
from tensorrt_llm import logger
|
|
from tensorrt_llm.quantization import QuantMode
|
|
from tensorrt_llm.runtime import (ModelConfig, SamplingConfig, Session,
|
|
TensorInfo)
|
|
|
|
|
|
def get_engine_name(rank):
|
|
return 'rank{}.engine'.format(rank)
|
|
|
|
|
|
def trt_dtype_to_torch(dtype):
|
|
if dtype == trt.float16:
|
|
return torch.float16
|
|
elif dtype == trt.float32:
|
|
return torch.float32
|
|
elif dtype == trt.int32:
|
|
return torch.int32
|
|
else:
|
|
raise TypeError("%s is not supported" % dtype)
|
|
|
|
|
|
class QWenInfer(object):
|
|
|
|
def __init__(self, tokenizer_dir, qwen_engine_dir, log_level, output_csv,
|
|
output_npy, num_beams):
|
|
self.tokenizer_dir = tokenizer_dir
|
|
self.qwen_engine_dir = qwen_engine_dir
|
|
self.log_level = log_level
|
|
self.global_max_input_len = 2048
|
|
self.decoder = None
|
|
self.tokenizer = None
|
|
self.config = None
|
|
self.sampling_config = None
|
|
self.output_csv = output_csv
|
|
self.output_npy = output_npy
|
|
self.num_beams = num_beams
|
|
self.model_config = None
|
|
|
|
def get_model(self):
|
|
# --load the tokenizer and engine #
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
self.tokenizer_dir,
|
|
legacy=False,
|
|
trust_remote_code=True,
|
|
)
|
|
config_path = os.path.join(self.qwen_engine_dir, 'config.json')
|
|
with open(config_path, 'r') as f:
|
|
config = json.load(f)
|
|
gen_config_path = os.path.join(self.tokenizer_dir,
|
|
'generation_config.json')
|
|
with open(gen_config_path, 'r') as f:
|
|
gen_config = json.load(f)
|
|
top_k = gen_config['top_k']
|
|
top_p = gen_config['top_p']
|
|
chat_format = gen_config['chat_format']
|
|
if chat_format == "raw":
|
|
eos_token_id = gen_config['eos_token_id']
|
|
pad_token_id = gen_config['pad_token_id']
|
|
elif chat_format == "chatml":
|
|
pad_token_id = eos_token_id = tokenizer.im_end_id
|
|
else:
|
|
raise Exception("unknown chat format ", chat_format)
|
|
|
|
use_gpt_attention_plugin = config['build_config']['plugin_config'][
|
|
'gpt_attention_plugin']
|
|
remove_input_padding = config['build_config']['plugin_config'][
|
|
'remove_input_padding']
|
|
dtype = config['pretrained_config']['dtype']
|
|
tp_size = config['pretrained_config']['mapping']['tp_size']
|
|
pp_size = config['pretrained_config']['mapping']['pp_size']
|
|
world_size = tp_size * pp_size
|
|
assert world_size == tensorrt_llm.mpi_world_size(), \
|
|
f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})'
|
|
num_heads = config['pretrained_config'][
|
|
'num_attention_heads'] // world_size
|
|
max_batch_size = config['build_config']['max_batch_size']
|
|
hidden_size = config['pretrained_config']['hidden_size'] // world_size
|
|
vocab_size = config['pretrained_config']['vocab_size']
|
|
num_layers = config['pretrained_config']['num_hidden_layers']
|
|
num_kv_heads = config['pretrained_config'].get('num_key_value_heads',
|
|
num_heads)
|
|
paged_kv_cache = config['build_config']['plugin_config'][
|
|
'paged_kv_cache']
|
|
tokens_per_block = config['build_config']['plugin_config'][
|
|
'tokens_per_block']
|
|
max_prompt_embedding_table_size = config['build_config'].get(
|
|
'max_prompt_embedding_table_size', 0)
|
|
quant_mode = QuantMode.from_quant_algo(
|
|
config['pretrained_config']['quantization']['quant_algo'],
|
|
config['pretrained_config']['quantization']['kv_cache_quant_algo'])
|
|
if config['pretrained_config'].get('multi_query_mode', False):
|
|
tensorrt_llm.logger.warning(
|
|
"`multi_query_mode` config is deprecated. Please rebuild the engine."
|
|
)
|
|
num_kv_heads = 1
|
|
# num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size
|
|
use_custom_all_reduce = config['build_config']['plugin_config'].get(
|
|
'use_custom_all_reduce', False)
|
|
|
|
runtime_rank = tensorrt_llm.mpi_rank()
|
|
runtime_mapping = tensorrt_llm.Mapping(world_size=world_size,
|
|
rank=runtime_rank,
|
|
tp_size=tp_size,
|
|
pp_size=pp_size)
|
|
torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
|
|
|
|
model_config = ModelConfig(
|
|
max_batch_size=max_batch_size,
|
|
num_heads=num_heads,
|
|
num_kv_heads=num_kv_heads,
|
|
hidden_size=hidden_size,
|
|
vocab_size=vocab_size,
|
|
num_layers=num_layers,
|
|
gpt_attention_plugin=use_gpt_attention_plugin,
|
|
paged_kv_cache=paged_kv_cache,
|
|
tokens_per_block=tokens_per_block,
|
|
remove_input_padding=remove_input_padding,
|
|
dtype=dtype,
|
|
quant_mode=quant_mode,
|
|
use_custom_all_reduce=use_custom_all_reduce,
|
|
max_prompt_embedding_table_size=max_prompt_embedding_table_size,
|
|
max_beam_width=self.num_beams)
|
|
sampling_config = SamplingConfig(
|
|
end_id=eos_token_id,
|
|
pad_id=pad_token_id,
|
|
num_beams=self.num_beams,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
temperature=1.0,
|
|
)
|
|
|
|
engine_name = get_engine_name(runtime_rank)
|
|
serialize_path = os.path.join(self.qwen_engine_dir, engine_name)
|
|
print(f'Loading engine from {serialize_path}')
|
|
return (model_config, sampling_config, runtime_mapping, runtime_rank,
|
|
serialize_path, tokenizer, eos_token_id, pad_token_id)
|
|
|
|
def qwen_model_init(self):
|
|
(model_config, sampling_config, runtime_mapping, runtime_rank,
|
|
serialize_path, tokenizer, eos_token_id,
|
|
pad_token_id) = self.get_model()
|
|
with open(serialize_path, 'rb') as f:
|
|
engine_buffer = f.read()
|
|
self.decoder = tensorrt_llm.runtime.GenerationSession(
|
|
model_config,
|
|
engine_buffer,
|
|
runtime_mapping,
|
|
)
|
|
self.tokenizer = tokenizer
|
|
self.sampling_config = sampling_config
|
|
self.model_config = model_config
|
|
self.config, _ = AutoConfig.from_pretrained(
|
|
self.tokenizer_dir,
|
|
return_unused_kwargs=True,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
def ptuning_setup(self, prompt_table, dtype, hidden_size, tasks, input_ids):
|
|
if prompt_table is not None:
|
|
task_vocab_size = torch.tensor([prompt_table.shape[1]],
|
|
dtype=torch.int32,
|
|
device="cuda")
|
|
prompt_table = prompt_table.view(
|
|
(prompt_table.shape[0] * prompt_table.shape[1],
|
|
prompt_table.shape[2]))
|
|
prompt_table = prompt_table.cuda().to(
|
|
dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
|
|
else:
|
|
prompt_table = torch.empty([1, hidden_size]).cuda()
|
|
task_vocab_size = torch.zeros([1]).cuda()
|
|
|
|
if tasks is not None:
|
|
tasks = torch.tensor([int(t) for t in tasks.split(',')],
|
|
dtype=torch.int32,
|
|
device="cuda")
|
|
assert tasks.shape[0] == input_ids.shape[
|
|
0], "Number of supplied tasks must match input batch size"
|
|
else:
|
|
tasks = torch.zeros([input_ids.size(0)], dtype=torch.int32).cuda()
|
|
|
|
return [prompt_table, tasks, task_vocab_size]
|
|
|
|
def make_context(
|
|
self,
|
|
query: str,
|
|
history: List[Tuple[str, str]] = None,
|
|
system: str = "You are a helpful assistant.",
|
|
max_window_size: int = 6144,
|
|
):
|
|
if history is None:
|
|
history = []
|
|
|
|
im_start, im_end = "<|im_start|>", "<|im_end|>"
|
|
im_start_tokens = [self.tokenizer.im_start_id] # 151644
|
|
im_end_tokens = [self.tokenizer.im_end_id] # [151645]
|
|
nl_tokens = self.tokenizer.encode("\n")
|
|
|
|
def _tokenize_str(role, content):
|
|
return f"{role}\n{content}", self.tokenizer.encode(
|
|
role, allowed_special=set(self.tokenizer.IMAGE_ST)
|
|
) + nl_tokens + self.tokenizer.encode(
|
|
content, allowed_special=set(self.tokenizer.IMAGE_ST))
|
|
|
|
system_text, system_tokens_part = _tokenize_str("system", system)
|
|
system_tokens = im_start_tokens + system_tokens_part + im_end_tokens
|
|
|
|
raw_text = ""
|
|
context_tokens = []
|
|
|
|
for turn_query, turn_response in reversed(history):
|
|
query_text, query_tokens_part = _tokenize_str("user", turn_query)
|
|
query_tokens = im_start_tokens + query_tokens_part + im_end_tokens
|
|
if turn_response is not None:
|
|
response_text, response_tokens_part = _tokenize_str(
|
|
"assistant", turn_response)
|
|
response_tokens = im_start_tokens + response_tokens_part + im_end_tokens
|
|
|
|
next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens
|
|
prev_chat = (
|
|
f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}"
|
|
)
|
|
else:
|
|
next_context_tokens = nl_tokens + query_tokens + nl_tokens
|
|
prev_chat = f"\n{im_start}{query_text}{im_end}\n"
|
|
|
|
current_context_size = (len(system_tokens) +
|
|
len(next_context_tokens) +
|
|
len(context_tokens))
|
|
if current_context_size < max_window_size:
|
|
context_tokens = next_context_tokens + context_tokens
|
|
raw_text = prev_chat + raw_text
|
|
else:
|
|
break
|
|
|
|
context_tokens = system_tokens + context_tokens
|
|
raw_text = f"{im_start}{system_text}{im_end}" + raw_text
|
|
context_tokens += (nl_tokens + im_start_tokens +
|
|
_tokenize_str("user", query)[1] + im_end_tokens +
|
|
nl_tokens + im_start_tokens +
|
|
self.tokenizer.encode("assistant") + nl_tokens)
|
|
raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n"
|
|
|
|
return raw_text, context_tokens
|
|
|
|
def generate_for_qwenvl(
|
|
self,
|
|
input_tokens,
|
|
max_new_tokens: int,
|
|
prompt_table=None,
|
|
tasks=None,
|
|
task_vocab_size=None,
|
|
num_beams=1,
|
|
):
|
|
input_ids = None
|
|
input_lengths = None
|
|
input_ids = torch.as_tensor(input_tokens,
|
|
device="cuda",
|
|
dtype=torch.int32)
|
|
input_lengths = torch.tensor([input_ids.size(1)],
|
|
device="cuda",
|
|
dtype=torch.int32)
|
|
max_input_length = torch.max(input_lengths).item()
|
|
max_new_tokens = min(max_new_tokens,
|
|
self.global_max_input_len - max_input_length)
|
|
self.decoder.setup(
|
|
batch_size=input_lengths.size(0),
|
|
max_context_length=max_input_length,
|
|
max_new_tokens=max_new_tokens,
|
|
beam_width=num_beams,
|
|
)
|
|
profiler.start("QWen")
|
|
run_time = 1
|
|
for _ in range(run_time):
|
|
output_ids = self.decoder.decode(input_ids, input_lengths,
|
|
self.sampling_config, prompt_table,
|
|
tasks, task_vocab_size)
|
|
torch.cuda.synchronize()
|
|
profiler.stop("QWen")
|
|
Qwen_time = profiler.elapsed_time_in_sec("QWen") / run_time
|
|
|
|
return output_ids, Qwen_time
|
|
|
|
def qwen_infer(self,
|
|
input_vit,
|
|
images_path,
|
|
input_text,
|
|
max_new_tokens,
|
|
num_beams=1,
|
|
history=None):
|
|
if images_path is None:
|
|
content_list = []
|
|
else:
|
|
content_list = images_path
|
|
if history is None:
|
|
history = []
|
|
content_list.append({'text': input_text})
|
|
query = self.tokenizer.from_list_format(content_list)
|
|
raw_text, context_tokens = self.make_context(query, history=history)
|
|
# context_tokens = self.tokenizer.encode(query)
|
|
input_ids = torch.tensor([context_tokens]).to('cuda')
|
|
bos_pos = torch.where(input_ids == self.config.visual['image_start_id'])
|
|
eos_pos = torch.where(
|
|
input_ids == self.config.visual['image_start_id'] + 1)
|
|
assert (bos_pos[0] == eos_pos[0]).all()
|
|
img_pos = torch.stack((bos_pos[0], bos_pos[1], eos_pos[1]), dim=1)
|
|
vocab_size = self.config.vocab_size
|
|
fake_prompt_id = torch.arange(vocab_size,
|
|
vocab_size +
|
|
input_vit.shape[0] * input_vit.shape[1],
|
|
device='cuda')
|
|
fake_prompt_id = fake_prompt_id.reshape(input_vit.shape[0],
|
|
input_vit.shape[1])
|
|
for idx, (i, a, b) in enumerate(img_pos):
|
|
input_ids[i][a + 1:b] = fake_prompt_id[idx]
|
|
input_ids = input_ids.contiguous().to(torch.int32).cuda()
|
|
input_lengths = torch.tensor(input_ids.size(1),
|
|
dtype=torch.int32).cuda()
|
|
dtype = self.model_config.dtype
|
|
prompt_table, tasks, task_vocab_size = self.ptuning_setup(
|
|
input_vit, dtype, self.config.hidden_size, None, input_ids)
|
|
|
|
output_ids, Qwen_time = self.generate_for_qwenvl(
|
|
input_ids, max_new_tokens, prompt_table, tasks, task_vocab_size,
|
|
num_beams)
|
|
|
|
runtime_rank = tensorrt_llm.mpi_rank()
|
|
input_lengths = torch.tensor([input_ids.size(1)],
|
|
device="cuda",
|
|
dtype=torch.int32)
|
|
effective_output_token = 0
|
|
if runtime_rank == 0:
|
|
if self.output_csv is None and self.output_npy is None:
|
|
for b in range(input_lengths.size(0)):
|
|
inputs = input_ids[b]
|
|
if content_list is not None:
|
|
print(f'Input: \"{content_list}\"')
|
|
print("\n")
|
|
if self.num_beams <= 1:
|
|
outputs = output_ids[b][0, len(inputs):].tolist()
|
|
try:
|
|
effective_output_token = effective_output_token + \
|
|
outputs.index(151643)
|
|
except:
|
|
effective_output_token = 1
|
|
output_text = self.tokenizer.decode(
|
|
outputs, skip_special_tokens=True)
|
|
print(f'Output: \"{output_text}\"')
|
|
print("\n")
|
|
else:
|
|
for beam in range(self.num_beams):
|
|
outputs = output_ids[b][beam, len(inputs):].tolist()
|
|
output_text = self.tokenizer.decode(
|
|
outputs, skip_special_tokens=True)
|
|
print(f'Output(beam: {beam}): \"{output_text}\"')
|
|
|
|
logger.info(f'TensorRT-LLM QWen time: {Qwen_time} sec ')
|
|
history.append((query, output_text))
|
|
return output_text
|
|
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--max_new_tokens', type=int, default=200)
|
|
parser.add_argument('--log_level', type=str, default='info')
|
|
parser.add_argument(
|
|
'--vit_engine_dir',
|
|
type=str,
|
|
default='qwen_outputs',
|
|
)
|
|
parser.add_argument(
|
|
'--qwen_engine_dir',
|
|
type=str,
|
|
default='qwen_outputs',
|
|
)
|
|
parser.add_argument('--tokenizer_dir',
|
|
type=str,
|
|
default=".",
|
|
help="Directory containing the tokenizer.model.")
|
|
parser.add_argument('--input_text',
|
|
type=str,
|
|
default="Describe the picture")
|
|
parser.add_argument('--images_path',
|
|
nargs='+',
|
|
type=json.loads,
|
|
default=[{
|
|
'image': './pics/demo.jpeg'
|
|
}])
|
|
parser.add_argument('--input_dir',
|
|
nargs='+',
|
|
type=json.loads,
|
|
default=[{
|
|
'image': 'image.pt'
|
|
}])
|
|
|
|
parser.add_argument(
|
|
'--input_tokens',
|
|
dest='input_file',
|
|
type=str,
|
|
help=
|
|
'CSV or Numpy file containing tokenized input. Alternative to text input.',
|
|
default=None)
|
|
parser.add_argument('--output_csv',
|
|
type=str,
|
|
help='CSV file where the tokenized output is stored.',
|
|
default=None)
|
|
parser.add_argument('--output_npy',
|
|
type=str,
|
|
help='Numpy file where the tokenized output is stored.',
|
|
default=None)
|
|
parser.add_argument('--num_beams',
|
|
type=int,
|
|
help="Use beam search if num_beams >1",
|
|
default=1)
|
|
parser.add_argument("--display", default=False, action='store_true')
|
|
parser.add_argument('--port', type=str, default='8006')
|
|
parser.add_argument("--local_machine", default=False, action='store_true')
|
|
return parser.parse_args()
|
|
|
|
|
|
def vit_process(image_path, engine_dir, stream):
|
|
vit_path = os.path.join(engine_dir,
|
|
'visual_encoder/visual_encoder_fp16.plan')
|
|
logger.info(f'Loading engine from {vit_path}')
|
|
with open(vit_path, 'rb') as f:
|
|
engine_buffer = f.read()
|
|
logger.info(f'Creating session from engine {vit_path}')
|
|
session_vit = Session.from_serialized_engine(engine_buffer)
|
|
device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
|
|
images_list = []
|
|
for img in image_path:
|
|
for v in img.values():
|
|
image = torch.load(v)
|
|
if image.device.type == 'cpu':
|
|
image = image.to(device)
|
|
images_list.append(image)
|
|
images = torch.cat(images_list)
|
|
batch_size = images.size(0)
|
|
images = images.expand(batch_size, -1, -1, -1).contiguous()
|
|
visual_inputs = {'input': images.float()}
|
|
visual_output_info = session_vit.infer_shapes(
|
|
[TensorInfo('input', trt.DataType.FLOAT, images.shape)])
|
|
visual_outputs = {
|
|
t.name: torch.empty(tuple(t.shape),
|
|
dtype=trt_dtype_to_torch(t.dtype),
|
|
device='cuda')
|
|
for t in visual_output_info
|
|
}
|
|
profiler.start("ViT")
|
|
|
|
run_time = 1
|
|
for _ in range(run_time):
|
|
ok = session_vit.run(visual_inputs, visual_outputs, stream)
|
|
profiler.stop("ViT")
|
|
Vit_time = profiler.elapsed_time_in_sec("ViT") / run_time
|
|
logger.info(f'TensorRT-LLM ViT latency: {Vit_time} sec ')
|
|
|
|
assert ok, "Runtime execution failed for vit session"
|
|
|
|
image_embeds = visual_outputs['output']
|
|
return image_embeds
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = parse_arguments()
|
|
stream = torch.cuda.current_stream().cuda_stream
|
|
tensorrt_llm.logger.set_level(args.log_level)
|
|
image_embeds = vit_process(args.input_dir, args.vit_engine_dir, stream)
|
|
qinfer = QWenInfer(args.tokenizer_dir, args.qwen_engine_dir, args.log_level,
|
|
args.output_csv, args.output_npy, args.num_beams)
|
|
qinfer.qwen_model_init()
|
|
qinfer.qwen_infer(image_embeds,
|
|
args.images_path,
|
|
args.input_text,
|
|
args.max_new_tokens,
|
|
args.num_beams,
|
|
history=[])
|