diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/README.md b/models/nlp/large_language_model/llama2-13b/trtllm/README.md new file mode 100755 index 0000000000000000000000000000000000000000..433730567bfb1f121fd4018530535d9e9e5ad181 --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/README.md @@ -0,0 +1,50 @@ +# Llama2_13b_gpu2 + +## Description +The Llama2 model is part of the Llama project which aims to unlock the power of large language models. The latest version of the Llama model is now accessible to individuals, creators, researchers, and businesses of all sizes. It includes model weights and starting code for pre-trained and fine-tuned Llama language models with parameters ranging from 7B to 70B. + +## Setup + +### Install +In order to run the model smoothly, we need the following dependency files: +1. ixrt-xxx.whl +2. tensorrt_llm-xxx.whl +3. ixformer-xxx.whl +Please contact the staff to obtain the relevant installation packages. + +```bash +yum install mesa-libGL +bash set_environment.sh +pip3 install Path/To/ixrt-xxx.whl +pip3 install Path/To/tensorrt_llm-xxx.whl +pip3 install Path/To/ixformer-xxx.whl +``` + +### Download +-Model: https://huggingface.co/meta-llama/Llama-2-13b-chat-hf + +-Dataset:https://huggingface.co/datasets/cnn_dailymail + +```bash +# Download model from the website and make sure the model's path is "data/llama2-13b-chat" +# Download dataset from the website and make sure the dataset's path is "data/datasets_cnn_dailymail" +mkdir data + +# Please download rouge.py to this path if your server can't attach huggingface.co. +mkdir -p rouge/ +wget --no-check-certificate https://raw.githubusercontent.com/huggingface/evaluate/main/metrics/rouge/rouge.py -P rouge +``` + +## Inference +```bash +export CUDA_VISIBLE_DEVICES=0,1 + +``` +### FP16 + +```bash +# Build Engine +bash scripts/test_trtllm_llama2_13b_gpu2_build.sh +# Inference +bash scripts/test_trtllm_llama2_13b_gpu2.sh +``` diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/build.py b/models/nlp/large_language_model/llama2-13b/trtllm/build.py new file mode 100644 index 0000000000000000000000000000000000000000..4ff0c9eaa0cedfd382783a5cfcca9175bf38acad --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/build.py @@ -0,0 +1,1163 @@ +# 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. +import argparse +import json +import math +import os +import sys +import time +from pathlib import Path + +# isort: off +import torch +import torch.multiprocessing as mp +import tensorrt as trt +# isort: on +from transformers import LlamaConfig, LlamaForCausalLM + +try: + from transformers import MixtralForCausalLM +except ImportError: + MixtralForCausalLM = None + +try: + from transformers import LlavaConfig, LlavaForConditionalGeneration +except ImportError: + pass + +import tensorrt_llm +from tensorrt_llm import profiler +from tensorrt_llm._common import check_max_num_tokens +from tensorrt_llm._utils import str_dtype_to_trt +from tensorrt_llm.builder import Builder +from tensorrt_llm.layers import MoeConfig +from tensorrt_llm.layers.attention import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models import quantize_model +from tensorrt_llm.network import net_guard +from tensorrt_llm.plugin.plugin import ContextFMHAType +from tensorrt_llm.quantization import QuantMode +from tensorrt_llm.runtime.lora_manager import LoraConfig + +from tensorrt_llm.models.llama.weight import ( # isort:skip + get_scaling_factors, load_from_awq_llama, load_from_binary, + load_from_gptq_llama, load_from_hf_checkpoint, load_from_hf_llama, + load_from_meta_llama, parse_bin_config) + +MODEL_NAME = "llama" + +# 2 routines: get_engine_name, serialize_engine +# are direct copy from gpt example, TODO: put in utils? + +import onnx +from onnx import TensorProto, helper + + +def trt_dtype_to_onnx(dtype): + if dtype == trt.float16: + return TensorProto.DataType.FLOAT16 + if dtype == trt.bfloat16: + return TensorProto.DataType.BFLOAT16 + elif dtype == trt.float32: + return TensorProto.DataType.FLOAT + elif dtype == trt.int32: + return TensorProto.DataType.INT32 + elif dtype == trt.int64: + return TensorProto.DataType.INT64 + elif dtype == trt.bool: + return TensorProto.DataType.BOOL + else: + raise TypeError("%s is not supported" % dtype) + + +def to_onnx(network, path): + inputs = [] + for i in range(network.num_inputs): + network_input = network.get_input(i) + inputs.append( + helper.make_tensor_value_info( + network_input.name, trt_dtype_to_onnx(network_input.dtype), + list(network_input.shape))) + + outputs = [] + for i in range(network.num_outputs): + network_output = network.get_output(i) + outputs.append( + helper.make_tensor_value_info( + network_output.name, trt_dtype_to_onnx(network_output.dtype), + list(network_output.shape))) + + nodes = [] + for i in range(network.num_layers): + layer = network.get_layer(i) + layer_inputs = [] + for j in range(layer.num_inputs): + ipt = layer.get_input(j) + if ipt is not None: + layer_inputs.append(layer.get_input(j).name) + layer_outputs = [ + layer.get_output(j).name for j in range(layer.num_outputs) + ] + nodes.append( + helper.make_node(str(layer.type), + name=layer.name, + inputs=layer_inputs, + outputs=layer_outputs, + domain="com.nvidia")) + + onnx_model = helper.make_model(helper.make_graph(nodes, + 'attention', + inputs, + outputs, + initializer=None), + producer_name='NVIDIA') + onnx.save(onnx_model, path) + + +def get_engine_name(model, dtype, tp_size, pp_size, rank): + if pp_size == 1: + return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) + return '{}_{}_tp{}_pp{}_rank{}.engine'.format(model, dtype, tp_size, + pp_size, rank) + + +def serialize_engine(engine, path): + logger.info(f'Serializing engine to {path}...') + tik = time.time() + with open(path, 'wb') as f: + f.write(engine) + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'Engine serialized. Total time: {t}') + + +def parse_arguments(cmd_args=None): + parser = argparse.ArgumentParser() + parser.add_argument('--world_size', type=int, default=1) + parser.add_argument('--tp_size', type=int, default=1) + parser.add_argument('--pp_size', type=int, default=1) + parser.add_argument('--model_dir', type=str, default=None) + parser.add_argument('--bin_model_dir', type=str, default=None) + parser.add_argument('--meta_ckpt_dir', type=str, default=None) + parser.add_argument('--quant_ckpt_path', type=str, default=None) + parser.add_argument('--dtype', + type=str, + default='float16', + choices=['float32', 'bfloat16', 'float16']) + parser.add_argument( + '--timing_cache', + type=str, + default='model.cache', + help= + 'The path of to read timing cache from, will be ignored if the file does not exist' + ) + parser.add_argument( + '--profiling_verbosity', + type=str, + default='layer_names_only', + choices=['layer_names_only', 'detailed', 'none'], + help= + 'The profiling verbosity for the generated TRT engine. Set to detailed can inspect tactic choices and kernel parameters.' + ) + parser.add_argument('--log_level', type=str, default='info') + parser.add_argument('--vocab_size', type=int, default=32000) + parser.add_argument('--n_layer', type=int, default=32) + parser.add_argument('--n_positions', type=int, default=2048) + parser.add_argument('--n_embd', type=int, default=4096) + parser.add_argument('--n_head', type=int, default=32) + parser.add_argument('--n_kv_head', type=int, default=None) + parser.add_argument('--multiple_of', type=int, default=256) + parser.add_argument('--ffn_dim_multiplier', type=float, default=1.0) + parser.add_argument('--inter_size', type=int, default=None) + parser.add_argument('--hidden_act', type=str, default='silu') + parser.add_argument('--rms_norm_eps', type=float, default=1e-06) + parser.add_argument('--max_batch_size', type=int, default=8) + parser.add_argument('--max_input_len', type=int, default=2048) + parser.add_argument('--max_output_len', type=int, default=512) + parser.add_argument('--max_beam_width', type=int, default=1) + parser.add_argument('--rotary_base', type=float, default=10000.0) + parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) + parser.add_argument('--use_gpt_attention_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'bfloat16', 'float32']) + parser.add_argument('--use_gemm_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'bfloat16', 'float32']) + parser.add_argument('--use_rmsnorm_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'float32', 'bfloat16']) + parser.add_argument('--use_lookup_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'bfloat16', 'float32']) + parser.add_argument('--use_gather_last_token_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'float32', 'bfloat16']) + parser.add_argument('--use_activation_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'float32', 'bfloat16']) + parser.add_argument('--use_elementwise_plugin', + nargs='?', + const='float16', + type=str, + default=False, + choices=['float16', 'float32', 'bfloat16']) + parser.add_argument("--use_cast_plugin", action="store_true") + + parser.add_argument('--parallel_build', default=False, action='store_true') + parser.add_argument('--enable_context_fmha', + default=False, + action='store_true') + parser.add_argument('--enable_context_fmha_fp32_acc', + default=False, + action='store_true') + parser.add_argument( + '--use_paged_context_fmha', + action='store_true', + help= + 'Activates paged context FMHA. This mode of the context FMHA is required for chunked context, speculative decoding and reuse of KV cache blocks. Context FMHA performance is worse when this mode is on.' + ) + parser.add_argument( + '--multi_block_mode', + default=False, + action='store_true', + help= + 'Split long kv sequence into multiple blocks (applied to generation MHA kernels). \ + It is beneficial when batch x num_heads cannot fully utilize GPU.' + ) + parser.add_argument( + '--disable_xqa', + default=False, + action='store_true', + help= + 'Disable XQA optimization for the generation MHA. See more details in docs/gpt_attention.' + ) + parser.add_argument('--visualize', default=False, action='store_true') + parser.add_argument('--load_by_shard', + action='store_true', + help='Load a pretrained model shard-by-shard.') + parser.add_argument('--enable_debug_output', + default=False, + action='store_true') + parser.add_argument('--gpus_per_node', type=int, default=8) + parser.add_argument('--builder_opt', type=int, default=None) + parser.add_argument( + '--output_dir', + type=str, + default='engine_outputs', + help= + 'The path to save the serialized engine files, timing cache file and model configs' + ) + parser.add_argument('--remove_input_padding', + default=False, + action='store_true') + parser.add_argument( + '--use_fused_mlp', + default=False, + action='store_true', + help= + 'Enable horizontal fusion in GatedMLP, reduces layer input traffic and potentially improves performance. ' + 'For FP8 PTQ, the downside is slight reduction of accuracy because one of the quantization scaling factors are discarded ' + '(0.45734 vs 0.45755 for LLaMA-v2 7B using ammo/examples/hf/instruct_eval/mmlu.py).' + ) + parser.add_argument('--enable_pos_shift', + default=False, + action='store_true', + help='Enable position shift for streamingllm method') + parser.add_argument( + '--dense_context_fmha', + default=False, + action='store_true', + help= + 'Enable dense fmha in context phase, otherwise sliding window attention.' + 'If dense_context_fmha=False, the sliding window size is the max attention window size.' + ) + # Arguments related to the quantization of the model. + parser.add_argument( + '--use_smooth_quant', + default=False, + action="store_true", + help= + 'Use the SmoothQuant method to quantize activations and weights for the various GEMMs.' + 'See --per_channel and --per_token for finer-grained quantization options.' + ) + parser.add_argument( + '--per_channel', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor for the GEMM\'s result. ' + 'per_channel instead uses a different static scaling factor for each channel. ' + 'The latter is usually more accurate, but a little slower.') + parser.add_argument( + '--per_token', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor to scale activations in the int8 range. ' + 'per_token chooses at run time, and for each token, a custom scaling factor. ' + 'The latter is usually more accurate, but a little slower.') + parser.add_argument( + '--per_group', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor to scale weights in the int4 range. ' + 'per_group chooses at run time, and for each group, a custom scaling factor. ' + 'The flag is built for GPTQ/AWQ quantization.') + parser.add_argument('--group_size', + type=int, + default=128, + help='Group size used in GPTQ/AWQ quantization.') + parser.add_argument( + '--int8_kv_cache', + default=False, + action="store_true", + help= + 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' + ) + parser.add_argument( + '--use_parallel_embedding', + action="store_true", + default=False, + help= + 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' + ) + parser.add_argument( + '--embedding_sharding_dim', + type=int, + default=1, # Meta does TP on hidden dim + choices=[0, 1], + help= + 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' + 'To shard it along hidden dimension, set embedding_sharding_dim=1' + 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' + ) + parser.add_argument( + '--enable_fp8', + default=False, + action='store_true', + help='Use FP8 Linear layer for Attention QKV/Dense and MLP.') + parser.add_argument( + '--fp8_kv_cache', + default=False, + action="store_true", + help= + 'By default, we use dtype for KV cache. fp8_kv_cache chooses int8 quantization for KV' + ) + parser.add_argument( + '--quantized_fp8_model_path', + type=str, + default=None, + help='Path of a quantized model checkpoint in .npz format') + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + parser.add_argument( + '--disable_weight_only_quant_plugin', + default=False, + action="store_true", + help= + 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + parser.add_argument( + '--weight_only_precision', + const='int8', + type=str, + nargs='?', + default='int8', + choices=['int8', 'int4', 'int4_awq', 'int4_gptq'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + parser.add_argument( + '--quantize_lm_head', + default=False, + action="store_true", + help='Quantize lm_head weights as well when using int4_awq.') + parser.add_argument( + '--use_inflight_batching', + action="store_true", + default=False, + help="Activates inflight batching mode of gptAttentionPlugin.") + parser.add_argument( + '--paged_kv_cache', + action="store_true", + default=False, + help= + 'By default we use contiguous KV cache. By setting this flag you enable paged KV cache' + ) + parser.add_argument('--tokens_per_block', + type=int, + default=128, + help='Number of tokens per block in paged KV cache') + parser.add_argument( + '--max_num_tokens', + type=int, + default=None, + help= + 'Define the max number of tokens supported by the engine, note that it takes no effect if --remove_input_padding is not set' + ) + parser.add_argument( + '--strongly_typed', + default=False, + action="store_true", + help= + 'This option is introduced with trt 9.1.0.1+ and will reduce the building time significantly for fp8.' + ) + parser.add_argument( + '--use_custom_all_reduce', + action='store_true', + help= + 'Activates latency-optimized algorithm for all-reduce instead of NCCL.') + parser.add_argument( + '--max_prompt_embedding_table_size', + type=int, + default=0, + help='Setting to a value > 0 enables support for prompt tuning.') + parser.add_argument( + '--gather_all_token_logits', + action='store_true', + default=False, + help='Enable both gather_context_logits and gather_generation_logits') + parser.add_argument('--gather_context_logits', + action='store_true', + default=False, + help='Gather context logits') + parser.add_argument('--gather_generation_logits', + action='store_true', + default=False, + help='Gather generation logits') + parser.add_argument( + '--use_lora_plugin', + nargs='?', + const=None, + default=False, + choices=['float16', 'float32', 'bfloat16'], + help="Activates the lora plugin which enables embedding sharing.") + parser.add_argument( + '--lora_target_modules', + nargs='+', + default=None, + choices=[ + "attn_qkv", + "attn_q", + "attn_k", + "attn_v", + "attn_dense", + "mlp_h_to_4h", + "mlp_gate", + "mlp_4h_to_h", + ], + help= + "Add lora in which modules. Only be activated when use_lora_plugin is enabled." + ) + parser.add_argument('--hf_lora_dir', type=str, default=None) + parser.add_argument( + '--max_lora_rank', + type=int, + default=64, + help='maximum lora rank for different lora modules. ' + 'It is used to compute the workspace size of lora plugin.') + parser.add_argument( + '--moe_num_experts', + default=0, + type=int, + help='Specify the number of experts to use for MOE layers') + parser.add_argument( + '--moe_top_k', + default=0, + type=int, + help= + 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' + ) + parser.add_argument( + '--moe_tp_mode', + default=MoeConfig.ParallelismMode.TENSOR_PARALLEL, + type=int, + help= + 'Controls how to distribute experts in TP. Check layers/moe.py for accepted values', + ) + parser.add_argument( + '--moe_renorm_mode', + default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, + type=int, + help= + 'Controls renormalization after gate logits. Check layers/moe.py for accepted values', + ) + parser.add_argument("--total_build_time_target", type=float, default=0) + + args = parser.parse_args(cmd_args) + logger.set_level(args.log_level) + + assert args.total_build_time_target >= 0, "total_build_time_target must bigger than 0" + + assert not ( + args.use_smooth_quant and args.use_weight_only + ), "You cannot enable both SmoothQuant and INT8 weight-only together." + + if not args.remove_input_padding: + if args.use_gpt_attention_plugin: + logger.warning( + f"It is recommended to specify --remove_input_padding when using GPT attention plugin" + ) + + if args.use_inflight_batching: + if not args.use_gpt_attention_plugin: + args.use_gpt_attention_plugin = 'float16' + logger.info( + f"Using GPT attention plugin for inflight batching mode. Setting to default '{args.use_gpt_attention_plugin}'" + ) + if not args.remove_input_padding: + args.remove_input_padding = True + logger.info( + "Using remove input padding for inflight batching mode.") + if not args.paged_kv_cache: + args.paged_kv_cache = True + logger.info("Using paged KV cache for inflight batching mode.") + + if args.use_smooth_quant: + args.quant_mode = QuantMode.use_smooth_quant(args.per_token, + args.per_channel) + elif args.use_weight_only: + args.quant_mode = QuantMode.from_description( + quantize_weights=True, + quantize_activations=False, + per_token=False, + per_channel=False, + per_group=args.per_group, + use_int4_weights="int4" in args.weight_only_precision) + else: + args.quant_mode = QuantMode(0) + + if args.int8_kv_cache: + args.quant_mode = args.quant_mode.set_int8_kv_cache() + elif args.fp8_kv_cache: + args.quant_mode = args.quant_mode.set_fp8_kv_cache() + if args.enable_fp8: + args.quant_mode = args.quant_mode.set_fp8_qdq() + + if args.rotary_scaling is not None: + assert args.use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." + rotary_scaling = { + "type": args.rotary_scaling[0], + "factor": float(args.rotary_scaling[1]) + } + assert rotary_scaling["type"] in ["linear", "dynamic"] + assert rotary_scaling["factor"] > 1.0 + args.rotary_scaling = rotary_scaling + + if args.model_dir is not None: + hf_config = LlamaConfig.from_pretrained(args.model_dir) + if hf_config.model_type == "llava": + # LLaVA = Vision model + Llama LLM + # We load a llava config and use its' text config as llama config + hf_config = LlavaConfig.from_pretrained(args.model_dir).text_config + hf_config.model_type = "llava" # Replace llama with llava + + args.inter_size = hf_config.intermediate_size # override the inter_size for LLaMA + args.n_embd = hf_config.hidden_size + args.n_head = hf_config.num_attention_heads + if hasattr(hf_config, "num_key_value_heads"): + args.n_kv_head = hf_config.num_key_value_heads + + # hf_config.num_hidden_layers = 1 # only for debug + args.n_layer = hf_config.num_hidden_layers + args.n_positions = hf_config.max_position_embeddings + args.vocab_size = hf_config.vocab_size if hf_config.vocab_size is not None else args.vocab_size + args.hidden_act = hf_config.hidden_act + args.rms_norm_eps = hf_config.rms_norm_eps + # These attributes only exists with Mixtral, for the moment + args.moe_num_experts = getattr(hf_config, "num_local_experts", + args.moe_num_experts) + args.moe_top_k = getattr(hf_config, "num_experts_per_tok", + args.moe_top_k) + args.rotary_base = getattr(hf_config, "rope_theta", args.rotary_base) + args.model_type = hf_config.model_type + if hf_config.model_type == "mixtral": + # HF LLaMA-type models are implicitly using gated activation. + # With our MoE implementation, we must make it explicit + args.hidden_act = "swiglu" + + elif args.meta_ckpt_dir is not None: + with open(Path(args.meta_ckpt_dir, "params.json")) as fp: + meta_config: dict = json.load(fp) + args.n_embd = meta_config["dim"] + args.n_head = meta_config["n_heads"] + args.n_layer = meta_config["n_layers"] + args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) + if "hidden_dim" in meta_config: + args.inter_size = meta_config["hidden_dim"] + else: + args.multiple_of = meta_config.get("multiple_of", 1) + n_embd = int(4 * args.n_embd * 2 / 3) + args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) + args.inter_size = args.multiple_of * ( + (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) + // args.multiple_of) + args.rms_norm_eps = meta_config["norm_eps"] + args.moe_num_experts = meta_config.get("moe", {}).get("num_experts", 0) + args.moe_top_k = meta_config.get("moe", {}).get("num_experts_per_tok", + 0) + elif args.bin_model_dir is not None: + n_embd, n_head, n_layer, n_positions, vocab_size, hidden_act, inter_size, n_kv_head = parse_bin_config( + Path(args.bin_model_dir) / "config.ini") + args.inter_size = inter_size # override the inter_size for LLaMA + args.n_kv_head = n_kv_head + args.n_embd = n_embd + args.n_head = n_head + args.n_layer = n_layer + args.n_positions = n_positions + args.vocab_size = vocab_size if args.vocab_size is None else args.vocab_size + args.hidden_act = hidden_act + args.rms_norm_eps = 1e-06 + logger.warning("Set rms_norm_eps to 1e-06 directly.") + if args.n_kv_head is None: + args.n_kv_head = args.n_head + elif args.n_kv_head != args.n_head: + assert (args.n_head % args.n_kv_head) == 0, \ + "MQA/GQA requires the number of heads to be divisible by the number of K/V heads." + assert (args.n_kv_head % args.tp_size) == 0 or (args.tp_size % args.n_kv_head) == 0, \ + "MQA/GQA requires either the number of K/V heads to be divisible by the tensor parallelism size OR " \ + "the tensor parallelism size to be divisible by the number of K/V heads." + + hf_modules_to_trtllm_modules = { + "q_proj": "attn_q", + "k_proj": "attn_k", + "v_proj": "attn_v", + "o_proj": "attn_dense", + "gate_proj": "mlp_h_to_4h", + "down_proj": "mlp_4h_to_h", + "up_proj": "mlp_gate" + } # lora modules on llama + + trtllm_modules_to_hf_modules = { + "attn_q": "q_proj", + "attn_k": "k_proj", + "attn_v": "v_proj", + "attn_dense": "o_proj", + "mlp_h_to_4h": "gate_proj", + "mlp_4h_to_h": "down_proj", + "mlp_gate": "up_proj", + } + + lora_config = LoraConfig.from_hf(args.hf_lora_dir, + hf_modules_to_trtllm_modules, + trtllm_modules_to_hf_modules) + + if lora_config.is_valid: + if args.lora_target_modules is None: + args.lora_target_modules = lora_config.lora_target_modules + # the lora checkpoint might finetune the embedding + if lora_config.vocab_size != 0: + args.vocab_size = lora_config.vocab_size + + args.lora_config = lora_config + + if args.weight_only_precision == 'int4_awq': + inter_alignment = args.tp_size * 128 + if args.inter_size % inter_alignment != 0: + args.inter_size = int((args.inter_size + inter_alignment - 1) / + inter_alignment) * inter_alignment + logger.info("To use awq we pad intermediate_size to {}.".format( + args.inter_size)) + + if args.quantize_lm_head: + vocab_alignment = args.tp_size * 64 + if args.vocab_size % vocab_alignment != 0: + args.vocab_size = int((args.vocab_size + vocab_alignment - 1) / + vocab_alignment) * vocab_alignment + logger.info("To use awq we pad vocab_size to {}.".format( + args.vocab_size)) + + assert args.pp_size * args.tp_size == args.world_size + + args.max_num_tokens = check_max_num_tokens( + max_num_tokens=args.max_num_tokens, + max_batch_size=args.max_batch_size, + max_input_len=args.max_input_len, + remove_input_padding=args.remove_input_padding) + + assert (math.log2(args.tokens_per_block).is_integer() + ), "tokens_per_block must be power of 2" + if args.enable_context_fmha or args.enable_context_fmha_fp32_acc: + assert (args.tokens_per_block >= + 128), "Context fMHA requires >= 128 tokens per block" + + if args.inter_size is None: + # this should not be need when loading a real model + # but it is helpful when creating a dummy model without loading any real weights + n_embd = int(4 * args.n_embd * 2 / 3) + args.inter_size = args.multiple_of * ( + (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) // + args.multiple_of) + logger.info(f"Setting inter_size to {args.inter_size}.") + + if args.enable_pos_shift: + assert args.use_gpt_attention_plugin, "Position shift is only support in the gpt attention plugin." + assert args.enable_context_fmha or args.enable_context_fmha_fp32_acc + + if args.moe_num_experts and args.moe_top_k == 0: + args.moe_top_k = 1 + args.moe_config = MoeConfig(args.moe_num_experts, args.moe_top_k, + args.moe_tp_mode, + args.moe_renorm_mode).validate() + + if args.gather_all_token_logits: + args.gather_context_logits = True + args.gather_generation_logits = True + + return args + + +def get_model_object(args, mapping, trt_dtype=None): + if trt_dtype is None: + trt_dtype = str_dtype_to_trt(args.dtype) + # Initialize Module + logger.debug("[Python]llama exampels, Initialize tensorrt_llm.models.LLaMAForCausalLM....") + tensorrt_llm_llama = tensorrt_llm.models.LLaMAForCausalLM( + num_layers=args.n_layer, + num_heads=args.n_head, + num_kv_heads=args.n_kv_head, + hidden_size=args.n_embd, + vocab_size=args.vocab_size, + hidden_act=args.hidden_act, + max_position_embeddings=args.n_positions, + dtype=trt_dtype, + mlp_hidden_size=args.inter_size, + position_embedding_type=PositionEmbeddingType.rope_gpt_neox, + mapping=mapping, + rotary_base=args.rotary_base, + rotary_scaling=args.rotary_scaling, + use_parallel_embedding=args.use_parallel_embedding, + embedding_sharding_dim=args.embedding_sharding_dim, + quant_mode=args.quant_mode, + rms_norm_eps=args.rms_norm_eps, + use_fused_mlp=args.use_fused_mlp, + use_prompt_tuning=args.max_prompt_embedding_table_size > 0, + enable_pos_shift=args.enable_pos_shift, + dense_context_fmha=args.dense_context_fmha, + moe_config=args.moe_config, + max_lora_rank=args.max_lora_rank) + quantize_kwargs = {} + if args.use_smooth_quant or args.use_weight_only: + if args.weight_only_precision == 'int4_awq': + exclude_modules = ['lm_head'] if not args.quantize_lm_head else [] + quantize_kwargs = { + "group_size": args.group_size, + "zero": False, + "pre_quant_scale": True, + "exclude_modules": exclude_modules, + } + elif args.weight_only_precision == 'int4_gptq': + quantize_kwargs = { + "group_size": args.group_size, + "zero": True, + "pre_quant_scale": False, + } + elif args.enable_fp8 or args.fp8_kv_cache: + logger.info(f'Loading scaling factors from ' + f'{args.quantized_fp8_model_path}') + quant_scales = get_scaling_factors(args.quantized_fp8_model_path, + num_layers=args.n_layer, + quant_mode=args.quant_mode) + quantize_kwargs = {"quant_scales": quant_scales} + + if args.use_weight_only and args.moe_config.has_moe(): + if 'exclude_modules' in quantize_kwargs: + quantize_kwargs['exclude_modules'].append('router') + else: + quantize_kwargs['exclude_modules'] = ['lm_head', 'router'] + + tensorrt_llm_llama = quantize_model(tensorrt_llm_llama, args.quant_mode, + **quantize_kwargs) + if args.per_group: + if args.weight_only_precision == 'int4_awq': + load_from_awq_llama(tensorrt_llm_llama=tensorrt_llm_llama, + quant_ckpt_path=args.quant_ckpt_path, + quantize_lm_head=args.quantize_lm_head, + mapping=mapping, + dtype=args.dtype, + bin_model_dir=args.bin_model_dir) + else: + load_from_gptq_llama(tensorrt_llm_llama=tensorrt_llm_llama, + quant_ckpt_path=args.quant_ckpt_path, + mapping=mapping, + dtype=args.dtype, + bin_model_dir=args.bin_model_dir) + elif args.meta_ckpt_dir is not None: + load_from_meta_llama(tensorrt_llm_llama, args.meta_ckpt_dir, mapping, + args.dtype) + elif args.model_dir is not None: + logger.info(f'Loading HF LLaMA ... from {args.model_dir}') + tik = time.time() + if not args.load_by_shard: + if args.model_type == "llava": + hf_llava = LlavaForConditionalGeneration.from_pretrained( + args.model_dir, torch_dtype="auto") + hf_llama = hf_llava.language_model + else: + hf_model = LlamaForCausalLM if args.model_type != "mixtral" else MixtralForCausalLM + hf_llama = hf_model.from_pretrained( + args.model_dir, + device_map={ + "model": "cpu", + "lm_head": "cpu", + "embed_tokens": "cpu", + "layers": "cpu", + "norm": "cpu", + }, # Load to CPU memory + torch_dtype='auto', + ) + use_gemm_woq_plugin = not args.disable_weight_only_quant_plugin + # hf_llama.config.num_hidden_layers = 1 # only for debug + load_from_hf_llama(tensorrt_llm_llama, + hf_llama, + mapping=mapping, + dtype=args.dtype, + use_gemm_woq_plugin=use_gemm_woq_plugin, + lora_config=args.lora_config) + del hf_llama + else: + load_from_hf_checkpoint(tensorrt_llm_llama, + args.model_dir, + mapping, + dtype=args.dtype, + lora_config=args.lora_config) + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'HF LLaMA loaded. Total time: {t}') + + elif args.bin_model_dir is not None: + load_from_binary(tensorrt_llm_llama, + args.bin_model_dir, + mapping, + fp16=(args.dtype == 'float16'), + multi_query_mode=(args.n_kv_head != args.n_head)) + + return tensorrt_llm_llama + + +def update_plugin_configs(args, network): + if args.use_gpt_attention_plugin: + network.plugin_config.set_gpt_attention_plugin( + dtype=args.use_gpt_attention_plugin) + if args.use_gemm_plugin: + if not args.enable_fp8: + network.plugin_config.set_gemm_plugin(dtype=args.use_gemm_plugin) + else: + logger.info( + "Gemm plugin does not support FP8. Disabled Gemm plugin.") + if args.use_rmsnorm_plugin: + network.plugin_config.set_rmsnorm_plugin(dtype=args.use_rmsnorm_plugin) + if args.use_lora_plugin: + network.plugin_config.set_lora_plugin(dtype=args.use_lora_plugin) + if args.use_lookup_plugin: + network.plugin_config.set_lookup_plugin(dtype=args.use_lookup_plugin) + if args.use_gather_last_token_plugin: + network.plugin_config.set_gather_last_token_plugin(dtype=args.use_gather_last_token_plugin) + if args.use_activation_plugin: + network.plugin_config.set_activation_plugin(dtype=args.use_activation_plugin) + if args.use_elementwise_plugin: + network.plugin_config.set_elementwise_plugin(dtype=args.use_elementwise_plugin) + if args.use_cast_plugin: + network.plugin_config.set_cast_plugin() + + # Quantization plugins. + if args.use_smooth_quant: + network.plugin_config.set_smooth_quant_gemm_plugin(dtype=args.dtype) + network.plugin_config.set_rmsnorm_quantization_plugin(dtype=args.dtype) + network.plugin_config.set_quantize_tensor_plugin() + network.plugin_config.set_quantize_per_token_plugin() + assert not (args.enable_context_fmha and args.enable_context_fmha_fp32_acc) + if args.enable_context_fmha: + network.plugin_config.set_context_fmha(ContextFMHAType.enabled) + if args.enable_context_fmha_fp32_acc: + network.plugin_config.set_context_fmha( + ContextFMHAType.enabled_with_fp32_acc) + if args.multi_block_mode: + network.plugin_config.enable_mmha_multi_block_mode() + if not args.disable_xqa: + network.plugin_config.enable_xqa_optimization() + + if args.use_weight_only and not args.disable_weight_only_quant_plugin: + if args.per_group: + network.plugin_config.set_weight_only_groupwise_quant_matmul_plugin( + dtype=args.dtype) + else: + network.plugin_config.set_weight_only_quant_matmul_plugin( + dtype=args.dtype) + if args.world_size > 1: + network.plugin_config.set_nccl_plugin(args.dtype, + args.use_custom_all_reduce) + if args.remove_input_padding: + network.plugin_config.enable_remove_input_padding() + if args.paged_kv_cache: + network.plugin_config.enable_paged_kv_cache(args.tokens_per_block) + return + + +def build_rank_engine(builder: Builder, + builder_config: tensorrt_llm.builder.BuilderConfig, + engine_name, rank, args): + ''' + @brief: Build the engine on the given rank. + @param rank: The rank to build the engine. + @param args: The cmd line arguments. + @return: The built engine. + ''' + dtype = str_dtype_to_trt(args.dtype) + mapping = Mapping(world_size=args.world_size, + rank=rank, + tp_size=args.tp_size, + pp_size=args.pp_size) + + assert args.n_layer % args.pp_size == 0, \ + f"num_layers {args.n_layer} must be a multiple of pipeline parallelism size {args.pp_size}" + + # FIXME (Not Support libnvidia-ml.so) + # profiler.print_memory_usage(f'Rank {rank} Engine build starts') + # Initialize Module + tensorrt_llm_llama = get_model_object(args, + mapping=mapping, + trt_dtype=dtype) + + # FIXME (Not Support libnvidia-ml.so) + # profiler.print_memory_usage(f'Rank {rank} model weight loaded.') + + # Module -> Network + logger.debug("[Python]llama exampels, convert module to network....") + network = builder.create_network() + network.trt_network.name = engine_name + update_plugin_configs(args, network) + + if args.use_paged_context_fmha: + assert args.enable_context_fmha or args.enable_context_fmha_fp32_acc, "context fmha must be enabled" + network.plugin_config.set_paged_context_fmha() + + logger.debug(f"[Python]llama exampels, network.plugin_config: \n{network.plugin_config}") + with net_guard(network): + # Prepare + network.set_named_parameters(tensorrt_llm_llama.named_parameters()) + + # Forward + inputs = tensorrt_llm_llama.prepare_inputs( + max_batch_size=args.max_batch_size, + max_input_len=args.max_input_len, + max_seq_len=args.max_input_len + args.max_output_len, + use_cache=True, + max_beam_width=args.max_beam_width, + max_num_tokens=args.max_num_tokens, + prompt_embedding_table_size=args.max_prompt_embedding_table_size, + gather_context_logits=args.gather_context_logits, + gather_generation_logits=args.gather_generation_logits, + lora_target_modules=args.lora_target_modules) + logger.info(f"[Python]llama exampels, forward....\n") + tensorrt_llm_llama(*inputs) + logger.info(f"[Python]llama exampels, forward finished\n") + if args.enable_debug_output: + # mark intermediate nodes' outputs + for k, v in tensorrt_llm_llama.named_network_outputs(): + logger.debug(f"enable_debug_output, debug tensor name: {k}") + v = v.trt_tensor + v.name = k + network.trt_network.mark_output(v) + v.dtype = dtype + if args.visualize: + model_path = os.path.join(args.output_dir, 'test.onnx') + to_onnx(network.trt_network, model_path) + + logger.debug("[Python]llama examples, tensorrt_llm.graph_rewriting.optimize....") + tensorrt_llm.graph_rewriting.optimize(network) + + engine = None + + # Network -> Engine + logger.debug("[Python]llama examples, builder.build_engine....") + engine = builder.build_engine(network, builder_config) + if rank == 0: + config_path = os.path.join(args.output_dir, 'config.json') + builder.save_config(builder_config, config_path) + + return engine + + +def get_builder_config_namespace(args, cache): + # NOTE: int8 flag is required to be true when INT8 tensors are exposed to TRT + # TRT-LLM has INT8 I/O when act/weights are quantized without group-scaling (AWQ, GPTQ) + # OR INT8 KV cache is set to contiguous (without paged KV cache enabled). + int8_trt_flag = (args.quant_mode.has_act_or_weight_quant() + and not args.quant_mode.has_per_group_scaling()) or ( + not args.paged_kv_cache + and args.quant_mode.has_int8_kv_cache()) + config = argparse.Namespace( + name=MODEL_NAME, + precision=args.dtype, + timing_cache=args.timing_cache if cache is None else cache, + profiling_verbosity=args.profiling_verbosity, + tensor_parallel=args.tp_size, + pipeline_parallel=args.pp_size, + parallel_build=args.parallel_build, + num_layers=args.n_layer, + num_heads=args.n_head, + num_kv_heads=args.n_kv_head, + hidden_size=args.n_embd, + vocab_size=args.vocab_size, + hidden_act=args.hidden_act, + max_position_embeddings=args.n_positions, + max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, + max_input_len=args.max_input_len, + max_output_len=args.max_output_len, + max_num_tokens=args.max_num_tokens, + int8=int8_trt_flag, + quant_mode=args.quant_mode, + strongly_typed=args.strongly_typed, + opt_level=args.builder_opt, + max_prompt_embedding_table_size=args.max_prompt_embedding_table_size, + gather_context_logits=args.gather_context_logits, + gather_generation_logits=args.gather_generation_logits, + lora_target_modules=args.lora_target_modules, + mlp_hidden_size=args.inter_size, + hf_modules_to_trtllm_modules=args.lora_config. + hf_modules_to_trtllm_modules, + trtllm_modules_to_hf_modules=args.lora_config. + trtllm_modules_to_hf_modules, + ) + return config + + +def build(rank, args): + torch.cuda.set_device(rank % args.gpus_per_node) + logger.set_level(args.log_level) + os.makedirs(args.output_dir, exist_ok=True) + + # when doing serializing build, all ranks share one engine + builder = Builder() + cache = None + for cur_rank in range(args.world_size): + # skip other ranks if parallel_build is enabled + if args.parallel_build and cur_rank != rank: + continue + tik = time.time() + + # NOTE: int8 flag is required to be true when INT8 tensors are exposed to TRT + # TRT-LLM has INT8 I/O when act/weights are quantized without group-scaling (AWQ, GPTQ) + # OR INT8 KV cache is set to contiguous (without paged KV cache enabled). + int8_trt_flag = (args.quant_mode.has_act_or_weight_quant() + and not args.quant_mode.has_per_group_scaling()) or ( + not args.paged_kv_cache + and args.quant_mode.has_int8_kv_cache()) + builder_config = builder.create_builder_config( + **vars(get_builder_config_namespace(args, cache))) + engine_name = get_engine_name(MODEL_NAME, args.dtype, args.tp_size, + args.pp_size, cur_rank) + logger.debug("[Python]llama example, build_rank_engine....") + engine = build_rank_engine(builder, builder_config, engine_name, + cur_rank, args) + assert engine is not None, f'Failed to build engine for rank {cur_rank}' + + local_num_kv_heads = (args.n_kv_head + args.world_size - + 1) // args.world_size + kv_dtype = str_dtype_to_trt(args.dtype) + if args.quant_mode.has_int8_kv_cache(): + kv_dtype = str_dtype_to_trt('int8') + elif args.quant_mode.has_fp8_kv_cache(): + kv_dtype = str_dtype_to_trt('fp8') + + # FIXME (Not Support libnvidia-ml.so) + # profiler.check_gpt_mem_usage( + # engine=engine, + # kv_dtype=kv_dtype, + # use_gpt_attention_plugin=args.use_gpt_attention_plugin, + # paged_kv_cache=args.paged_kv_cache, + # max_batch_size=args.max_batch_size, + # max_beam_width=args.max_beam_width, + # max_seq_len=args.max_input_len + args.max_output_len, + # local_num_kv_heads=local_num_kv_heads, + # head_size=args.n_embd / args.n_head, + # num_layers=args.n_layer) + + if cur_rank == 0: + # Use in-memory timing cache for multiple builder passes. + if not args.parallel_build: + cache = builder_config.trt_builder_config.get_timing_cache() + + serialize_engine(engine, os.path.join(args.output_dir, engine_name)) + del engine + # FIXME (Not Support libnvidia-ml.so) + # profiler.print_memory_usage(f'Rank {cur_rank} Engine serialized') + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info( + f'Rank {cur_rank} Engine build time: {t} - {tok - tik} (sec)') + + if rank == 0: + ok = builder.save_timing_cache( + builder_config, os.path.join(args.output_dir, "model.cache")) + assert ok, "Failed to save timing cache." + + +if __name__ == '__main__': + args = parse_arguments() + print(args) + tik = time.time() + if args.parallel_build and args.world_size > 1 and \ + torch.cuda.device_count() >= args.world_size: + logger.warning( + f'Parallelly build TensorRT engines. Please make sure that all of the {args.world_size} GPUs are totally free.' + ) + mp.spawn(build, nprocs=args.world_size, args=(args, )) + else: + args.parallel_build = False + logger.info('Serially build TensorRT engines.') + build(0, args) + + tok = time.time() + build_engine_time = tok - tik + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'Total time of building all {args.world_size} engines: {t}') + + if args.total_build_time_target != 0: + status = build_engine_time <= args.total_build_time_target + if status: + print("successful.") + else: + print(f"Build engine time check failed! Target: {args.total_build_time_target}, Actual: {build_engine_time}") + sys.exit(int(not status)) diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/run.py b/models/nlp/large_language_model/llama2-13b/trtllm/run.py new file mode 100644 index 0000000000000000000000000000000000000000..3899ec9d55a33bca6eeeac4840353345467b474d --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/run.py @@ -0,0 +1,539 @@ +# 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. + +import argparse +import ast +import csv +from pathlib import Path +import sys +import time + +import numpy as np +import torch +import tensorrt_llm +import tensorrt_llm.profiler +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner + +from utils import (DEFAULT_HF_MODEL_DIRS, DEFAULT_PROMPT_TEMPLATES, + load_tokenizer, read_model_name, throttle_generator) + +if PYTHON_BINDINGS: + from tensorrt_llm.runtime import ModelRunnerCpp + + +def parse_arguments(args=None): + parser = argparse.ArgumentParser() + parser.add_argument('--max_output_len', type=int, required=True) + parser.add_argument( + '--max_attention_window_size', + type=int, + default=None, + help= + 'The attention window size that controls the sliding window attention / cyclic kv cache behaviour' + ) + parser.add_argument('--sink_token_length', + type=int, + default=None, + help='The sink token length.') + parser.add_argument('--log_level', type=str, default='error') + parser.add_argument('--engine_dir', type=str, default='engine_outputs') + parser.add_argument('--use_py_session', + default=False, + action='store_true', + help="Whether or not to use Python runtime session") + parser.add_argument( + '--input_text', + type=str, + nargs='+', + default=["Born in north-east France, Soyer trained as a"]) + parser.add_argument( + '--no_prompt_template', + dest='use_prompt_template', + default=True, + action='store_false', + help= + "Whether or not to use default prompt template to wrap the input text.") + parser.add_argument( + '--input_file', + type=str, + help= + 'CSV or Numpy file containing tokenized input. Alternative to text input.', + default=None) + parser.add_argument('--max_input_length', type=int, default=923) + 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( + '--output_logits_npy', + type=str, + help= + 'Numpy file where the generation logits are stored. Use only when num_beams==1', + default=None) + parser.add_argument('--tokenizer_dir', + help="HF tokenizer config path", + default='gpt2') + parser.add_argument( + '--tokenizer_type', + help= + 'Specify that argument when providing a .model file as the tokenizer_dir. ' + 'It allows AutoTokenizer to instantiate the correct tokenizer type.') + parser.add_argument('--vocab_file', + help="Used for sentencepiece tokenizers") + parser.add_argument('--num_beams', + type=int, + help="Use beam search if num_beams >1", + default=1) + parser.add_argument('--temperature', type=float, default=1.0) + parser.add_argument('--top_k', type=int, default=1) + parser.add_argument('--top_p', type=float, default=0.0) + parser.add_argument('--length_penalty', type=float, default=1.0) + parser.add_argument('--repetition_penalty', type=float, default=1.0) + parser.add_argument('--presence_penalty', type=float, default=0.0) + parser.add_argument('--frequency_penalty', type=float, default=0.0) + parser.add_argument('--debug_mode', + default=False, + action='store_true', + help="Whether or not to turn on the debug mode") + parser.add_argument('--no_add_special_tokens', + dest='add_special_tokens', + default=True, + action='store_false', + help="Whether or not to add special tokens") + 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_path', + type=str, + help="Path to .npy file, exported by nemo_prompt_convert.py") + parser.add_argument( + '--prompt_tasks', + help="Comma-separated list of tasks for prompt tuning, e.g., 0,3,1,0") + parser.add_argument('--lora_dir', + type=str, + default=None, + nargs="+", + help="The directory of LoRA weights") + parser.add_argument( + '--lora_task_uids', + type=str, + default=None, + nargs="+", + help="The list of LoRA task uids; use -1 to disable the LoRA module") + parser.add_argument('--lora_ckpt_source', + type=str, + default="hf", + choices=["hf", "nemo"], + help="The source of lora checkpoint.") + parser.add_argument( + '--num_prepend_vtokens', + nargs="+", + type=int, + help="Number of (default) virtual tokens to prepend to each sentence." + " For example, '--num_prepend_vtokens=10' will prepend the tokens" + " [vocab_size, vocab_size + 1, ..., vocab_size + 9] to the sentence.") + parser.add_argument( + '--run_profiling', + default=False, + action='store_true', + help="Run several 10 iterations to profile the inference latencies.") + parser.add_argument( + '--medusa_choices', + type=str, + default=None, + help="Medusa choice to use, if not none, will use Medusa decoding." + " E.g.: [[0, 0, 0, 0], [0, 1, 0], [1, 0], [1, 1]] for 9 medusa tokens." + ) + parser.add_argument('--target_load_engine_time', + type=float, + default=0) + parser.add_argument('--target_qps', + type=float, + default=0) + + return parser.parse_args(args=args) + + +def parse_input(tokenizer, + input_text=None, + prompt_template=None, + input_file=None, + add_special_tokens=True, + max_input_length=923, + pad_id=None, + num_prepend_vtokens=[], + model_name=None): + if pad_id is None: + pad_id = tokenizer.pad_token_id + + batch_input_ids = [] + if input_file is None: + for curr_text in input_text: + if prompt_template is not None: + curr_text = prompt_template.format(input_text=curr_text) + input_ids = tokenizer.encode(curr_text, + add_special_tokens=add_special_tokens, + truncation=True, + max_length=max_input_length) + batch_input_ids.append(input_ids) + 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_ids = np.array(line, dtype='int32') + batch_input_ids.append(input_ids[-max_input_length:]) + elif input_file.endswith('.npy'): + inputs = np.load(input_file) + for row in inputs: + input_ids = row[row != pad_id] + batch_input_ids.append(input_ids[-max_input_length:]) + elif input_file.endswith('.txt'): + with open(input_file, 'r', encoding='utf-8', + errors='replace') as txt_file: + input_text = txt_file.read() + input_ids = tokenizer.encode( + input_text, + add_special_tokens=add_special_tokens, + truncation=True, + max_length=max_input_length) + batch_input_ids.append(input_ids) + else: + print('Input file format not supported.') + raise SystemExit + + if num_prepend_vtokens: + assert len(num_prepend_vtokens) == len(batch_input_ids) + base_vocab_size = tokenizer.vocab_size - len( + tokenizer.special_tokens_map.get('additional_special_tokens', [])) + for i, length in enumerate(num_prepend_vtokens): + batch_input_ids[i] = list( + range(base_vocab_size, + base_vocab_size + length)) + batch_input_ids[i] + if model_name == 'glm_10b': + for ids in batch_input_ids: + ids.append(tokenizer.sop_token_id) + batch_input_ids = [ + torch.tensor(x, dtype=torch.int32) for x in batch_input_ids + ] + return batch_input_ids + + +def print_output(tokenizer, + output_ids, + input_lengths, + sequence_lengths, + output_csv=None, + output_npy=None, + context_logits=None, + generation_logits=None, + output_logits_npy=None): + batch_size, num_beams, _ = output_ids.size() + if output_csv is None and output_npy is None: + for batch_idx in range(batch_size): + inputs = output_ids[batch_idx][0][:input_lengths[batch_idx]].tolist( + ) + input_text = tokenizer.decode(inputs) + print(f'Input [Text {batch_idx}]: \"{input_text}\"') + for beam in range(num_beams): + output_begin = input_lengths[batch_idx] + output_end = sequence_lengths[batch_idx][beam] + outputs = output_ids[batch_idx][beam][ + output_begin:output_end].tolist() + output_text = tokenizer.decode(outputs) + print( + f'Output [Text {batch_idx} Beam {beam}]: \"{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) + + # Save context logits + if context_logits is not None and output_logits_npy is not None: + context_logits = torch.cat(context_logits, axis=0) + vocab_size_padded = context_logits.shape[-1] + context_logits = context_logits.reshape([1, -1, vocab_size_padded]) + + output_context_logits_npy = output_logits_npy.split( + '.npy')[0] + "_context" + output_context_logits_file = Path(output_context_logits_npy) + context_outputs = np.array( + context_logits.squeeze(0).cpu().contiguous(), + dtype='float32') # [promptLengthSum, vocabSize] + np.save(output_context_logits_file, context_outputs) + + # Save generation logits + if generation_logits is not None and output_logits_npy is not None and num_beams == 1: + output_generation_logits_npy = output_logits_npy.split( + '.npy')[0] + "_generation" + output_generation_logits_file = Path(output_generation_logits_npy) + generation_outputs = np.array(generation_logits.cpu().contiguous(), + dtype='float32') + np.save(output_generation_logits_file, generation_outputs) + + +def check_status(args, load_engine_time, qps): + print("==================== check status ====================") + successful = True + if args.target_load_engine_time != 0 and load_engine_time > args.target_load_engine_time: + print(f"Load engine time check failed! Target: {args.target_load_engine_time}, Actual: {load_engine_time}") + successful = False + if args.target_qps != 0 and qps < args.target_qps: + print(f"Performance check failed! Target: {args.target_qps}, Actual: {qps}") + successful = False + return successful + + +def main(args): + runtime_rank = tensorrt_llm.mpi_rank() + logger.set_level(args.log_level) + + model_name = read_model_name(args.engine_dir) + if args.tokenizer_dir is None: + args.tokenizer_dir = DEFAULT_HF_MODEL_DIRS[model_name] + + tokenizer, pad_id, end_id = load_tokenizer( + tokenizer_dir=args.tokenizer_dir, + vocab_file=args.vocab_file, + model_name=model_name, + tokenizer_type=args.tokenizer_type, + ) + + # # An example to stop generation when the model generate " London" on first sentence, " eventually became" on second sentence + # stop_words_list = [[" London"], ["eventually became"]] + # stop_words_list = tensorrt_llm.runtime.to_word_list_format(stop_words_list, tokenizer) + # stop_words_list = torch.Tensor(stop_words_list).to(torch.int32).to("cuda").contiguous() + stop_words_list = None + + # # An example to prevent generating " chef" on first sentence, " eventually" and " chef before" on second sentence + # bad_words_list = [[" chef"], [" eventually, chef before"]] + # bad_words_list = tensorrt_llm.runtime.to_word_list_format(bad_words_list, tokenizer) + # bad_words_list = torch.Tensor(bad_words_list).to(torch.int32).to("cuda").contiguous() + bad_words_list = None + + prompt_template = None + if args.use_prompt_template and model_name in DEFAULT_PROMPT_TEMPLATES: + prompt_template = DEFAULT_PROMPT_TEMPLATES[model_name] + batch_input_ids = parse_input(tokenizer=tokenizer, + input_text=args.input_text, + prompt_template=prompt_template, + input_file=args.input_file, + add_special_tokens=args.add_special_tokens, + max_input_length=args.max_input_length, + pad_id=pad_id, + num_prepend_vtokens=args.num_prepend_vtokens, + model_name=model_name) + input_lengths = [x.size(0) for x in batch_input_ids] + + if not PYTHON_BINDINGS and not args.use_py_session: + logger.warning( + "Python bindings of C++ session is unavailable, fallback to Python session." + ) + args.use_py_session = True + if args.debug_mode and not args.use_py_session: + logger.warning( + "Debug mode is not supported in C++ session for now, fallback to Python session." + ) + args.use_py_session = True + runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp + runner_kwargs = dict(engine_dir=args.engine_dir, + lora_dir=args.lora_dir, + rank=runtime_rank, + debug_mode=args.debug_mode, + lora_ckpt_source=args.lora_ckpt_source) + if args.medusa_choices is not None: + args.medusa_choices = ast.literal_eval(args.medusa_choices) + assert args.use_py_session, "Medusa is only supported by py_session" + assert args.temperature == 0, "Medusa should use temperature == 0" + assert args.num_beams == 1, "Medusa should use num_beams == 1" + runner_kwargs.update(medusa_choices=args.medusa_choices) + if not args.use_py_session: + runner_kwargs.update( + max_batch_size=len(batch_input_ids), + max_input_len=max(input_lengths), + max_output_len=args.max_output_len, + max_beam_width=args.num_beams, + max_attention_window_size=args.max_attention_window_size, + sink_token_length=args.sink_token_length, + ) + runner = runner_cls.from_dir(**runner_kwargs) + + torch.cuda.synchronize() + start_time = time.time() + with torch.no_grad(): + outputs = runner.generate( + batch_input_ids, + max_new_tokens=args.max_output_len, + max_attention_window_size=args.max_attention_window_size, + sink_token_length=args.sink_token_length, + end_id=end_id, + pad_id=pad_id, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + num_beams=args.num_beams, + length_penalty=args.length_penalty, + repetition_penalty=args.repetition_penalty, + presence_penalty=args.presence_penalty, + frequency_penalty=args.frequency_penalty, + stop_words_list=stop_words_list, + bad_words_list=bad_words_list, + lora_uids=args.lora_task_uids, + prompt_table_path=args.prompt_table_path, + prompt_tasks=args.prompt_tasks, + streaming=args.streaming, + output_sequence_lengths=True, + return_dict=True, + medusa_choices=args.medusa_choices) + torch.cuda.synchronize() + + status = False + end_time = time.time() + if runtime_rank == 0: + num_inputs = sum([torch.numel(x) for x in batch_input_ids]) + num_outputs = torch.numel(outputs["output_ids"]) + num_gens = num_outputs - num_inputs + + load_engine_time = tensorrt_llm.profiler.elapsed_time_in_sec("load tensorrt_llm engine") + qps = num_gens/(end_time-start_time) + logger.info(f'Load engine takes: {load_engine_time} sec') + print(f"input tokens: {num_inputs}, generate tokens: {num_gens}, QPS: {qps}") + status = check_status(args, load_engine_time, qps) + else: + status = True + + if args.streaming: + for curr_outputs in throttle_generator(outputs, + args.streaming_interval): + if runtime_rank == 0: + output_ids = curr_outputs['output_ids'] + sequence_lengths = curr_outputs['sequence_lengths'] + print_output(tokenizer, + output_ids, + input_lengths, + sequence_lengths, + output_csv=args.output_csv, + output_npy=args.output_npy) + else: + if runtime_rank == 0: + output_ids = outputs['output_ids'] + sequence_lengths = outputs['sequence_lengths'] + context_logits = None + generation_logits = None + if runner.gather_context_logits: + context_logits = outputs['context_logits'] + if runner.gather_generation_logits: + generation_logits = outputs['generation_logits'] + print_output(tokenizer, + output_ids, + input_lengths, + sequence_lengths, + output_csv=args.output_csv, + output_npy=args.output_npy, + context_logits=context_logits, + generation_logits=generation_logits, + output_logits_npy=args.output_logits_npy) + + if args.run_profiling: + ite = 10 + # warmup + for _ in range(ite): + with torch.no_grad(): + outputs = runner.generate( + batch_input_ids, + max_new_tokens=args.max_output_len, + max_attention_window_size=args.max_attention_window_size, + end_id=end_id, + pad_id=pad_id, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + num_beams=args.num_beams, + length_penalty=args.length_penalty, + repetition_penalty=args.repetition_penalty, + presence_penalty=args.presence_penalty, + frequency_penalty=args.frequency_penalty, + stop_words_list=stop_words_list, + bad_words_list=bad_words_list, + lora_uids=args.lora_task_uids, + prompt_table_path=args.prompt_table_path, + prompt_tasks=args.prompt_tasks, + streaming=args.streaming, + output_sequence_lengths=True, + return_dict=True) + torch.cuda.synchronize() + + tensorrt_llm.profiler.start("tmp") + for _ in range(ite): + with torch.no_grad(): + outputs = runner.generate( + batch_input_ids, + max_new_tokens=args.max_output_len, + max_attention_window_size=args.max_attention_window_size, + end_id=end_id, + pad_id=pad_id, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + num_beams=args.num_beams, + length_penalty=args.length_penalty, + repetition_penalty=args.repetition_penalty, + presence_penalty=args.presence_penalty, + frequency_penalty=args.frequency_penalty, + stop_words_list=stop_words_list, + bad_words_list=bad_words_list, + lora_uids=args.lora_task_uids, + prompt_table_path=args.prompt_table_path, + prompt_tasks=args.prompt_tasks, + streaming=args.streaming, + output_sequence_lengths=True, + return_dict=True) + torch.cuda.synchronize() + tensorrt_llm.profiler.stop("tmp") + + print( + f"batch_size: {len(batch_input_ids)}, avg latency of {ite} iterations: : {tensorrt_llm.profiler.elapsed_time_in_sec('tmp') / ite} sec" + ) + if status: + print("successful.") + else: + print("failed.") + sys.exit(int(not status)) + + +if __name__ == '__main__': + args = parse_arguments() + print(args) + main(args) diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/scripts/requirements.txt b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..05ba0afcd5cf0fd9dd88900b82e094b47a009496 --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/requirements.txt @@ -0,0 +1,45 @@ +# Copyright (c) 2024, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. +# All Rights Reserved. +# +# 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. + +accelerate +build +colored +# cuda-python # Do not override the custom version of cuda-python installed in the NGC PyTorch image. +diffusers +lark +mpi4py +numpy +onnx>=1.12.0 +polygraphy +psutil +pybind11 +pynvml>=11.5.0 +sentencepiece>=0.1.99 +# tensorrt==9.2.0.post12.dev5 +# torch +# nvidia-ammo~=0.5.0; platform_machine=="x86_64" +transformers +wheel +optimum +evaluate +janus +parameterized +scikit-learn + +# special +scipy==1.11.4 +pandas==1.5.3 +nltk +rouge_score diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/scripts/set_environment.sh b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/set_environment.sh new file mode 100644 index 0000000000000000000000000000000000000000..557998a8fa5f5692015e2eb075ffd76ce2abc30e --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/set_environment.sh @@ -0,0 +1,22 @@ +# Copyright (c) 2024, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. +# All Rights Reserved. +# +# 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. + +#!/bin/bash +set -e + +PROJECT_DIR=$1 + +pip3 install -r "requirements.txt" + diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/scripts/test_trtllm_llama2_13b_gpu2.sh b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/test_trtllm_llama2_13b_gpu2.sh new file mode 100644 index 0000000000000000000000000000000000000000..2e0f26dd76cbf12ea7789b21727b20615e34b5ec --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/test_trtllm_llama2_13b_gpu2.sh @@ -0,0 +1,56 @@ +# Copyright (c) 2024, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. +# All Rights Reserved. +# +# 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. + +#!/bin/bash + +EXIT_STATUS=0 +LOG_LEVEL=info +BS=${BS:-1} +DTYPE=${DTYPE:-"float16"} + +PROJECT_DIR="./" + +DATASET_DIR=${DATASET_DIR:-"${PROJECT_DIR}/data/datasets_cnn_dailymail"} +MODEL_DIR=${MODEL_DIR:-"${PROJECT_DIR}/data/llama2-13b-chat"} +ENGINE_DIR=${ENGINE_DIR:-"${PROJECT_DIR}/checkpoints/"} + +export TLLM_LOG_LEVEL=${LOG_LEVEL} +export PLUGIN_DTYPE="float16" + +check_status() +{ + if ((${PIPESTATUS[0]} != 0));then + EXIT_STATUS=1 + fi +} + +export TASK_DATA_PATH=${DATASET_DIR} + +# target is 95% of best (load engine time: 41.74, rouge1: 29.21, tps: 15.23) +mpirun -n 2 --allow-run-as-root \ +python3 ${PROJECT_DIR}/summarize.py \ +--test_trt_llm \ +--log_level ${LOG_LEVEL} \ +--batch_size ${BS} \ +--data_type ${DTYPE} \ +--hf_model_dir ${MODEL_DIR} \ +--tokenizer_dir ${MODEL_DIR} \ +--tokenizer_type "llama" \ +--engine_dir ${ENGINE_DIR} \ +--target_load_engine_time 43.94 \ +--tensorrt_llm_rouge1_threshold 27.74 \ +--target_tps 14.46 \ +--use_py_session "$@"; check_status +exit ${EXIT_STATUS} diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/scripts/test_trtllm_llama2_13b_gpu2_build.sh b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/test_trtllm_llama2_13b_gpu2_build.sh new file mode 100644 index 0000000000000000000000000000000000000000..4919ffa3cefd7cf5c9f3c74a1186a24d00eb3a96 --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/scripts/test_trtllm_llama2_13b_gpu2_build.sh @@ -0,0 +1,53 @@ +# Copyright (c) 2024, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. +# All Rights Reserved. +# +# 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. + +#!/bin/bash + +EXIT_STATUS=0 +LOG_LEVEL=info +BS=${BS:-1} +DTYPE=${DTYPE:-"float16"} + +PROJECT_DIR="./" + +MODEL_DIR=${MODEL_DIR:-"${PROJECT_DIR}/data/llama2-13b-chat"} +OUTPUT_DIR=${OUTPUT_DIR:-"${PROJECT_DIR}/checkpoints/"} + +echo PROJECT_DIR : ${PROJECT_DIR} + +export TLLM_LOG_LEVEL=${LOG_LEVEL} +export PLUGIN_DTYPE="float16" + +check_status() +{ + if ((${PIPESTATUS[0]} != 0));then + EXIT_STATUS=1 + fi +} + +# best(build engine time: 223.33) is 95% of target +python3 ${PROJECT_DIR}/build.py \ +--log_level ${LOG_LEVEL} \ +--dtype ${DTYPE} \ +--model_dir ${MODEL_DIR} \ +--remove_input_padding \ +--use_gpt_attention_plugin float16 --use_gemm_plugin float16 \ +--enable_context_fmha \ +--disable_xqa \ +--world_size 2 \ +--tp_size 2 \ +--total_build_time_target 235.1 \ +--output_dir ${OUTPUT_DIR} "$@"; check_status +exit ${EXIT_STATUS} diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/summarize.py b/models/nlp/large_language_model/llama2-13b/trtllm/summarize.py new file mode 100644 index 0000000000000000000000000000000000000000..8e9437f506a09aefb4c6e63d2827e8bdea2814e4 --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/summarize.py @@ -0,0 +1,724 @@ +# 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. + +import argparse +import ast +import os +from pathlib import Path +import sys +import time + +import evaluate +import numpy as np +import torch +from datasets import load_dataset, load_from_disk +from transformers import (AutoModel, AutoModelForCausalLM, + AutoModelForSeq2SeqLM, GenerationConfig) +from utils import DEFAULT_HF_MODEL_DIRS, load_tokenizer, read_model_name + +import tensorrt_llm +import tensorrt_llm.profiler as profiler +from tensorrt_llm._utils import str_dtype_to_torch +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner +from tensorrt_llm.tools.ppl import ppl + +if PYTHON_BINDINGS: + from tensorrt_llm.runtime import ModelRunnerCpp + + +def check_status(args, load_engine_time, rouge1, tps): + print("==================== check status ====================") + successful = True + if args.target_load_engine_time != 0 and load_engine_time > args.target_load_engine_time: + print(f"Load engine time check failed! Target: {args.target_load_engine_time}, Actual: {load_engine_time}") + successful = False + if rouge1 < args.tensorrt_llm_rouge1_threshold: + print(f"Accuracy check failed! Target: {args.tensorrt_llm_rouge1_threshold}%, Actual: {rouge1}%") + successful = False + if args.target_tps != 0 and tps < args.target_tps: + print(f"Performance check failed! Target: {args.target_tps}, Actual: {tps}") + successful = False + return successful + + +def main(args): + runtime_rank = tensorrt_llm.mpi_rank() + logger.set_level(args.log_level) + + model_name = read_model_name(args.engine_dir) + if args.hf_model_dir is None: + args.hf_model_dir = DEFAULT_HF_MODEL_DIRS[model_name] + if args.tokenizer_dir is None: + args.tokenizer_dir = args.hf_model_dir + + test_hf = args.test_hf and runtime_rank == 0 # only run hf on rank 0 + test_trt_llm = args.test_trt_llm + profiler.start('load tokenizer') + tokenizer, pad_id, end_id = load_tokenizer( + tokenizer_dir=args.tokenizer_dir, + vocab_file=args.vocab_file, + model_name=model_name, + tokenizer_type=args.tokenizer_type, + ) + profiler.stop('load tokenizer') + logger.info( + f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' + ) + + if args.eval_task == 'code_completion': + dataset_name = "openai_humaneval" + dataset_revision = None + dataset_input_key = 'prompt' + dataset_output_key = 'canonical_solution' + dataset_split = 'test' + elif args.eval_task == 'summarize': + dataset_name = "ccdv/cnn_dailymail" + dataset_revision = "3.0.0" + dataset_input_key = 'article' + dataset_output_key = 'highlights' + dataset_split = 'test' + elif args.eval_task == 'summarize_long': + dataset_name = "tau/zero_scrolls" + dataset_revision = 'squality' + dataset_input_key = 'input' + dataset_output_key = 'output' + dataset_split = 'validation' # only this split contains reference strings + + + logger.info(f"prepare datasets....") + if os.getenv("TASK_DATA_PATH"): + dataset = load_from_disk(os.getenv("TASK_DATA_PATH"))[dataset_split] + else: + # dataset = load_dataset(dataset_name, + # dataset_revision, + # cache_dir=args.dataset_path, + # split=dataset_split, + # trust_remote_code=True) + + dataset = load_dataset(dataset_name, + dataset_revision, + cache_dir=args.dataset_path, + split=dataset_split) + + logger.info(f"datasets is ready.") + max_batch_size = args.batch_size + + # runtime parameters + top_k = args.top_k + top_p = args.top_p + output_len = args.output_len + test_token_num = args.max_input_length + max_attention_window_size = args.max_attention_window_size + sink_token_length = args.sink_token_length + + # random_seed = 5 + temperature = args.temperature + num_beams = args.num_beams + length_penalty = args.length_penalty + repetition_penalty = args.repetition_penalty + presence_penalty = args.presence_penalty + frequency_penalty = args.frequency_penalty + + if test_trt_llm: + if not PYTHON_BINDINGS and not args.use_py_session: + logger.warning( + "Python bindings of C++ session is unavailable, fallback to Python session." + ) + args.use_py_session = True + runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp + runner_kwargs = dict(engine_dir=args.engine_dir, + rank=runtime_rank, + debug_mode=args.debug_mode) + if args.medusa_choices is not None: + args.medusa_choices = ast.literal_eval(args.medusa_choices) + assert args.use_py_session, "Medusa is only supported by py_session" + assert args.temperature == 0, "Medusa should use temperature == 0" + assert args.num_beams == 1, "Medusa should use num_beams == 1" + runner_kwargs.update(medusa_choices=args.medusa_choices) + if not args.use_py_session: + runner_kwargs.update( + max_batch_size=max_batch_size, + max_input_len=test_token_num, + max_output_len=output_len, + max_beam_width=num_beams, + max_attention_window_size=max_attention_window_size, + sink_token_length=sink_token_length) + runner = runner_cls.from_dir(**runner_kwargs) + assert not (args.eval_ppl and not (runner.gather_context_logits and runner.gather_generation_logits)), \ + "PPL evaluation requires engine built with gather_all_token_logits enabled" + + if test_hf: + profiler.start('load HF model') + dtype_alias_mapping = { + 'fp32': 'float32', + 'fp16': 'float16', + 'bf16': 'bfloat16' + } + args.data_type = dtype_alias_mapping.get(args.data_type, args.data_type) + if model_name.startswith('chatglm'): + auto_model_cls = AutoModel + elif model_name.startswith('glm'): + auto_model_cls = AutoModelForSeq2SeqLM + else: + auto_model_cls = AutoModelForCausalLM + model = auto_model_cls.from_pretrained( + args.hf_model_dir, + trust_remote_code=True, + torch_dtype=str_dtype_to_torch(args.data_type), + device_map='auto' if args.hf_device_map_auto else None) + try: + model.to_bettertransformer() + except ValueError as e: + logger.warning( + f'Fail to call model.to_bettertransformer(), exception:\n{str(e)}' + ) + if not args.hf_device_map_auto: + model.cuda() + if model_name == 'qwen': + model.generation_config = GenerationConfig.from_pretrained( + args.hf_model_dir, trust_remote_code=True) + profiler.stop('load HF model') + logger.info( + f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' + ) + + output_dir = Path(args.output_dir) if args.output_dir else None + if output_dir is not None: + output_dir.mkdir(exist_ok=True, parents=True) + if test_trt_llm: + with (output_dir / 'trtllm.out').open('w') as f: + f.write(f'Engine path: {args.engine_dir}\n') + f.write(f'Tokenizer path: {args.tokenizer_dir}\n') + if test_hf: + with (output_dir / 'hf.out').open('w') as f: + f.write(f'Model path: {args.hf_model_dir}\n') + f.write(f'Tokenizer path: {args.tokenizer_dir}\n') + + def _prepare_inputs(batch_input_texts, + eval_task='summarize', + add_special_tokens=True): + batch_size = len(batch_input_texts) + append_str = ' TL;DR: ' if eval_task == 'summarize' else '' + batch_input_ids = [] + for i in range(batch_size): + curr_text = batch_input_texts[i] + append_str + curr_text = curr_text.strip().replace(" n't", "n't") + + # TODO: The below lines are used to be compatible with the original code; may need fix + if model_name.startswith(('chatglm2', 'chatglm3')): + input_ids = tokenizer.encode(curr_text, + return_tensors='pt').squeeze(0) + input_ids = input_ids[:test_token_num] + elif model_name == 'qwen': + from qwen.utils.utils import make_context + # use make_content to generate prompt + system_prompt = "You are a useful assistant, please directly output the corresponding summary according to the article entered by the user." + _, input_id_list = make_context( + tokenizer=tokenizer, + query=curr_text, + history=[], + system=system_prompt, + max_input_length=test_token_num, + ) + input_ids = torch.tensor(input_id_list) + else: + input_ids = tokenizer.encode( + curr_text, + return_tensors='pt', + add_special_tokens=add_special_tokens, + truncation=True, + max_length=test_token_num).squeeze(0) + + batch_input_ids.append(input_ids) + return batch_input_ids + + def eval_trt_llm(datapoint, + eval_task='summarize', + eval_ppl=False, + add_special_tokens=True): + batch_size = len(datapoint[dataset_input_key]) + batch_input_ids = _prepare_inputs(datapoint[dataset_input_key], + eval_task=eval_task, + add_special_tokens=add_special_tokens) + input_lengths = [x.size(0) for x in batch_input_ids] + + with torch.no_grad(): + outputs = runner.generate( + batch_input_ids, + max_new_tokens=output_len, + max_attention_window_size=max_attention_window_size, + sink_token_length=sink_token_length, + end_id=end_id, + pad_id=pad_id, + temperature=temperature, + top_k=top_k, + top_p=top_p, + num_beams=num_beams, + length_penalty=length_penalty, + repetition_penalty=repetition_penalty, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + output_sequence_lengths=True, + return_dict=True, + medusa_choices=args.medusa_choices) + torch.cuda.synchronize() + + # Extract a list of tensors of shape beam_width x output_ids. + if runtime_rank == 0: + output_ids = outputs['output_ids'] + output_beams_list = [ + tokenizer.batch_decode(output_ids[batch_idx, :, + input_lengths[batch_idx]:], + skip_special_tokens=True) + for batch_idx in range(batch_size) + ] + output_ids_list = [ + output_ids[batch_idx, :, input_lengths[batch_idx]:] + for batch_idx in range(batch_size) + ] + + ppls = [[] for _ in range(batch_size)] + seq_lengths_array = outputs["sequence_lengths"].cpu().tolist() + lengths_info = { + 'input_lengths': input_lengths, + 'seq_lengths': seq_lengths_array + } + if eval_ppl: + seq_lengths = outputs['sequence_lengths'] + context_logits = outputs['context_logits'] + # Remove the first generation logits which are same to last context logits + generation_logits = outputs['generation_logits'][:, :, 1:] + for batch_idx in range(batch_size): + # [batch, beam, step] + for beam_idx in range(num_beams): + curr_len = seq_lengths[batch_idx, beam_idx] + curr_ctx_len = input_lengths[batch_idx] + curr_gen_len = curr_len - curr_ctx_len + + curr_ids = output_ids[batch_idx, beam_idx, 1:curr_len] + curr_logits = torch.cat([ + context_logits[batch_idx], + generation_logits[batch_idx, + beam_idx, :curr_gen_len - 1] + ], + dim=0) + curr_ppl = ppl(curr_logits, curr_ids) + logger.debug( + f"TensorRT-LLM PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" + ) + ppls[batch_idx].append(curr_ppl) + + return output_beams_list, output_ids_list, ppls, lengths_info + return [], [], [], {} + + def eval_hf(datapoint, + eval_task='summarize', + eval_ppl=False, + add_special_tokens=True): + batch_size = len(datapoint[dataset_input_key]) + if batch_size > 1: + logger.warning( + f"HF does not support batch_size > 1 to verify correctness due to padding. Current batch size is {batch_size}" + ) + batch_input_ids = _prepare_inputs(datapoint[dataset_input_key], + eval_task=eval_task, + add_special_tokens=add_special_tokens) + input_lengths = [x.size(0) for x in batch_input_ids] + # Left padding for HF + max_length = max(input_lengths) + paddings = [ + torch.ones(max_length - l, dtype=torch.int32) * pad_id + for l in input_lengths + ] + batch_input_ids = [ + torch.cat([pad, x]) for x, pad in zip(batch_input_ids, paddings) + ] + batch_input_ids = torch.stack(batch_input_ids) + batch_input_ids = batch_input_ids.cuda() + + with torch.no_grad(): + outputs = model.generate(batch_input_ids, + max_new_tokens=output_len, + top_k=top_k, + temperature=temperature, + eos_token_id=end_id, + pad_token_id=pad_id, + num_beams=num_beams, + num_return_sequences=num_beams, + early_stopping=True, + length_penalty=length_penalty, + output_scores=True, + return_dict_in_generate=True) + if eval_ppl and batch_size == 1: + # model.generate cannot return context logits? + # Will cause additional latency + context_outputs = model(batch_input_ids) + + output_ids = outputs['sequences'] + tokens_list = output_ids[:, len(batch_input_ids[0]):].tolist() + output_ids = output_ids.reshape([batch_size, num_beams, -1]) + output_lines_list = [ + tokenizer.batch_decode(output_ids[:, i, + len(batch_input_ids[0]):], + skip_special_tokens=True) + for i in range(num_beams) + ] + + ppls = [[] for _ in range(batch_size)] + if eval_ppl and batch_size == 1: + # Only for batch size of 1 + seq_lens = (output_ids != end_id).logical_and( + output_ids != pad_id).sum(dim=-1) + context_logits = context_outputs['logits'] + # Remove the first generation logits which are same to last context logits + generation_logits = torch.stack(outputs['scores'][1:], dim=1) + _, max_gen_len, voc_size = generation_logits.size() + generation_logits = generation_logits.view(batch_size, num_beams, + max_gen_len, voc_size) + for batch_idx in range(batch_size): + for beam_idx in range(num_beams): + curr_len = seq_lens[batch_idx, beam_idx] + curr_ctx_len = input_lengths[batch_idx] + curr_gen_len = curr_len - curr_ctx_len + + curr_ids = output_ids[batch_idx, beam_idx, 1:curr_len] + curr_logits = torch.cat([ + context_logits[batch_idx], + generation_logits[batch_idx, + beam_idx, :curr_gen_len - 1] + ], + dim=0) + curr_ppl = ppl(curr_logits, curr_ids) + logger.debug( + f"HF PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" + ) + ppls[batch_idx].append(curr_ppl) + + return output_lines_list, tokens_list, ppls + + if test_trt_llm: + datapoint = dataset[0:1] + output, *_ = eval_trt_llm(datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + if runtime_rank == 0: + logger.info( + "---------------------------------------------------------") + logger.info("TensorRT-LLM Generated : ") + logger.info(f" Input : {datapoint[dataset_input_key]}") + logger.info(f"\n Reference : {datapoint[dataset_output_key]}") + logger.info(f"\n Output : {output}") + logger.info( + "---------------------------------------------------------") + if test_hf: + datapoint = dataset[0:1] + output, *_ = eval_hf(datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + logger.info("---------------------------------------------------------") + logger.info("HF Generated : ") + logger.info(f" Input : {datapoint[dataset_input_key]}") + logger.info(f"\n Reference : {datapoint[dataset_output_key]}") + logger.info(f"\n Output : {output}") + logger.info("---------------------------------------------------------") + + # TODO: Add random_seed flag in gptj + metric_tensorrt_llm = [evaluate.load("rouge") for _ in range(num_beams)] + metric_hf = [evaluate.load("rouge") for _ in range(num_beams)] + for i in range(num_beams): + metric_tensorrt_llm[i].seed = 0 + metric_hf[i].seed = 0 + ppls_trt_llm = [[] for _ in range(num_beams)] + ppls_hf = [[] for _ in range(num_beams)] + + ite_count = 0 + data_point_idx = 0 + total_output_token_count_trt_llm = 0 # only valid for runtime_rank == 0 + + if args.stability_test: + logger.info(f"stability test, need {args.stability_test_hours} hours") + else: + logger.info(f"dataset size: {len(dataset)}, max_ite: {args.max_ite}") + stability_start_time = time.time() + while (data_point_idx < len(dataset)) and (ite_count < args.max_ite): + if runtime_rank == 0: + logger.debug( + f"run data_point {data_point_idx} ~ {data_point_idx + max_batch_size}" + ) + datapoint = dataset[data_point_idx:(data_point_idx + max_batch_size)] + + if test_trt_llm: + profiler.start('tensorrt_llm') + output_tensorrt_llm, output_ids_trt_llm, curr_ppls_trt_llm, lengths_info = eval_trt_llm( + datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + profiler.stop('tensorrt_llm') + if runtime_rank == 0: + input_lengths = lengths_info['input_lengths'] + seq_lengths = lengths_info['seq_lengths'] + output_token_count_trt_llm = sum( + seq_lengths[idx][0] - input_lengths[idx] + for idx in range(len(input_lengths))) + total_output_token_count_trt_llm += output_token_count_trt_llm + + if test_hf: + profiler.start('hf') + output_hf, _, curr_ppls_hf = eval_hf( + datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + profiler.stop('hf') + + if runtime_rank == 0: + if test_trt_llm: + for batch_idx in range(len(output_tensorrt_llm)): + for beam_idx in range(num_beams): + metric_tensorrt_llm[beam_idx].add_batch( + predictions=[ + output_tensorrt_llm[batch_idx][beam_idx] + ], + references=[ + datapoint[dataset_output_key][batch_idx] + ]) + if args.eval_ppl: + ppls_trt_llm[beam_idx].append( + curr_ppls_trt_llm[batch_idx][beam_idx]) + if output_dir is not None: + # yapf: disable + for i in range(len(output_tensorrt_llm[0])): + for beam_idx in range(num_beams): + with (output_dir / 'trtllm.out').open('a') as f: + f.write(f'[{data_point_idx + i}] [Beam {beam_idx}] {output_tensorrt_llm[beam_idx][i]}\n') + # yapf: enable + if test_hf: + for beam_idx in range(num_beams): + for batch_idx in range(len(output_hf[beam_idx])): + metric_hf[beam_idx].add_batch( + predictions=[output_hf[beam_idx][batch_idx]], + references=[ + datapoint[dataset_output_key][batch_idx] + ]) + if args.eval_ppl and args.batch_size == 1: + ppls_hf[beam_idx].append( + curr_ppls_hf[batch_idx][beam_idx]) + if output_dir is not None: + # yapf: disable + for i in range(len(output_hf[0])): + for beam_idx in range(num_beams): + with (output_dir / 'hf.out').open('a') as f: + f.write(f'[{data_point_idx + i}] [Beam {beam_idx}] {output_hf[beam_idx][i]}\n') + # yapf: enable + + logger.debug('-' * 100) + logger.debug(f"Input : {datapoint[dataset_input_key]}") + if test_trt_llm: + logger.debug(f'TensorRT-LLM Output: {output_tensorrt_llm}') + if test_hf: + logger.debug(f'HF Output: {output_hf}') + logger.debug(f"Reference : {datapoint[dataset_output_key]}") + + data_point_idx += max_batch_size + ite_count += 1 + + if args.stability_test: + test_time_hours = round((time.time() - stability_start_time)/3600, 1) + if test_time_hours > args.stability_test_hours: + if runtime_rank == 0: + logger.info(f"Stability Test Finished. Total run {test_time_hours} hours.") + break + else: + data_point_idx = data_point_idx % len(dataset) + ite_count = ite_count % args.max_ite + if runtime_rank == 0 and ite_count % 1000 == 0: + logger.info(f"stability test, remain {round(args.stability_test_hours - test_time_hours, 1)} hours") + elif runtime_rank == 0 and ite_count % 10 == 0: + logger.info(f"data_point_idx: {data_point_idx}, ite_count: {ite_count}") + + if runtime_rank == 0: + if test_trt_llm: + np.random.seed(0) # rouge score use sampling to compute the score + logger.info( + f'TensorRT-LLM (total latency: {profiler.elapsed_time_in_sec("tensorrt_llm")} sec)' + ) + logger.info( + f'TensorRT-LLM (total output tokens: {total_output_token_count_trt_llm})' + ) + logger.info( + f'TensorRT-LLM (tokens per second: {total_output_token_count_trt_llm / profiler.elapsed_time_in_sec("tensorrt_llm")})' + ) + + rouge1 = 0 + tps = total_output_token_count_trt_llm / profiler.elapsed_time_in_sec("tensorrt_llm") + + for beam_idx in range(num_beams): + logger.info(f"TensorRT-LLM beam {beam_idx} result") + computed_metrics_tensorrt_llm = metric_tensorrt_llm[ + beam_idx].compute() + for key in computed_metrics_tensorrt_llm.keys(): + logger.info( + f' {key} : {computed_metrics_tensorrt_llm[key]*100}') + + if args.check_accuracy and beam_idx == 0: + assert computed_metrics_tensorrt_llm[ + 'rouge1'] * 100 > args.tensorrt_llm_rouge1_threshold + + if beam_idx == 0: + rouge1 = computed_metrics_tensorrt_llm['rouge1'] * 100 + + if args.eval_ppl: + logger.info( + f" Per-token perplexity: {np.mean(ppls_trt_llm[beam_idx])}" + ) + if args.check_accuracy and beam_idx == 0: + assert np.mean(ppls_trt_llm[beam_idx] + ) < args.tensorrt_llm_ppl_threshold + + load_engine_time = tensorrt_llm.profiler.elapsed_time_in_sec("load tensorrt_llm engine") + logger.info(f'Load engine takes: {load_engine_time} sec') + + status = check_status(args, load_engine_time, rouge1, tps) + if status: + print("successful.") + else: + print("failed.") + + sys.exit(int(not status)) + + if test_hf: + np.random.seed(0) # rouge score use sampling to compute the score + logger.info( + f'Hugging Face (total latency: {profiler.elapsed_time_in_sec("hf")} sec)' + ) + for beam_idx in range(num_beams): + logger.info(f"HF beam {beam_idx} result") + computed_metrics_hf = metric_hf[beam_idx].compute() + for key in computed_metrics_hf.keys(): + logger.info(f' {key} : {computed_metrics_hf[key]*100}') + if args.eval_ppl and args.batch_size == 1: + logger.info( + f" Per-token perplexity: {np.mean(ppls_hf[beam_idx])}") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--hf_model_dir', '--model_dir', type=str, default=None) + parser.add_argument( + '--tokenizer_dir', + default=None, + help='tokenizer path; defaults to hf_model_dir if left unspecified') + parser.add_argument( + '--tokenizer_type', + help= + 'Specify that argument when providing a .model file as the tokenizer_dir. ' + 'It allows AutoTokenizer to instantiate the correct tokenizer type.') + parser.add_argument('--vocab_file') + parser.add_argument('--test_hf', action='store_true') + parser.add_argument('--test_trt_llm', action='store_true') + parser.add_argument( + '--data_type', + type=str, + choices=['fp32', 'fp16', 'bf16', 'float32', 'float16', 'bfloat16'], + default='fp16') + parser.add_argument('--engine_dir', type=str, default='engine_outputs') + parser.add_argument('--use_py_session', + default=False, + action='store_true', + help="Whether or not to use Python runtime session") + parser.add_argument( + '--eval_task', + type=str, + default='summarize', + choices=['summarize', 'summarize_long', 'code_completion']) + parser.add_argument('--check_accuracy', action='store_true') + parser.add_argument('--tensorrt_llm_rouge1_threshold', + type=float, + default=15.0) + parser.add_argument('--eval_ppl', action='store_true') + parser.add_argument('--tensorrt_llm_ppl_threshold', + type=float, + default=15.0) + parser.add_argument('--target_load_engine_time', + type=float, + default=0) + parser.add_argument('--target_tps', + type=float, + default=0) + parser.add_argument('--dataset_path', type=str, default='') + parser.add_argument('--log_level', type=str, default='info') + parser.add_argument('--batch_size', type=int, default=1) + parser.add_argument('--max_ite', type=int, default=20) + parser.add_argument('--output_len', type=int, default=100) + parser.add_argument('--max_input_length', type=int, default=923) + parser.add_argument( + '--max_attention_window_size', + type=int, + default=None, + help= + 'The attention window size that controls the sliding window attention / cyclic kv cache behaviour' + ) + parser.add_argument('--sink_token_length', + type=int, + default=None, + help='The sink token length.') + parser.add_argument('--num_beams', type=int, default=1) + parser.add_argument('--temperature', type=float, default=1.0) + parser.add_argument('--top_k', type=int, default=1) + parser.add_argument('--top_p', type=float, default=0.0) + parser.add_argument('--length_penalty', type=float, default=1.0) + parser.add_argument('--repetition_penalty', type=float, default=1.0) + parser.add_argument('--presence_penalty', type=float, default=0.0) + parser.add_argument('--frequency_penalty', type=float, default=0.0) + parser.add_argument('--debug_mode', + default=False, + action='store_true', + help="Whether or not to turn on the debug mode") + parser.add_argument('--no_add_special_tokens', + dest='add_special_tokens', + default=True, + action='store_false', + help="Whether or not to add special tokens") + parser.add_argument( + '--hf_device_map_auto', + action='store_true', + help="Use device map 'auto' to load a pretrained HF model. This may " + "help to test a large model that cannot fit into a singlue GPU.") + parser.add_argument( + '--output_dir', + type=str, + default=None, + help="Directory where to save output sentences. 'trtllm.out' for " + "TensorRT-LLM outputs, and 'hf.out' for HF outputs. If None, do not " + "save outputs.") + parser.add_argument( + '--medusa_choices', + type=str, + default=None, + help="Medusa choice to use, if not none, will use Medusa decoding." + " E.g.: [[0, 0, 0, 0], [0, 1, 0], [1, 0], [1, 1]] for 9 medusa tokens." + ) + parser.add_argument('--stability_test', + default=False, + action='store_true', + help="Whether or not to run stability test for tensorrt_llm.") + parser.add_argument('--stability_test_hours', type=float, default=24.0) + args = parser.parse_args() + print(args) + main(args) diff --git a/models/nlp/large_language_model/llama2-13b/trtllm/utils.py b/models/nlp/large_language_model/llama2-13b/trtllm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..44042d9e2dcb44dd6cd917ab16a00010e4005202 --- /dev/null +++ b/models/nlp/large_language_model/llama2-13b/trtllm/utils.py @@ -0,0 +1,119 @@ +# 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. + +import json +from pathlib import Path +from typing import Optional + +from transformers import AutoTokenizer, T5Tokenizer + +import tensorrt_llm + +DEFAULT_HF_MODEL_DIRS = { + 'baichuan': 'baichuan-inc/Baichuan-13B-Chat', + 'bloom': 'bigscience/bloom-560m', + 'chatglm_6b': 'THUDM/chatglm-6b', + 'chatglm2_6b': 'THUDM/chatglm2-6b', + 'chatglm2_6b_32k': 'THUDM/chatglm2-6b-32k', + 'chatglm3_6b': 'THUDM/chatglm3-6b', + 'chatglm3_6b_base': 'THUDM/chatglm3-6b-base', + 'chatglm3_6b_32k': 'THUDM/chatglm3-6b-32k', + 'falcon': 'tiiuae/falcon-rw-1b', + 'glm_10b': 'THUDM/glm-10b', + 'gpt': 'gpt2-medium', + 'gptj': 'EleutherAI/gpt-j-6b', + 'gptneox': 'EleutherAI/gpt-neox-20b', + 'internlm': 'internlm/internlm-chat-7b', + 'llama': 'meta-llama/Llama-2-7b-hf', + 'mpt': 'mosaicml/mpt-7b', + 'phi': 'microsoft/phi-2', + 'opt': 'facebook/opt-350m', + 'qwen': 'Qwen/Qwen-7B', +} + +DEFAULT_PROMPT_TEMPLATES = { + 'internlm': + "<|User|>:{input_text}\n<|Bot|>:", + 'qwen': + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{input_text}<|im_end|>\n<|im_start|>assistant\n", +} + + +def read_model_name(engine_dir: str): + engine_version = tensorrt_llm.runtime.engine.get_engine_version(engine_dir) + + with open(Path(engine_dir) / "config.json", 'r') as f: + config = json.load(f) + + if engine_version is None: + return config['builder_config']['name'] + + return config['pretrained_config']['architecture'] + + +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 load_tokenizer(tokenizer_dir: Optional[str] = None, + vocab_file: Optional[str] = None, + model_name: str = 'gpt', + tokenizer_type: Optional[str] = None): + if vocab_file is None: + use_fast = True + if tokenizer_type is not None and tokenizer_type == "llama": + use_fast = False + # Should set both padding_side and truncation_side to be 'left' + tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir, + legacy=False, + padding_side='left', + truncation_side='left', + trust_remote_code=True, + tokenizer_type=tokenizer_type, + use_fast=use_fast) + else: + # For gpt-next, directly load from tokenizer.model + assert model_name == 'gpt' + tokenizer = T5Tokenizer(vocab_file=vocab_file, + padding_side='left', + truncation_side='left') + + if model_name == 'qwen': + with open(Path(tokenizer_dir) / "generation_config.json") as f: + gen_config = json.load(f) + chat_format = gen_config['chat_format'] + if chat_format == 'raw': + pad_id = gen_config['pad_token_id'] + end_id = gen_config['eos_token_id'] + elif chat_format == 'chatml': + pad_id = tokenizer.im_end_id + end_id = tokenizer.im_end_id + else: + raise Exception(f"unknown chat format: {chat_format}") + elif model_name == 'glm_10b': + pad_id = tokenizer.pad_token_id + end_id = tokenizer.eop_token_id + else: + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + pad_id = tokenizer.pad_token_id + end_id = tokenizer.eos_token_id + + return tokenizer, pad_id, end_id