4.1 图像预处理优化
在多模态推理中Vision Transformer (ViT) 是一个关键的模块,图像的预处理是将图像转换为适合ViT模型输入数据的过程。主要包括图像颜色空间转换、尺寸调整 (Resize)、划分图像块 (Patch Partitioning)、归一化(Normalize)等步骤。在LMDeploy框架中,图像预处理过程中主要通过PIL(Pillow)的Image模块在CPU上对图像进行处理,在图像Resize及Partition过程中,效率较低,耗时占整个ViT过程的20%以上。为了提升系统吞吐能力,减少图像预处理耗时,我们分别使用Pillow与OpenCV进行预处理测试,具体表现如下:
- CPU: Intel(R) Xeon(R) Silver 4410Y
使用OpenCV可以极大的减少图像预处理的耗时,平均处理单张图片的耗时由23.67ms减少到12.03ms,性能提升49.18%。在Resize过程中,虽然两个处理库对应的插值方式均使用BICUBIC,但当图像进行下采样时效果存在明显差异,使用OpenCV进行处理的图像存在波纹。如下:图2:Pillow与OpenCV效果对比通过对比源码实现,发现二者在插值与边界处理实现上有所差异:- 插值计算方式有差异:二者均使用4x4的卷积核进行插值计算,OpenCV直接使用三次多项式公式计算每个像素的权重,并对周围 16 个像素进行加权平均;而Pillow将三次卷积操作分解为两个一维卷积,先对水平方向进行卷积,然后再对垂直方向进行卷积。
- 边界处理的差异:OpenCV 供多种边界处理方式,例如 BORDER_REPLICATE, BORDER_REFLECT, BORDER_WRAP 等;Pillow通常使用边界复制的方式进行处理,即边缘像素值被复制到图像外部,以避免在边缘出现伪影。
针对这个问题,OpenCV说明文档中提供了相应的解决方案:
To shrink an image, it will generally look best with INTER_AREA interpolation, whereas to enlare an image, it will generally look best with INTER_CUBIC (slow) or INTER_LINEAR (faster but still looks OK).于是我们根据不同的图像采样对插值方式进行动态调整,对图像降采样时,使用INTER_AREA插值,上采样时,使用INTER_CUBIC(速度较慢,但效果最好),调整后,Resize结果如下:图3:OpenCV优化前后与Pillow效果对比4.2 ViT模块支持TensorRT
ViT模块是多模态推理框架中一个必不可少的组成模块,主要负责图像相关处理及编码工作。ViT模块的处理速度,直接影响整个框架的整体推理效率。为了进一步提升框架的推理效率,我们对ViT模块的耗时进行了分块分析,结果如下:
图4:vision 模型推理耗时及内存占用情况内存拷贝相关逻辑:
图5:LMdeploy VIT阶段内存拷贝代码截图经过验证,内存拷贝耗时主要是等待GPU异步处理结果,所以实际上主要耗时模块为图像预处理及特征提取两部分。具体定位步骤如下:
- lmdeploy/vl/engine.py 取消结果拷贝至cpu操作
- lmdeploy/serve/vl_async_engine.py 取消拷贝到cpu及转换numpy操作
- lmdeploy/pytorch/message.py中修改InputEmbeddings及类型为Torch.Tensor(GPU)
逻辑调整后,推理结果异常。在vl/engine.py forward增加输出结果日志后,推理正常。经验证输出结果日志操作起到同步等待作用,使用torch.cuda.synchronize()或者sleep验证猜想正确。后续在模型内增加日志输出结果或者以上两个操作,推理结果均正常。推理结果正常后定位耗时模块,定位到ViT中extract_feature为主要耗时模块。为了进一步提升推理效率,我们借鉴了TensorRT-LLM中的推理加速方案TensorRT。TensorRT是一个高性能的深度学习推理(Inference)优化器,可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可对多种应用场景进行推理加速,并且支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架。将TensorRT和NVIDIA的GPU结合起来,能在几乎所有的框架中进行快速和高效的部署推理。
在对ViT模块进行TensorRT改造时,主要包含模型转换、模型优化和推理部署三个阶段。模型转化支持 TensorFlow、PyTorch、ONNX 等主流深度学习框架的模型转换和优化,本文以ONNX为例进行说明。导出ONNX时可能会遇到不支持的算子,如在导出快速傅里叶变换(FFT)和快速傅里叶逆变换(IFFT)时会遇到如下错误,
Exporting the operator 'aten::fft_rfftn' to ONNX opset version 17 is not supported这时需要调整模型网络结构或者自定义算子。在对ViT模块进行ONNX转换过程中,部分多模态模型的ViT中使用了FlashAttention2进行注意力加速,而FlashAttention2中的flash_attn_func是作为独立的内核实现的,不是torch.nn.Module的实例,导致导出器无法捕获计算图,如下:/usr/local/lib/python3.10/dist-packages/flash_attn/flash_attn_interface.py:90: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! 因此,对Attention模块进行了调整,使用PyTorch内部实现的缩放点积注意力(Scaled Dot-Product Attention, SDPA),如下图,至此模型便可成功转换成ONNX格式。该阶段主要完成模型优化,如下图所示,在模型优化过程中会完成层间融合,精度校准等。这一步的输出是一个针对特定GPU平台和网络模型的优化过的TensorRT模型,这个TensorRT模型可以序列化存储到磁盘或内存中,存储到磁盘中的文件为TensorRT planfile。部署阶段将上一个步骤中的plan文件反序列化,并创建一个runtime engine,输入对应的图像数据,输出推理结果。经过TRT加速后,ViT模块feature_extract速度缩减45%左右(不包含图片预处理),feature_extract耗时在ViT中占比从60%减少至45.36%,整体推理耗时耗时缩减在70ms左右。4.3 ViT模块支持CudaGraph
推理框架lmdeploy在0.6.0版本引入了CUDA Graph,并提升了近30%的推理性能:
Employ CUDA graph to boost the inference performance (30%)不过受多方因素限制,目前lmdeploy只在语言模型中引入了CUDA Graphs。为了进一步提升推理速度,我们在ViT模块中引入了CUDA Graphs。CUDA Graphs可以用于优化执行过程中的CUDA操作,在GPU上实现更加高效的深度学习模型推理。在使用CUDA Graphs时需要对CUDA操作进行录制(capture)和重放(replay),以此来减少CPU到GPU的调度开销,提高整体的执行效率。如下图,简单展示了CUDA Graphs的优势。在顶部,CPU 逐个启动一系列短内核。CPU 启动开销导致内核之间出现明显间隙。如果我们用 CUDA 图替换此内核序列,最初我们需要花费一些额外的时间来构建图并在第一次启动整个图时一次性启动整个图,但后续执行将非常快,因为内核之间的间隙将非常小。当多次重复相同的操作序列时,例如在许多训练步骤中重复,差异会更加明显。
首先,在ViT支持CUDA Graphs时,需要torch.cuda.CUDAGraph创建对应的图,然后使用torch.cuda.graph()对ViT的推理过程进行录制,在推理过程中,使用刚创建的图对录制的过程进行重放CUDAGraph.play()。但是要注意,由于 CUDA Graphs 不支持动态控制流(如条件语句和循环),因此在设计算法时应尽量避免使用这些结构;其次,确保输入张量的形状在图创建时是固定的,因为 CUDA Graphs 的设计是基于静态形状的张量结构,创建 Graph 时,所有操作及其输入输出的形状必须在图创建时确定。
而ViT模块在进行图像处理时,输入的图像数张量的形状是 [batch_size, channel, width, height],其中batch_size是可变的且各视觉模型均已限定最大值。于是,我们在框架内部维护了Graphs Pool,推理时使用batch_size索引至相应的graph,再执行重放操作。
增加CUDA Graphs后ViT模块平均耗时减少30ms左右。虽然CUDA Graphs可以在一定程度上提升推理的效率,但是在构建graphs也需要占用一些额外的显存,在使用时需要综合衡量具体的业务场景及硬件资源。
4.4 图像Token化处理
输入token的长度对推理耗时影响很大,多模态模型中,图像部分占据了很大比例的token数,降低图像转换的Token数可提升推理性能。如下是结果对比:(1)根据图像宽高比和分辨率大小将原图拆分成若干个448*448的patch,拆分的原则是尽量保持图像不失真。拆分代码如下:图13:VLLM中InternVL2-8B模型拆图代码截图 上述代码基本流程是,给定动态拆分的阈值范围,穷举出所有可能的目标比例,再根据原图比例匹配最佳的拆分规则,拆分逻辑图示如下图左上部分,图示中会被拆分成6个path块和一张缩略图。(2)一个448*448的patch生成的token数计算方式如下:image_tokens_per_patch=(force_image_size // patch_size)**2 * (downsample_ratio**2))force_image_size=448,patch_size=14,downsample_ratio=0.5,这个计算后结果为256。不同的模型值可能会有所差异。(3)分辨率为896*1344的图像,经过步骤1处理,会拆分成2*3=6个patch,再加上一张缩略图(可选,有效果会更好),最终堆叠后shape是[7,3,448,448],图像转换的token数为7*256=1792。部署到线上时,单卡吞吐量上不去,其中一个原因是拆图规则导致拆分后的图片数量比较多,如分辨率612*464,最合适的宽高比是 (4, 3),按模型的图片拆分规则,图像将被拆分成[13,3,448,448],转化后的token数达到3328,再加上prompt的token,总token数会达到3400+,太长的输入token对模型推理速度影响很大,再加上显存和算力的限制,无法做到更大batch的推理,使得单卡推理的吞吐量很低。基于此原因,我们的优化思路是降低图像的总token数,经实验分析,官方代码在实现上存在比较大的冗余设计,如图像分辨率为480*360,也会转换成3328个token数,对于低分辨率图像生成太多的token存在资源浪费。在保持图像内容不拉伸前提下,对图像的宽高比做调整,以适应vit的要求,优化后,480*320的图像只转换成512个token数,这样在推理时能做到更大的batch处理。在我们实际落地场景中,处理后吞吐量能提升1倍。4.5 prefixcache在多模态模型里应用
在PagedAttention中,KV Cache只是在一个请求内复用,而没有做到跨请求的KV Cache复用。长prompt的场景,prompt在不同的请求中是相同的,KV Cache的计算也是相同的,如果能把prompt的KV Cache保存下来,留给后续的请求复用,将会极大地降低首Token的耗时。在LLM模型里,prefixcache分二个阶段,第一个阶段,当prompt第一次被推理时,是按block_size(通常是64)大小对input tokens从前往后进行分块,计算每个分块的hash作为唯一标识,每个分块的token_id作为key进行缓存,这里不足block_size长度的块不会被缓存;第二阶段,当新prompt被推理时,会进行prefix cache matching,命中就直接复用kvcache,只计算未命中部分的input tokens。多模态模型区别在于,一次任务的输入tokens组成由纯文本变成了文本+图片,由system+prompt变成了system+image+prompt,在计算prefix cache时,image对应的只是 padding tokens,那么在计算prefix cache matching时,不同图片可能匹配到一样的 prefix 上,这样推理结果就会出现错误。针对这个问题,在input tokens中对image进行范围标记,在计算prefix cache时不对image token进行kvcache,只cache image之前的部分;在prefix cache matching时,也同样保证image token不会被复用。经实验验证,修改后能保证在开启prefix cache时,推理结果是正确的。需要注意,Prefix Caching只节省了prefill阶段的耗时(也就是降低了TTFT,Time To First Token),并不能节省解码阶段的耗时(也就是TPOT,Time Per Output Token)。如果请求的主要耗时是在解码阶段(例如prompt很短而completion很长),或者多个请求的prompt并没有公共的前缀,那么Prefix Caching就对于整个LLM推理的性能提升帮助不大。4.6 模型量化
量化是大模型领域中的一项关键技术,它通过降低模型参数的精度,将浮点数转换为整数或定点数从而实现模型的压缩和优化。模型量化可以减少模型尺寸,进而减少在推理时的显存消耗,并且在一些低精度运算较快的处理器上可以增加推理速度。量化分很多情况。从量化对象来说,量化可以是权重、激活、kv cache和梯度;从量化的形式上来说分为线性量化和非线性量化,其中线性量化又分为对称量化和非对称量化;根据应用量化压缩模型的阶段,又可以将模型量化分为量化感知训练、量化感知微调、训练后量化。我们现阶段使用的量化方式是AWQ和GPTQ,这两种量化都属于训练后量化,是针对权重的线性量化,其中AWQ采用对称量化,GPTQ采用非对称量化。AWQ量化的原理是对于LLM,权重不是同等重要的,通过保留1%的显著权重可以大大减少量化误差。在此基础上采用激活感知的方法,考虑更大的激活幅度应该对应更重要的权重通道,在处理重要特征时起关键作用,逐通道确定最佳缩放因子。从而在量化所有权重的同时,最小化量化误差。GPTQ对模型的每一层(通常是线性层或卷积层)进行单独处理,考虑了量化带来的误差,并通过调整未量化的权重来补偿这些误差。利用了二阶偏导Hessian矩阵的逆,来指导权重的调整,以减少整体的量化误差。将权重矩阵分成多个子矩阵(block),对每个子矩阵中的权重逐个进行量化,同时调整同一子矩阵内其他权重,以保持模型输出的相似性。其量化后的误差依赖一份高质量的校准数据。整体上来看,AWQ相较于GPTQ量化的算法更直接,对校准数据依赖小;GPTQ则更容易有比较好的量化效果,但是算法相对复杂,对校准数据依赖比较大,实际过程中用哪个更合适需要根据实际的场景选用。在实际测试中,不论是AWQ还是GPTQ实际采用的都是w4A16的量化策略,在推理的时候,性能差异比较小,在RTX4090显卡下,我们使用vllm,对应不同参数,并且设置最优batch,实际测试值如下:
从测试结果看:在4090下,大batch的计算,使用gemm内核,速度不如原精度,原因是在大batch的情况下,增加了反量化的时间。使用marlin内核,计算的速度有优化,但是在大batch下,优化速度不明显。低batch的计算原精度是计算最慢的,gemm的内核计算与marlin计算差别不是很大,都比原生的有大幅提高。原因是gemm在低batch下,也做了内核优化,这一点可以从原代码中验证:图16:VLLM中awq量化模型mul计算逻辑代码