使用 Torch-TensorRT torch.compile 前端編譯 GPT2

此範例說明了使用 Torch-TensorRT 的 torch.compile 前端優化的最新模型 GPT2。在編譯之前安裝以下依賴項

pip install -r requirements.txt

GPT2 是一種因果 (單向) 轉換器，使用語言建模在非常大量的文字資料語料庫上進行預訓練。在此範例中，我們使用 HuggingFace 上提供的 GPT2 模型，並對其應用 torch.compile 以取得圖形模組的圖形表示。Torch-TensorRT 會將此圖形轉換為最佳化的 TensorRT 引擎。

匯入必要的程式庫

import torch
import torch_tensorrt
from transformers import AutoModelForCausalLM, AutoTokenizer

定義必要的參數

Torch-TensorRT 需要 GPU 才能成功編譯模型。 MAX_LENGTH 是產生的權杖可以具有的最大長度。這對應於輸入提示的長度 + 產生的新權杖數量

MAX_LENGTH = 32
DEVICE = torch.device("cuda:0")

模型定義

我們使用 AutoModelForCausalLM 類別從 Hugging Face 載入預訓練的 GPT2 模型。Torch-TRT 目前不支援 kv_cache，因此 use_cache=False

with torch.no_grad():
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    model = (
        AutoModelForCausalLM.from_pretrained(
            "gpt2",
            pad_token_id=tokenizer.eos_token_id,
            use_cache=False,
            attn_implementation="eager",
        )
        .eval()
        .cuda()
    )

PyTorch 推論

將範例輸入提示語進行 Tokenize (分詞)，並取得 PyTorch 模型的輸出

prompt = "I enjoy walking with my cute dog"
model_inputs = tokenizer(prompt, return_tensors="pt")
input_ids = model_inputs["input_ids"].cuda()

AutoModelForCausalLM 類別的 generate() API 用於使用貪婪解碼 (greedy decoding) 進行自迴歸生成 (auto-regressive generation)。

pyt_gen_tokens = model.generate(
    input_ids,
    max_length=MAX_LENGTH,
    use_cache=False,
    pad_token_id=tokenizer.eos_token_id,
)

Torch-TensorRT 編譯和推論

輸入序列長度是動態的，因此我們使用 torch._dynamo.mark_dynamic API 來標記它。我們提供此值的 (最小, 最大) 範圍，以便 TensorRT 預先知道要優化的值。通常，這將是模型的上下文長度。由於 0/1 specialization，我們從 min=2 開始。

torch._dynamo.mark_dynamic(input_ids, 1, min=2, max=1023)
model.forward = torch.compile(
    model.forward,
    backend="tensorrt",
    dynamic=None,
    options={
        "enabled_precisions": {torch.float32},
        "disable_tf32": True,
        "min_block_size": 1,
    },
)

使用 TensorRT 模型進行貪婪解碼的自迴歸生成迴圈。第一個 Token 生成會使用 TensorRT 編譯模型，而第二個 Token 則會遇到重新編譯 (目前這是一個問題，未來將會解決)。

trt_gen_tokens = model.generate(
    inputs=input_ids,
    max_length=MAX_LENGTH,
    use_cache=False,
    pad_token_id=tokenizer.eos_token_id,
)

解碼 PyTorch 和 TensorRT 的輸出句子

print(
    "Pytorch model generated text: ",
    tokenizer.decode(pyt_gen_tokens[0], skip_special_tokens=True),
)
print("=============================")
print(
    "TensorRT model generated text: ",
    tokenizer.decode(trt_gen_tokens[0], skip_special_tokens=True),
)

輸出句子應如下所示

"""
Pytorch model generated text:  I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with my dog. I'm not sure if I'll
=============================
TensorRT model generated text:  I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with my dog. I'm not sure if I'll
"""