Zenn
😽

TensorRT-LLMで推論してみる。

に公開
Python
LLM
tensorrtllm
tech

TensorRT-LLM

TensorRT-LLM の機能

pythonで実装されている。
LLMの推論最適化ライブラリ

  • TensorRT-LLMはTensorRT(C++)に依存している。

筆者の実行環境

databricks notebook

  • TensorRT-LLM version: 0.14.0
  • TensorRT-LLM version 0.14.0 が動くコンテナ

構築された環境を確かめる

import tensorrt_llm

推論用のエンジンファイル（.engine)が必要になる

HFからのモデル変換

HFモデル→ チェックポイント→エンジンファイルの順に変換していく。

https://github.com/NVIDIA/TensorRT-LLM/blob/v0.14.0/examples/llama/convert_checkpoint.py

チェックポイントの作成。↑はllamaベースの場合

エンジンファイルが含まれているフォルダにtokenizerのファイルもコピーして含めておく。
(tokenizer.json/tokenizer_config.json)

!python convert_checkpoint.py --model_dir '' \
                            --output_dir ''　\
                            --dtype float16 \
                            --tp_size 1

チェックポイントのファイルからビルド

!trtllm-build --checkpoint_dir '' \
            --output_dir '' \
            --gemm_plugin auto

推論の実行

runnerの読み込み

tokenizerはHFのものを使用している。

from transformers import AutoTokenizer
import tensorrt_llm
from tensorrt_llm.runtime import ModelRunnerCpp

engine_dir = "..." 
tokenizer = AutoTokenizer.from_pretrained(engine_dir, device_map="cuda")
runtime_rank = tensorrt_llm.mpi_rank()

runner = ModelRunnerCpp.from_dir(engine_dir=engine_dir, rank=runtime_rank)

推論の実行

import torch

def generate_trt(runner, tokenizer, input_text):
    EOS_TOKEN = tokenizer.eos_token_id
    input_tokens = [tokenizer.encode(input_text, add_special_tokens=False)]
    input_ids = torch.nested.to_padded_tensor(
        torch.nested.nested_tensor(input_tokens, dtype=torch.int32),
        EOS_TOKEN).cuda()
    outputs = runner.generate(
        input_ids,
        pad_id=tokenizer.pad_token_id,
        end_id=tokenizer.eos_token_id,
        max_new_tokens=256,
        temeprature=0.7,
        repetition_penalty=1.15,
        do_sample=True,
        return_dict=True,
        output_sequence_lengths=True,
        streaming=False)
    
    torch.cuda.synchronize()
    output_gen_ids = outputs["output_ids"]
    sequence_length = outputs["sequence_lengths"][0][0]
    output_text = tokenizer.decode(output_gen_ids[0][0][:sequence_length])
    return output_text

generate_trt(runner,tokenizer,"日本語で答えてください。ドラえもんについて教えて")

参考

https://developer.nvidia.com/ja-jp/blog/optimizing-inference-on-llms-with-tensorrt-llm-now-publicly-available/

Discussion