mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-27 06:03:09 +08:00
d879430b04
* Update TensorRT-LLM --------- Co-authored-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
import asyncio
|
|
import inspect
|
|
import sys
|
|
from argparse import ArgumentParser
|
|
from typing import List, Optional
|
|
|
|
from tensorrt_llm import LLM, ModelConfig
|
|
|
|
# NOTE, Currently, the following examples are only available for LLaMA models.
|
|
|
|
|
|
def run_llm_from_huggingface_model(prompts: List[str],
|
|
llama_model_dir: str,
|
|
dump_engine_dir: Optional[str] = None):
|
|
''' Loading a HuggingFace model. '''
|
|
# Load the model from a local HuggingFace model
|
|
config = ModelConfig(model_dir=llama_model_dir)
|
|
llm = LLM(config)
|
|
if dump_engine_dir:
|
|
llm.save(dump_engine_dir)
|
|
|
|
for output in llm.generate(prompts):
|
|
print(output)
|
|
|
|
|
|
def run_llm_from_tllm_engine(prompts: List[str], llama_engine_dir: str):
|
|
''' Loading a built TensorRT-LLM engine. '''
|
|
config = ModelConfig(model_dir=llama_engine_dir)
|
|
llm = LLM(config)
|
|
|
|
for output in llm.generate(prompts):
|
|
print(output)
|
|
|
|
|
|
def run_llm_on_tensor_parallel(prompts: List[str], llama_model_dir: str):
|
|
''' Running LLM with Tensor Parallel on multiple GPUs. '''
|
|
config = ModelConfig(model_dir=llama_model_dir)
|
|
config.parallel_config.tp_size = 2 # 2 GPUs
|
|
|
|
llm = LLM(config)
|
|
|
|
for output in llm.generate(prompts):
|
|
print(output)
|
|
|
|
|
|
def run_llm_generate_async_example(prompts: List[str],
|
|
llama_model_dir: str,
|
|
streaming: bool = False):
|
|
''' Running LLM generation asynchronously. '''
|
|
config = ModelConfig(model_dir=llama_model_dir)
|
|
|
|
llm = LLM(config, async_mode=True)
|
|
|
|
async def task(prompt: str):
|
|
outputs = []
|
|
async for output in llm.generate_async(prompt, streaming=streaming):
|
|
outputs.append(output.text)
|
|
print(' '.join(outputs))
|
|
|
|
async def main():
|
|
tasks = [task(prompt) for prompt in prompts]
|
|
await asyncio.gather(*tasks)
|
|
|
|
asyncio.run(main())
|
|
|
|
|
|
def run_llm_with_quantization(prompts: List[str],
|
|
llama_model_dir: str,
|
|
engine_dump_dir: Optional[str] = None,
|
|
quant_type: str = 'int4_awq'):
|
|
''' Running LLM with quantization.
|
|
quant_type could be 'int4_awq' or 'fp8'.
|
|
'''
|
|
|
|
config = ModelConfig(model_dir=llama_model_dir)
|
|
if quant_type == 'int4_awq':
|
|
config.quant_config.init_from_description(quantize_weights=True,
|
|
use_int4_weights=True,
|
|
per_group=True)
|
|
else:
|
|
config.quant_config.set_fp8_qdq()
|
|
config.quant_config.set_fp8_kv_cache()
|
|
|
|
llm = LLM(config)
|
|
|
|
for output in llm.generate(prompts):
|
|
print(output)
|
|
|
|
|
|
def _parse_arguments():
|
|
parser = ArgumentParser()
|
|
parser.add_argument('--task', type=str, choices=_get_functions())
|
|
parser.add_argument('--hf_model_dir',
|
|
type=str,
|
|
help='The directory of the model.')
|
|
parser.add_argument('--dump_engine_dir',
|
|
type=str,
|
|
help='The directory to dump the engine.',
|
|
default=None)
|
|
parser.add_argument('--quant_type', type=str, choices=['int4_awq', 'fp8'])
|
|
parser.add_argument('--prompt', type=str, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def _get_functions():
|
|
cur_module = sys.modules[__name__]
|
|
function_names = [
|
|
name for name, _ in inspect.getmembers(cur_module, inspect.isfunction)
|
|
if not name.startswith('_')
|
|
]
|
|
return function_names
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = _parse_arguments()
|
|
|
|
tasks = dict(
|
|
run_llm_from_huggingface_model=lambda: run_llm_from_huggingface_model(
|
|
[args.prompt], args.hf_model_dir, args.dump_engine_dir),
|
|
run_llm_from_tllm_engine=lambda: run_llm_from_tllm_engine(
|
|
[args.prompt], args.dump_engine_dir),
|
|
run_llm_on_tensor_parallel=lambda: run_llm_on_tensor_parallel(
|
|
[args.prompt], args.hf_model_dir),
|
|
run_llm_generate_async_example=lambda: run_llm_generate_async_example(
|
|
[args.prompt], args.hf_model_dir),
|
|
run_llm_with_quantization=lambda: run_llm_with_quantization([
|
|
args.prompt
|
|
], args.hf_model_dir, args.dump_engine_dir, args.quant_type),
|
|
)
|
|
|
|
print(f'Running {args.task} ...')
|
|
|
|
tasks[args.task]()
|