mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
587d063e6d
* Update TensorRT-LLM --------- Co-authored-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
376 lines
15 KiB
Python
376 lines
15 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 csv
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import LlamaTokenizerFast
|
|
|
|
import tensorrt_llm
|
|
from tensorrt_llm.quantization import QuantMode
|
|
from tensorrt_llm.runtime import ModelConfig, SamplingConfig
|
|
|
|
from build import get_engine_name # isort:skip
|
|
|
|
EOS_TOKEN = 2
|
|
PAD_TOKEN = 2
|
|
|
|
|
|
def throttle_generator(generator, stream_interval):
|
|
for i, out in enumerate(generator):
|
|
if not i % stream_interval:
|
|
yield out
|
|
|
|
if i % stream_interval:
|
|
yield out
|
|
|
|
|
|
def read_config(config_path: Path):
|
|
with open(config_path, 'r') as f:
|
|
config = json.load(f)
|
|
use_gpt_attention_plugin = config['plugin_config']['gpt_attention_plugin']
|
|
remove_input_padding = config['plugin_config']['remove_input_padding']
|
|
dtype = config['builder_config']['precision']
|
|
gather_all_token_logits = config['builder_config'][
|
|
'gather_all_token_logits']
|
|
tp_size = config['builder_config']['tensor_parallel']
|
|
pp_size = config['builder_config']['pipeline_parallel']
|
|
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['builder_config']['num_heads']
|
|
hidden_size = config['builder_config']['hidden_size']
|
|
vocab_size = config['builder_config']['vocab_size']
|
|
num_layers = config['builder_config']['num_layers']
|
|
num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads)
|
|
paged_kv_cache = config['plugin_config']['paged_kv_cache']
|
|
tokens_per_block = config['plugin_config']['tokens_per_block']
|
|
quant_mode = QuantMode(config['builder_config']['quant_mode'])
|
|
if config['builder_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
|
|
assert (num_heads % tp_size) == 0
|
|
num_heads = num_heads // tp_size
|
|
hidden_size = hidden_size // tp_size
|
|
use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce',
|
|
False)
|
|
max_prompt_embedding_table_size = config['builder_config'].get(
|
|
'max_prompt_embedding_table_size', 0)
|
|
|
|
model_config = ModelConfig(
|
|
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,
|
|
gather_all_token_logits=gather_all_token_logits,
|
|
use_custom_all_reduce=use_custom_all_reduce,
|
|
max_prompt_embedding_table_size=max_prompt_embedding_table_size)
|
|
|
|
return model_config, tp_size, pp_size, dtype
|
|
|
|
|
|
def parse_input(input_text: str, input_file: str, tokenizer, end_id: int,
|
|
remove_input_padding: bool, input_tokens_limit: Union[int,
|
|
None]):
|
|
input_tokens = []
|
|
if input_file is None:
|
|
input_tokens.append(
|
|
tokenizer.encode(input_text, add_special_tokens=False))
|
|
else:
|
|
if input_file.endswith('.csv'):
|
|
with open(input_file, 'r') as csv_file:
|
|
csv_reader = csv.reader(csv_file, delimiter=',')
|
|
for line in csv_reader:
|
|
input_tokens.append(np.array(line, dtype='int32'))
|
|
elif input_file.endswith('.npy'):
|
|
inputs = np.load(input_file)
|
|
for row in inputs:
|
|
row = row[row != end_id]
|
|
input_tokens.append(row)
|
|
elif input_file.endswith('.txt'):
|
|
with open(input_file, 'r', encoding='utf-8',
|
|
errors='replace') as txt_file:
|
|
input_text = txt_file.read()
|
|
input_tokens.append(
|
|
tokenizer.encode(input_text, add_special_tokens=False))
|
|
else:
|
|
print('Input file format not supported.')
|
|
raise SystemExit
|
|
|
|
# Cap max input tokens
|
|
if input_tokens_limit is not None:
|
|
print(
|
|
f"Maximum input number of tokens found as {max([len(x) for x in input_tokens])};"
|
|
f" will be capped to {input_tokens_limit}")
|
|
input_tokens = [x[-input_tokens_limit:] for x in input_tokens]
|
|
|
|
input_ids = None
|
|
input_lengths = torch.tensor([len(x) for x in input_tokens],
|
|
dtype=torch.int32,
|
|
device='cuda')
|
|
if remove_input_padding:
|
|
input_ids = np.concatenate(input_tokens)
|
|
input_ids = torch.tensor(input_ids, dtype=torch.int32,
|
|
device='cuda').unsqueeze(0)
|
|
else:
|
|
input_ids = torch.nested.to_padded_tensor(
|
|
torch.nested.nested_tensor(input_tokens, dtype=torch.int32),
|
|
end_id).cuda()
|
|
|
|
return input_ids, input_lengths
|
|
|
|
|
|
def ptuning_setup(prompt_table, dtype, hidden_size, tasks, input_ids,
|
|
input_lengths, remove_input_padding):
|
|
if prompt_table is not None:
|
|
prompt_table = torch.from_numpy(np.load(prompt_table))
|
|
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()
|
|
|
|
num_sequences = input_lengths.size(
|
|
0) if remove_input_padding else input_ids.size(0)
|
|
|
|
if tasks is not None:
|
|
tasks = torch.tensor([int(t) for t in tasks.split(',')],
|
|
dtype=torch.int32,
|
|
device="cuda")
|
|
assert tasks.shape[
|
|
0] == num_sequences, "Number of supplied tasks must match input batch size"
|
|
else:
|
|
tasks = torch.zeros([num_sequences]).cuda()
|
|
|
|
return [prompt_table, tasks, task_vocab_size]
|
|
|
|
|
|
def print_output(output_ids, input_lengths, max_output_len, tokenizer,
|
|
output_csv, output_npy, sequence_lengths):
|
|
num_beams = output_ids.size(1)
|
|
if output_csv is None and output_npy is None:
|
|
for b in range(input_lengths.size(0)):
|
|
inputs = output_ids[b][0][:input_lengths[b]].tolist()
|
|
input_text = tokenizer.decode(inputs)
|
|
print(f'Input: \"{input_text}\"')
|
|
for beam in range(num_beams):
|
|
output_begin = input_lengths[b]
|
|
output_length = sequence_lengths[b][beam] - input_lengths[b]
|
|
output_end = input_lengths[b] + output_length
|
|
outputs = output_ids[b][beam][output_begin:output_end].tolist()
|
|
output_text = tokenizer.decode(outputs)
|
|
print(f'Output: \"{output_text}\"')
|
|
|
|
output_ids = output_ids.reshape((-1, output_ids.size(2)))
|
|
|
|
if output_csv is not None:
|
|
output_file = Path(output_csv)
|
|
output_file.parent.mkdir(exist_ok=True, parents=True)
|
|
outputs = output_ids.tolist()
|
|
with open(output_file, 'w') as csv_file:
|
|
writer = csv.writer(csv_file, delimiter=',')
|
|
writer.writerows(outputs)
|
|
|
|
if output_npy is not None:
|
|
output_file = Path(output_npy)
|
|
output_file.parent.mkdir(exist_ok=True, parents=True)
|
|
outputs = np.array(output_ids.cpu().contiguous(), dtype='int32')
|
|
np.save(output_file, outputs)
|
|
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--max_output_len', type=int, required=True)
|
|
parser.add_argument('--max_kv_cache_len',
|
|
type=int,
|
|
default=None,
|
|
help='The max kv cache length. \
|
|
If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \
|
|
If it is set to None, we will use the max sequence length.')
|
|
parser.add_argument('--log_level', type=str, default='error')
|
|
parser.add_argument('--engine_dir', type=str, default='llama_outputs')
|
|
parser.add_argument('--tokenizer_dir',
|
|
type=str,
|
|
default=".",
|
|
help="Directory containing the tokenizer.model.")
|
|
parser.add_argument('--input_text',
|
|
type=str,
|
|
default='Born in north-east France, Soyer trained as a')
|
|
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(
|
|
'--input_tokens_limit',
|
|
type=int,
|
|
help='Truncate input tokens if number exceeds the set limit value',
|
|
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('--streaming', default=False, action='store_true')
|
|
parser.add_argument('--streaming_interval',
|
|
type=int,
|
|
help="How often to return tokens when streaming.",
|
|
default=5)
|
|
parser.add_argument(
|
|
'--prompt_table',
|
|
type=Path,
|
|
help="Path to .npy file, exported by nemo_prompt_convert.py")
|
|
parser.add_argument(
|
|
'--tasks',
|
|
help="Comma-separated list of tasks for prompt tuning: ex 0,3,1,0")
|
|
return parser.parse_args()
|
|
|
|
|
|
def generate(
|
|
max_output_len: int,
|
|
log_level: str = 'error',
|
|
engine_dir: str = 'llama_outputs',
|
|
input_text: str = 'Born in north-east France, Soyer trained as a',
|
|
input_file: str = None,
|
|
output_csv: str = None,
|
|
output_npy: str = None,
|
|
tokenizer_dir: str = None,
|
|
max_kv_cache_len: int = None,
|
|
num_beams: int = 1,
|
|
streaming: bool = False,
|
|
streaming_interval: int = 5,
|
|
prompt_table: Path = None,
|
|
tasks: str = None,
|
|
input_tokens_limit: Union[None, int] = None,
|
|
):
|
|
tensorrt_llm.logger.set_level(log_level)
|
|
|
|
engine_dir = Path(engine_dir)
|
|
config_path = engine_dir / 'config.json'
|
|
model_config, tp_size, pp_size, dtype = read_config(config_path)
|
|
world_size = tp_size * pp_size
|
|
|
|
runtime_rank = tensorrt_llm.mpi_rank()
|
|
runtime_mapping = tensorrt_llm.Mapping(world_size,
|
|
runtime_rank,
|
|
tp_size=tp_size,
|
|
pp_size=pp_size)
|
|
torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
|
|
|
|
tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_dir, legacy=False)
|
|
|
|
sampling_config = SamplingConfig(end_id=EOS_TOKEN,
|
|
pad_id=PAD_TOKEN,
|
|
num_beams=num_beams)
|
|
|
|
engine_name = get_engine_name('llama', dtype, tp_size, pp_size,
|
|
runtime_rank)
|
|
serialize_path = engine_dir / engine_name
|
|
with open(serialize_path, 'rb') as f:
|
|
engine_buffer = f.read()
|
|
decoder = tensorrt_llm.runtime.GenerationSession(model_config,
|
|
engine_buffer,
|
|
runtime_mapping,
|
|
debug_mode=False,
|
|
debug_tensors_to_save=None)
|
|
if runtime_rank == 0:
|
|
print(f"Running the {dtype} engine ...")
|
|
|
|
input_ids, input_lengths = parse_input(
|
|
input_text,
|
|
input_file,
|
|
tokenizer,
|
|
EOS_TOKEN,
|
|
model_config.remove_input_padding,
|
|
input_tokens_limit=input_tokens_limit)
|
|
|
|
max_input_length = torch.max(input_lengths).item()
|
|
decoder.setup(input_lengths.size(0),
|
|
max_input_length,
|
|
max_output_len,
|
|
num_beams,
|
|
max_kv_cache_length=max_kv_cache_len)
|
|
|
|
ptuning_args = [] if model_config.max_prompt_embedding_table_size == 0 else ptuning_setup(
|
|
prompt_table, dtype, model_config.hidden_size, tasks, input_ids,
|
|
input_lengths, model_config.remove_input_padding)
|
|
|
|
outputs = decoder.decode(input_ids,
|
|
input_lengths,
|
|
sampling_config,
|
|
*ptuning_args,
|
|
streaming=streaming,
|
|
output_sequence_lengths=True,
|
|
return_dict=True)
|
|
torch.cuda.synchronize()
|
|
if streaming:
|
|
for outputs_dict in throttle_generator(outputs, streaming_interval):
|
|
if runtime_rank == 0:
|
|
output_ids = outputs_dict['output_ids']
|
|
sequence_lengths = outputs_dict['sequence_lengths']
|
|
print_output(output_ids, input_lengths, max_output_len,
|
|
tokenizer, output_csv, output_npy,
|
|
sequence_lengths)
|
|
else:
|
|
if runtime_rank == 0:
|
|
output_ids = outputs['output_ids']
|
|
sequence_lengths = outputs['sequence_lengths']
|
|
print_output(output_ids, input_lengths, max_output_len, tokenizer,
|
|
output_csv, output_npy, sequence_lengths)
|
|
|
|
if model_config.gather_all_token_logits:
|
|
if runtime_mapping.is_last_pp_rank():
|
|
print(
|
|
f"context_logits.shape: {outputs['context_logits'].shape}")
|
|
print(
|
|
f"generation_logits.shape: {len(outputs['generation_logits']), outputs['generation_logits'][0].shape}"
|
|
)
|
|
print(outputs['context_logits'])
|
|
print(outputs['generation_logits'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = parse_arguments()
|
|
generate(**vars(args))
|