An intermediate representation (IR) is the practical middle layer between a model’s development framework and its target hardware. Instead of writing a conversion for every framework–hardware pair, developers export models to formats such as TorchScript, MLIR, or ONNX, then use hardware-specific software to produce machine code. section 6.1.1
The IR approach enables:
- Reuse across GPUs, TPUs, and CPUs without implementing every framework-to-device combination. section 6.1.1
- Optimization before deployment, including graph optimization and hardware-specific kernel compilation. section 6.1.1
- Production workflows that export, trace, compile, and configure models for static or dynamic input sizes. section 6.1.1
Compiling a language model
A GPT-2 model can be loaded with PyTorch, traced using torch.jit.trace, and compiled with Torch-TensorRT. The compilation settings specify tensor data types and either a fixed input shape or minimum, optimal, and maximum dynamic shapes. section 6.1.1
Listing 6.1 Compiling a model with TensorRT
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch_tensorrt
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokens = tokenizer("The cat is on the table.", return_tensors="pt")[
"input_ids"
].cuda()
model = GPT2LMHeadModel.from_pretrained(
"gpt2", use_cache=False, return_dict=False, torchscript=True
).cuda()
model.eval()
traced_model = torch.jit.trace(model, tokens) #1
compile_settings = { #2
"inputs": [
torch_tensorrt.Input(
# For static size
shape=[1, 7],
# For dynamic sizing:
# min_shape=[1, 3],
# opt_shape=[1, 128],
# max_shape=[1, 1024],
dtype=torch.int32, # Datatype of input tensor.
# Allowed options torch.(float|half|int8|int32|bool)
)
],
"truncate_long_and_double": True,
"enabled_precisions": {torch.half}, #3
"ir": "torchscript",
}
trt_model = torch_tensorrt.compile(traced_model, **compile_settings)
torch.jit.save(trt_model, "trt_model.ts") #4
trt_model = torch.jit.load("trt_model.ts") #5
tokens.half()
tokens = tokens.type(torch.int)
logits = trt_model(tokens)
results = torch.softmax(logits[-1], dim=-1).argmax(dim=-1)
print(tokenizer.batch_decode(results))
Inspect a concrete TensorRT compilation workflow, including model tracing, input configuration, and deployment-oriented settings. section 6.1.1
