趋近智
优化目标硬件上的LLM推理性能需要专门推理运行时的实践应用。这涉及硬件映射、内存管理、优化过的核函数以及编译器技术等概念。理论模型优化(如量化或剪枝)只有与能够在硬件上高效执行优化后模型图的运行时配合时,才能充分发挥其潜力。我们将考察ONNX Runtime、NVIDIA TensorRT和vLLM等运行时如何将优化后的模型转化为更快的执行。
我们的目标是获取一个预训练Transformer模型,并使用不同运行时环境衡量其推理性能,同时与标准深度学习框架(如PyTorch)中的基准实现进行比较。我们将侧重于GPU加速,这是LLM部署的常见情况。
假设我们有一个中等大小的Transformer模型(例如GPT-2的蒸馏版本或类似架构,可能已使用第二章中的方法量化为INT8),并以PyTorch格式提供。我们的第一步是设定一个基准性能指标。这通常涉及在PyTorch中加载模型,并在目标GPU上对代表性输入数据批次运行推理,仔细测量延迟或吞吐量。
# 占位符:PyTorch基准推理测量
import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
# 假设模型和分词器已加载并位于GPU上
# model = AutoModelForCausalLM.from_pretrained(...)
# tokenizer = AutoTokenizer.from_pretrained(...)
# model.to('cuda')
# model.eval()
# 示例输入
# inputs = tokenizer("Example input text", return_tensors="pt").to('cuda')
# input_ids = inputs['input_ids']
# --- 基准测量 ---
num_runs = 50
warmup_runs = 5
latencies = []
with torch.no_grad():
# Warmup
for _ in range(warmup_runs):
_ = model.generate(input_ids, max_new_tokens=50) # Or model(**inputs) for classification
# Timed runs
for _ in range(num_runs):
torch.cuda.synchronize() # 确保GPU上计时准确
start_time = time.perf_counter()
_ = model.generate(input_ids, max_new_tokens=50) # Or model(**inputs)
torch.cuda.synchronize()
end_time = time.perf_counter()
latencies.append((end_time - start_time) * 1000) # 毫秒
baseline_avg_latency = sum(latencies) / num_runs
print(f"PyTorch基准平均延迟: {baseline_avg_latency:.2f} 毫秒")
这个基准提供了一个参考点,我们可以据此评估专用运行时带来的改进。
ONNX Runtime (ORT) 是一个跨平台推理和训练加速器,它通过开放神经网络交换(ONNX)格式支持来自各种框架(PyTorch、TensorFlow、scikit-learn等)的模型。ORT应用图优化并使用称为执行提供程序(EPs)的硬件专用加速库。
模型转换: 第一步是将我们的PyTorch模型导出为ONNX格式。torch.onnx.export函数促进了这一点。我们需要提供一个示例输入来追踪模型的执行图。
# 占位符:将PyTorch模型导出到ONNX
# 假设已定义 'model' 和 'dummy_input'(匹配预期输入形状)
# dummy_input = tokenizer("Example input text", return_tensors="pt")['input_ids'].to('cuda')
torch.onnx.export(model, # 运行的模型
dummy_input, # 模型输入(或多个输入的元组)
"model.onnx", # 模型保存路径
export_params=True, # 将训练好的参数权重存储在模型文件中
opset_version=14, # 导出模型到ONNX的版本
do_constant_folding=True, # 执行常量折叠以进行优化
input_names = ['input_ids'], # 指定输入名称
output_names = ['output'], # 指定输出名称
# dynamic_axes={'input_ids' : {0 : 'batch_size', 1: 'sequence'}, # 使轴动态
# 'output' : {0 : 'batch_size'}})
print("模型已导出到 model.onnx")
注意: 处理动态轴对于可变输入大小(批次、序列长度)很重要,但这会增加复杂性。为简单起见,我们可能会从固定形状开始。
使用ONNX Runtime进行推理: 现在我们可以加载ONNX模型并使用ORT运行推理,指定所需的执行提供程序。对于GPU加速,我们使用CUDAExecutionProvider。
# 占位符:ONNX Runtime推理
import onnxruntime as ort
import numpy as np
# 加载ONNX模型
sess_options = ort.SessionOptions()
# 可选:启用图优化
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
ort_session = ort.InferenceSession("model.onnx", sess_options, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
# 准备输入(将PyTorch张量转换为NumPy)
# numpy_input = inputs['input_ids'].cpu().numpy()
ort_inputs = {ort_session.get_inputs()[0].name: numpy_input}
# --- ONNX Runtime测量 ---
ort_latencies = []
# Warmup
for _ in range(warmup_runs):
_ = ort_session.run(None, ort_inputs)
# Timed runs
for _ in range(num_runs):
start_time = time.perf_counter()
_ = ort_session.run(None, ort_inputs) # ORT通常隐式处理GPU同步
end_time = time.perf_counter()
ort_latencies.append((end_time - start_time) * 1000)
ort_avg_latency = sum(ort_latencies) / num_runs
print(f"ONNX Runtime (CUDA EP) 平均延迟: {ort_avg_latency:.2f} 毫秒")
ORT自动应用各种图优化(算子融合、常量折叠)并使用CUDA EP提供的优化过的CUDA核函数,通常会带来相比PyTorch基准实现显著的加速。如果导出的模型已经量化,ORT可以使用专门的INT8核函数获得进一步的提升。
NVIDIA TensorRT是一款专门用于NVIDIA GPU的高性能深度学习推理优化器和运行时。它执行积极的图优化、层融合、核函数自动调优和精度校准(FP32、FP16、INT8),以最大限度地使用底层硬件,包括Tensor Cores。
构建TensorRT引擎: TensorRT通常接收ONNX模型(或与TensorFlow/PyTorch等框架集成)。核心步骤是构建一个TensorRT“引擎”,这是模型的某个高度优化版本,特定于目标GPU和TensorRT配置(例如,精度模式)。这个构建过程可能需要时间,因为TensorRT会检查最优核函数和配置。
# 占位符:构建TensorRT引擎(通常通过'trtexec'命令行工具或Python API完成)
# 使用trtexec(命令行示例):
# trtexec --onnx=model.onnx --saveEngine=model.engine --fp16 --workspace=4096
import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()
parser = trt.OnnxParser(network, TRT_LOGGER)
# 解析ONNX模型
with open("model.onnx", "rb") as model_file:
if not parser.parse(model_file.read()):
print("ERROR: Failed to parse the ONNX file.")
for error in range(parser.num_errors):
print(parser.get_error(error))
# 适当处理错误
print("ONNX文件解析完成")
# 配置构建器(例如,启用FP16、INT8、设置工作区)
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * (1024**3)) # 4GB工作区
if builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
print("FP16已启用")
# 构建序列化引擎
# plan = builder.build_serialized_network(network, config)
# if not plan:
# print("ERROR: Failed to build engine.")
# # 处理错误
#
# with open("model.engine", "wb") as f:
# f.write(plan)
# print("TensorRT引擎已保存到 model.engine")
```
*注意:* TensorRT引擎构建需要仔细处理动态形状、INT8精度校准和足够的工作区内存。
2. TensorRT推理: 引擎构建完成后,推理涉及加载并使用输入数据执行它。这可以使用TensorRT Python/C++ API完成,或者通常更方便地通过ONNX Runtime的TensorRT执行提供程序进行,后者在后台处理引擎管理。
```python
# 占位符:使用ONNX Runtime与TensorRT EP进行推理
import onnxruntime as ort
import numpy as np
# 确保环境中可访问TensorRT库
ort_session_trt = ort.InferenceSession("model.onnx", # 从ONNX加载,ORT构建/缓存TRT引擎
providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'],
provider_options=[
{'device_id': 0,
'trt_fp16_enable': True, # 启用FP16精度
'trt_engine_cache_enable': True,
'trt_engine_cache_path': './trt_cache'}, # 缓存目录
{}, {} # CUDA EP和CPU EP的选项
])
# 准备输入(如前所述)
# ort_inputs = {ort_session_trt.get_inputs()[0].name: numpy_input}
# --- TensorRT测量(通过ORT)---
trt_latencies = []
# Warmup
for _ in range(warmup_runs):
_ = ort_session_trt.run(None, ort_inputs)
# Timed runs
for _ in range(num_runs):
start_time = time.perf_counter()
_ = ort_session_trt.run(None, ort_inputs)
end_time = time.perf_counter()
trt_latencies.append((end_time - start_time) * 1000)
trt_avg_latency = sum(trt_latencies) / num_runs
print(f"ONNX Runtime (TensorRT EP) 平均延迟: {trt_avg_latency:.2f} 毫秒")
```
TensorRT通常在NVIDIA GPU上提供最高的性能,这得益于其硬件专用优化,特别是在使用FP16或INT8等降低精度时。
对于文本生成任务,标准批处理方法可能效率低下,因为批次中的序列在不同时间完成,导致对填充令牌的计算浪费。vLLM是一个专门为解决LLM推理中此问题而设计的库。其核心创新是PagedAttention,它能更高效地管理注意力键和值的内存,类似于操作系统中虚拟内存和分页的工作方式。这使得内存使用接近最优,减少了碎片,并实现了连续批处理,即新序列在其他序列完成后立即添加到批次中。
设置与使用: vLLM提供一个高级Python API,通常允许直接从Hugging Face Hub等来源加载模型。
# 占位符:用于文本生成的vLLM推理
# 确保vllm已安装:pip install vllm
from vllm import LLM, SamplingParams
# 加载模型(vLLM处理优化加载)
# 使用与基准相同的模型标识符
# llm = LLM(model="gpt2", # 替换为您的模型
# tensor_parallel_size=1) # 根据可用GPU调整
# 定义采样参数
sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=50)
# 准备提示
prompts = ["Example input text", "Another prompt for batching"]
# --- vLLM测量 ---
vllm_latencies = [] # 测量每个请求的延迟或批处理的总时间
# 预热可能由内部处理或需要几次初始运行
start_time = time.perf_counter()
outputs = llm.generate(prompts, sampling_params)
end_time = time.perf_counter()
total_time_ms = (end_time - start_time) * 1000
avg_latency_per_seq = total_time_ms / len(prompts)
# 吞吐量计算将取决于令牌数量和时间
print(f"vLLM每个序列的平均延迟: {avg_latency_per_seq:.2f} 毫秒")
# 为了公平比较,测量吞吐量(令牌/秒)
# total_output_tokens = sum(len(output.outputs[0].token_ids) for output in outputs)
# throughput_tok_sec = total_output_tokens / (total_time_ms / 1000)
# print(f"vLLM吞吐量: {throughput_tok_sec:.2f} 令牌/秒")
# 打印生成文本(可选)
# for output in outputs:
# prompt = output.prompt
# generated_text = output.outputs[0].text
# print(f"提示: {prompt!r}, 生成文本: {generated_text!r}")
vLLM在生成任务的吞吐量方面表现出色,特别是在可变输出长度时,这归因于PagedAttention和连续批处理最大限度地减少了GPU空闲时间。
准确的基准测试是根本。总是在测量前执行预热运行,运行多次迭代以平均结果,并测量相关指标(例如,p50、p90、p99等不同百分位数的延迟,以及每秒令牌数或每秒请求数的吞吐量)。如果构建服务系统,请确保在真实的负载条件下进行测量。
让我们以可视化方式比较这些运行时产生的结果:
文本生成任务在不同运行时之间的推理性能比较。更低的延迟和更高的吞吐量表示更好的性能。vLLM通常在生成方面表现出显著的吞吐量提升,这归因于优化过的内存管理和批处理。
分析:
从理论模型压缩转向,使用专用推理运行时对于部署高性能LLMs是必不可少的。ONNX Runtime、TensorRT和vLLM等工具实现了本章讨论的许多硬件感知优化,例如核函数融合、优化过的内存管理(如PagedAttention)、高效硬件映射以及对降低精度的支持。通过转换模型并在这些优化环境中执行它们,可以实现延迟的显著降低和吞吐量的显著增加。这一步弥合了压缩模型文件与高效可部署的AI应用之间的差距。尝试使用这些运行时及其配置选项是LLM部署工作流程的标准部分。
这部分内容有帮助吗?
© 2026 ApX Machine Learning用心打造