第一篇主要讲了怎么搭环境,以及如何使用TensorRT-LLM跑起来。而本篇主要是简析下trt-llm的内部原理。简单讲讲其内部架构以及一些周边库,旨在可以参考快速上手trt-llm。
话不多说,开始吧。
基本架构
TensorRT-LLM的前身是FasterTransformer,现在独立出来为一个git仓库:
官方推荐搭配Triton Inference Server去部署,有对应的backend:
TensorRT-LLM库(下文简称trt-llm)主要由这些组成:
- TensorRT(下文简称trt)
- FasterTransformer的一些设计和kernel
- NCCL(scalable inference,TP AllReduce PP send&Recv)
- cutlass、triton等其他组件
整体源码大部分开源,涉及硬件信息的kernel和依赖的tensorrt不开源。
以下带DataFun标记的图来自 揭秘NVIDIA大模型推理框架:TensorRT-LLM-人工智能-PHP中文网
因为现在LLM发展超级快,trt-llm的git仓库分为两个分支:Release branch和Dev branch。Release branch每月更新一次,而Dev branch会更频繁地更新来自官方或社区中的功能:
- https://github.com/NVIDIA/TensorRT-LLM/discussions 这里可以看最新的dev分支
因为trt-llm中很多op的实现都是来源于TensorRT(或者可以理解trt-llm是在trt上又包了一层),所以使用trt-llm前建议了解TensorRT的基本概念。
TensorRT基本概念
简单了解下trt这个推理加速库:
- trt可以加速模型推理,是模型推理框架,只能推理不能训练
- trt只支持单卡
- 类似于Pytorch,有自己的网络标准,提供了python-api去搭建网络
- 类似于编译器,输入我们的网络(比如onnx),输出优化好的模型(经过pass、fuse、kernel gen等步骤)
- 针对特定硬件优化(不同架构的显卡build出来的engine在trt-8.6版本之前不兼容)
trt如何构建网络
一般来说我们使用trt都是通过onnx → trt这条路经,trt官方提供了parser去解析onnx网络进一步构建为trt的engine。
在大模型这边则不用onnx去转trt,一是模型比较大,权重动不动就20G往上,onnx对大文件的支持不是很好;另一个是大模型的结构基本都比较相似,搭建修改的成本不是很高;而且大模型一般runtime都比较复杂(kv cache、paged attention),需要和模型op交互的地方比较多。
除了onnx这条路,据说trt-llm也尝试过torch-dynamo的trace的方式,只是因为某些原因放弃了
trt构建网络过程相比Pytorch构建稍微有点绕,比如:
# 摘自 tensorrtx/centernet/centernet.py
def add_batchnorm_2d(self, input_tensor, parent):
gamma = self.weights[parent + '.weight'].numpy()
beta = self.weights[parent + '.bias'].numpy()
mean = self.weights[parent + '.running_mean'].numpy()
var = self.weights[parent + '.running_var'].numpy()
eps = 1e-5
scale = gamma / np.sqrt(var + eps)
shift = beta - mean * gamma / np.sqrt(var + eps)
power = np.ones_like(scale)
return self.network.add_scale(input=input_tensor.get_output(0), mode=trt.ScaleMode.CHANNEL, shift=shift, scale=scale, power=power)
def add_basic_block(self, input_tensor, out_channels, residual=None, stride=1, dilation=1, parent=''):
conv1_w = self.weights[parent + '.conv1.weight'].numpy()
conv1 = self.network.add_convolution(input=input_tensor.get_output(
0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv1_w)
conv1.stride = (stride, stride)
conv1.padding = (dilation, dilation)
conv1.dilation = (dilation, dilation)
bn1 = self.add_batchnorm_2d(conv1, parent + '.bn1')
ac1 = self.network.add_activation(
input=bn1.get_output(0), type=trt.ActivationType.RELU)
conv2_w = self.weights[parent + '.conv2.weight'].numpy()
conv2 = self.network.add_convolution(input=ac1.get_output(
0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv2_w)
conv2.padding = (dilation, dilation)
conv2.dilation = (dilation, dilation)
out = self.add_batchnorm_2d(conv2, parent + '.bn2')
if residual is None:
out = self.network.add_elementwise(input_tensor.get_output(
0), out.get_output(0), trt.ElementWiseOperation.SUM)
else:
out = self.network.add_elementwise(residual.get_output(
0), out.get_output(0), trt.ElementWiseOperation.SUM)
return self.network.add_activation(input=out.get_output(0), type=trt.ActivationType.RELU)
可以看到:
- 所有带权重的层的权重需要手动放进去
- 某些常用的op需要自行组合(比如上述的bn层)
- 输入输出需要显式的指定
上述是搭建trt中INetworkDefinition
的过程,看起来没有Pytorch搭建网络那么直观,搭建完INetworkDefinition
后就可以利用trt的builder
去构建engine
,这一步trt做了很多优化。
来看下官方的流程介绍:
- Populate a
tensorrt.INetworkDefinition
either with a parser or by using the TensorRT Network API (seetensorrt.INetworkDefinition
for more details). Thetensorrt.Builder
can be used to generate an emptytensorrt.INetworkDefinition
. - Use the
tensorrt.Builder
to build atensorrt.ICudaEngine
using the populatedtensorrt.INetworkDefinition
. - Create a
tensorrt.IExecutionContext
from thetensorrt.ICudaEngine
and use it to perform optimized inference.
trt-llm 构建网络
trt-llm的网络构建中,是在trt的python-api上又包了一层,提升了易用性。直接使用trt-llm构建网络的感觉大概和pytorch一样。
首先将常用的函数用比较巧妙地方式包起来:
# In tensorrt_llm.functional:
def activation(input: Tensor, act_type: trt.ActivationType) -> Tensor:
layer = default_trtnet().add_activation(input.trt_tensor, act_type) # default_trtnet() -> INetworkDefinition
return _create_tensor(layer.get_output(0), layer)
然后利用partial
特性拓展出几个其他的激活op:
# In tensorrt_llm.functional:
relu = partial(activation, act_type=trt.ActivationType.RELU)
sigmoid = partial(activation, act_type=trt.ActivationType.SIGMOID)
对于进阶组合的激活op,比如silu,可以直接以最直观的方式组合起来:
# In tensorrt_llm.functional:
def silu(input: Tensor) -> Tensor:
return input * sigmoid(input)
上述的乘法操作(input * sigmoid(input)中的 *
)在组合的时候最终会跳转到trt-python-api中的elementwise_binary
操作,显得很自然:
class Tensor(object):
'''
The class to represent dense tensors.
A dense tensor is named, has a shape and contains typed elements. Each
dimension of a tensor can either be static or dynamic. Static dimensions
are known at engine compilation by TensorRT. Dynamic dimensions can take
values determined at runtime. The tensor can be located on the host (CPU)
or the device (GPU).
'''
...
def __mul__(self, b):
'''
See functional.mul.
'''
return mul(self, b)
mul = partial(elementwise_binary, op=trt.ElementWiseOperation.PROD)
对于常见的复杂op,是这么构建的:
class RmsNorm(Module):
def __init__(self,
normalized_shape,
eps=1e-06,
elementwise_affine=True,
dtype=None):
super().__init__()
if isinstance(normalized_shape, int):
normalized_shape = (normalized_shape, )
self.normalized_shape = tuple(normalized_shape)
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = Parameter(shape=self.normalized_shape, dtype=dtype)
else:
self.register_parameter('weight', None)
self.eps = eps
def forward(self, x):
weight = None if self.weight is None else self.weight.value
return rms_norm(x, self.normalized_shape, weight, self.eps)
def rms_norm(input: Tensor,
normalized_shape: Union[int, Tuple[int]],
weight: Optional[Tensor] = None,
eps: float = 1e-06) -> Tensor:
normalized_shape = [normalized_shape] if isinstance(
normalized_shape, int) else normalized_shape
dim = tuple([-i - 1 for i in range(len(normalized_shape))])
if default_net().strongly_typed:
input_dtype = input.dtype
fp32_input = cast(input, "float32")
varx = pow(fp32_input, 2.0)
varx = varx.mean(dim, keepdim=True)
denom = varx + eps
denom = denom.sqrt()
fp32_y = fp32_input / denom
y = cast(fp32_y, input_dtype)
else:
with precision("float32"):
varx = pow(input, 2.0)
varx = varx.mean(dim, keepdim=True)
denom = varx + eps
denom = denom.sqrt()
y = input / denom
if weight is not None:
y = y * weight
return y
整体流程是通过封装好的python-API,人工搭建出网络,搭建好的网络类型就是trt中的INetworkDefinition类,刚提到过的,由很多trt的layer组成。
@property
def trt_network(self) -> trt.INetworkDefinition:
return self._trt_network
When the TensorRT-LLM’s Python API is utilized, a graph of the network is assembled. The graph can later be traversed or transformed using the graph traversal API exposed by the
tensorrt.ILayer
class. That graph will also be optimized by TensorRT during the compilation of the engine, as explained in the next section.
trt-llm编译流程
类似于ONNX转trt,首先要转换weights,在上述构建好的network中进行weight的绑定,在将weights从pytorch模型中抽出来之后(会按照名字组成一个dict)再按照名称顺序一一填到network中的layer中:
# The Linear operator exposes two parameters (see tensorrt_llm/layers/linear.py):
class Linear(Module):
def __init__(self, ...):
self.weight = Parameter(shape=(self.out_features, self.in_features), dtype=dtype)
self.bias = Parameter(shape=(self.out_features, ), dtype=dtype)
# The parameters are bound to the weights before compiling the model. See examples/gpt/weight.py:
tensorrt_llm_gpt.layers[i].mlp.fc.weight.value = fromfile(...)
tensorrt_llm_gpt.layers[i].mlp.fc.bias.value = fromfile(...)
Once populated, the instance of the
tensorrt.INetworkDefinition
, can be compiled into an efficient engine by thetensorrt.Builder
In TensorRT-LLM, it is done through thebuild_engine
member function of thetensorrt_llm.Builder
class that calls thebuild_serialized_network
method of thetensorrt.Builder
object. That call, if everything works as expected, produces an instance of thetensorrt.IHostMemory
class. That object is an optimized TensorRT engine that can be stored as a binary file.
得到tensorrt.INetworkDefinition
后就可以开始build,流程和trt类似,最终生成优化好的engine,和TensorRT的engine的区别是,这个engine一般很大,plugin很多,支持多卡并行。
In-flight Batching && Paged KV Cache
In-flight Batching就是动态batch,我们CV中的图片都是静态的,比如(64,3,512,512)这种的,但是对于LLM来说,每次对话长度肯定不一样。
LLM 推理场景中,一个 batch 中每个 sample/request 的输出长度是无法预测的。如果按照静态batching的方法,一个batch的时延取决于 sample/request 中输出最长的那个,因此,虽然输出较短的 sample/request 已经结束,但是并未释放计算资源,其时延与输出最长的那个 sample/request 时延相同。
In-flight batching的做法是在已经结束的 sample/request 处插入新的 sample/request。这样,不但减少了单个 sample/request 的延时,避免了资源浪费问题,同时也提升了整个系统的吞吐量。
In-flight batching配合paged kv-cache使用,类似于vllm中的方式。
多卡
TensorRT-LLM支持Tensor Parallelism和Pipeline Parallelism。
Tensor Parallelism 通过在多个设备上垂直划分模型参数并同时处理各个切片的计算任务来实现并行性,这导致设备间需要进行频繁的数据交换,通常适用于具有高速互连技术(如NVLINK)的环境。
Pipeline Parallelism是采用了一种横向分割的策略,将模型按顺序切分为多个阶段,每个阶段在不同的设备上独立计算并将其结果传递给下一阶段,这种方式在设备间的通信带宽较低时表现更为有效,因为它仅在阶段间进行点对点的数据传递。
trt-llm中开启多卡功能通过plugin实现:
def allreduce(tensor: Tensor,
group: List[int],
strategy: Optional[AllReduceStrategy] = None) -> Tensor:
allreduce_plg_creator = trt.get_plugin_registry().get_plugin_creator(
'AllReduce', '1', TRT_LLM_PLUGIN_NAMESPACE)
if strategy is None:
if default_net().plugin_config.use_custom_all_reduce:
strategy = AllReduceStrategy.AUTO
else:
strategy = AllReduceStrategy.RING
counter = 0
workspace = None
if strategy != AllReduceStrategy.RING:
counter = current_all_reduce_helper().gen_id()
workspace = current_all_reduce_helper().workspace
assert allreduce_plg_creator is not None
group = trt.PluginField("group", np.array(group, dtype=np.int32),
trt.PluginFieldType.INT32)
p_dtype = default_net().plugin_config.nccl_plugin
pf_dtype = trt.PluginField(
"type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32),
trt.PluginFieldType.INT32)
pfc = [group, pf_dtype]
p_strategy = trt.PluginField("strategy", np.array([int(strategy)], np.int8),
trt.PluginFieldType.INT8)
pfc.append(p_strategy)
p_counter = trt.PluginField("counter", np.array([counter], np.int32),
trt.PluginFieldType.INT32)
pfc.append(p_counter)
pfc = trt.PluginFieldCollection(pfc)
ar_plug = allreduce_plg_creator.create_plugin("allreduce", pfc)
plug_inputs = [tensor.cast(p_dtype).trt_tensor]
if strategy != AllReduceStrategy.RING:
plug_inputs.append(workspace.trt_tensor)
layer = default_trtnet().add_plugin_v2(plug_inputs, ar_plug)
_add_plugin_info(layer, allreduce_plg_creator, "allreduce", pfc)
return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype)
对于tp多卡的情况,会把模型参数拆分成多个分给每个TensorRT engine,比如这里的embedding函数:
class Embedding(Module):
def __init__(self,
num_embeddings,
embedding_dim,
dtype=None,
tp_size=1,
tp_group=None):
super().__init__()
# num_embeddings records the total vocab size no matter using TP or not
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
self.tp_size = tp_size
self.tp_group = tp_group
# When TP are involved (tp_size>1),
# num_embeddings_tp is the size of the embedding numbers per process.
self.weight = Parameter(shape=(math.ceil(
self.num_embeddings / self.tp_size), self.embedding_dim),
dtype=dtype)
self.tp_size = tp_size
self.tp_group = tp_group
...
TensorRT-LLM中的kernel
trt-llm的kernel性能应该是目前最强的,除了TensorRT本身自带的CUDA kerne就比较强,trt-llm也实现了很多高性能的kernel,有着内部的黑科技加成,给你的就是cubin编译好的二进制文件:
第一个是FMHA(fused multi-head attention) kernel。由于 Transformer 中最为耗时的部分是 self-attention 的计算,因此trt-llm设计了 FMHA 来优化 self-attention 的计算,并提供了累加器分别为 fp16 和 fp32 不同的版本(针对不用的卡按需选择,某些卡是有阉割累加器的)。另外除了速度提升,内存的占用也大大降低。同时还提供了基于 flash attention 的实现,可以将 sequence-length 扩展到任意长度。
另外一个 Kernel则是MMHA(Masked Multi-Head Attention)。刚才的FMHA 主要用在 context phase 阶段的计算(prefill阶段),而 MMHA 主要提供 generation phase 阶段 attention 的加速,也就是decode阶段。性能比之前的FastTransformer更快。
llm在计算的时候,初始输入的prompt在prefill阶段计算,llm生成答案在decode阶段计算
在构建plugin的时候,通过控制num_kv_heads
就可以控制使用MQA/GQA,而最终plugin执行的code,就是上述提到的kernel。
template <typename T, typename KVCacheBuffer>
int GPTAttentionPluginCommon::enqueueContext(const EnqueueContextParams<T, KVCacheBuffer>& params, cudaStream_t stream)
{
const int num_heads = mNumHeads;
const int num_kv_heads = mNumKVHeads;
const int head_size = getHeadSize();
const int local_hidden_units_qo = num_heads * head_size;
const int local_hidden_units_kv = num_kv_heads * head_size;
const PositionEmbeddingType position_embedding_type = mPositionEmbeddingType;
const float q_scaling = mQScaling;
const bool* finished = nullptr;
const bool has_ia3 = false;
KVCacheBuffer kv_cache_buffer;
const auto elem_size = mKVCacheQuantMode.hasKvCacheQuant() ? sizeof(int8_t) : sizeof(T);
int64_t* host_kv_cache_block_ptrs = nullptr;
if (mPagedKVCache)
{
using BufferDataType = typename KVCacheBufferDataType<KVCacheBuffer>::Type;
kv_cache_buffer = KVCacheBuffer(params.batch_size, params.max_blocks_per_sequence, mTokensPerBlock,
num_kv_heads * head_size * elem_size, params.cyclic_attention_window_size, params.sink_token_length, false);
kv_cache_buffer.data = reinterpret_cast<BufferDataType*>(params.block_pointers);
host_kv_cache_block_ptrs = reinterpret_cast<int64_t*>(params.host_block_pointers);
}
...
if (params.sink_token_length > 0)
{
TLLM_LOG_ERROR("Cannot support StreamingLLM now when enabling paged KV context FMHA.");
}
mFMHARunner->setup_paged_kv(params.batch_size, params.input_seq_length, params.max_past_kv_len,
params.max_blocks_per_sequence, mTokensPerBlock, params.cyclic_attention_window_size, params.num_tokens,
isALiBi(), isAliBiWithScale(), mTpSize, mTpRank);
mFMHARunner->run_paged_kv(q_buf_2_, paged_kv_tma_desc, host_kv_cache_block_ptrs,
reinterpret_cast<KVBlockArray&>(kv_cache_buffer), cu_q_seqlens, cu_kv_seqlens, params.context_buf,
stream);
}
上述run_paged_kv
函数的实现在上述提到的contextFusedMultiHeadAttention
下的fmhaRunner
中定义,调用的就是cubin版本的kernel,黑科技。
trt-llm中的量化
量化是很常见提升模型性能的手段,小模型量化都没什么问题,大模型就更容易量化了,方式也更多。量化主要提升点是:
- INT8的算力比fp16和fp32的要高,可以充分利用tensor core的INT8算力
- 量化后的权重占用的显存小,可以开大batch,吞吐量就上来了
常用量化方式主要分为PTQ(训练后量化)和 QAT(训练感知量化),量化后拿到量化版的权重,之后推理方式都是一样的。
对于 LLM 量化技术,一个重要的特点是算法设计和工程实现的 co-design,即对应量化方法设计之初,就要考虑硬件的特性。否则,有可能达不到预期的推理速度提升。
TensorRT 中 PTQ 量化步骤一般分为如下几步:
- 首先对模型做量化,然后把权重和模型转化成 TensorRT-LLM 的形式
- 对于一些定制化的操作,需要自己编写 kernel
比如下述代码,在转换huggingface权重的时候同时做量化操作:
# /tensorrt_llm/models/llama/convert.py
@torch.no_grad()
def apply_smoothing(scales,
gemm_weights,
layernorm_weights=None,
layernorm_bias=None,
dtype=torch.float32,
layernorm_1p=False):
if not isinstance(gemm_weights, list):
gemm_weights = [gemm_weights]
if layernorm_weights is not None:
assert layernorm_weights.numel() == scales.numel()
layernorm_weights.div_(scales).to(dtype)
if layernorm_bias is not None:
assert layernorm_bias.numel() == scales.numel()
layernorm_bias.div_(scales).to(dtype)
if layernorm_1p:
layernorm_weights += (1 / scales) - 1
for gemm in gemm_weights:
gemm.mul_(scales.view(1, -1)).to(dtype)
常用的 PTQ 量化方法包括 INT8 weight-only、SmoothQuant、GPTQ 和 AWQ,这些方法都是典型的 co-design 的方法。
INT8 weight-only
INT8 weight-only直接把权重量化到 INT8,但是激活值还是保持为 FP16。好处就是模型存储2x减小,加载 weights 的存储带宽减半,同时可以开大batch,达到了提升推理性能的目的。
一般叫做W8A16,即权重为 INT8,激活值为 FP16/BF16——以 INT8 精度存储,以 FP16/BF16 格式计算。该方法直观,不改变 weights,容易实现,具有较好的泛化性能。
SmoothQuant
SmoothQuant观察到权重通常服从高斯分布,容易量化,但是激活值存在离群点,量化比特位利用不高。
- SmoothQuant 通过先对激活值做平滑操作即除以一个scale将对应分布进行压缩,同时为了保证等价性,需要对权重乘以相同的 scale。
- 权重和激活都可以量化
- 对应的存储和计算精度都可以是 INT8 或者 FP8,可以利用 INT8 或者 FP8 的 TensorCore 进行计算。
- 权重支持 Per-tensor 和 Per-channel 的量化,激活值支持 Per-tensor 和 Per-token 的量化。
GPTQ
一种逐层(layer-wise)量化的方法,通过最小化重构损失来实现。GPTQ 属于 weight-only 的方式,计算采用 FP16 的数据格式,适配W4A16。
该方法用在量化大模型时,由于量化本身开销就比较大,所以作者设计了一些 trick 来降低量化本身的开销,比如:
- Lazy batch-updates
- 以相同顺序量化所有行的权重
GPTQ 还可以与其他方法结合使用如 grouping 策略。针对不同的情况,trt-llm提供了不同的实现优化性能。比如 batch size 较小的情况,用cuda core实现;batch size 大的时候用 tensor core。
AWQ
该方法认为不是所有权重都是同等重要的,其中只有 0.1%-1% 的权重(salient weights)对模型精度贡献更大,并且这些权重取决于激活值分布而不是权重分布。该方法的量化过程类似于 SmoothQuant,差异主要在于 scale 是基于激活值分布计算得到的。
目前 TensorRT-LLM 提供了两类方法,即 FP8 和刚才提到的 INT4/INT8 量化方法。低精度如果 INT8 做 GEMM 时,累加器会采用高精度数据类型,如 fp16,甚至 fp32 以防止 overflow。
关于反量化,以 fp8 量化为例,TensorRT-LLM 优化计算图时,可能会自动移动反量化结点,合并到其它的操作中达到优化目的(QDQ合并,减少qdq的过程)。
对于前面介绍的 GPTQ 和 QAT,目前是通过硬编码写在 kernel 中,没有统一量化或反量化节点的处理。
和vllm的不同
- 添加新模型新pipeline(比如多模态)稍微麻烦些
- 需要提前确定输入尺度相关信息(batch、sequence_len,output_len等),毕竟涉及到构建trt的过程,有些参数一旦build好就不能变了,牺牲灵活性带来性能
- debug稍微繁琐一些,毕竟用了trt,需要显示指定output从而避免fuse带来的精度问题
运行
运行的整理流程:
TensorRT-LLM Backend
服务端的话,可以配合triton的backend来使用,官方已经提供了backend代码:tensorrtllm_backend/inflight_batcher_llm/src/libtensorrtllm.cc
,其实也就是用了trt-llm中batch_manager的接口,配合in-flight batching和kv-cache来实现高性能推理服务,可惜batch_manager部分是闭源的:
# from tensorrtllm_backend/inflight_batcher_llm/src/libtensorrtllm.cc
#include "tensorrt_llm/batch_manager/GptManager.h"
#include "tensorrt_llm/batch_manager/NamedTensor.h"
#include "tensorrt_llm/batch_manager/callbacks.h"
#include "tensorrt_llm/batch_manager/inferenceRequest.h"
#include "tensorrt_llm/common/logger.h"
#include "tensorrt_llm/plugins/api/tllmPlugin.h"
#include "tensorrt_llm/runtime/tllmLogger.h"
总结
一般来说,LLM推理主要看两个方面,一方面是极致的kernel性能(主要是降低latency),另一方是是极致的调度(主要提升throughput)。TensorRT-LLM的kernel性能大概是SOTA的,不过调度方面其他开源的框架说不定可以赶上或者超越。
可惜就是TensorRT-LLM的部分调度代码没有开源,用户没法改进和自定义。如果某些调度方式trt-llm官方没来的及支持(比如多模态llava),那么这个模型benchmark的时候估计就比较差(测试throughput的时候),比不过vllm或者其他框架。只能等trt-llm更新才行。
不过主要你的模型是比较常见的(llama、ChatGLM),trt-llm在kernel和调度上已经做很好了,性能肯定是比vllm或者其他框架强的,可以无脑使用。