投递人 itwriter 发布于 2022-04-27 09:38 [收藏] « »

  作者:王申领 

  MLPerf 是一套衡量机器学习系统性能的权威标准,将在标准目标下训练或推理机器学习模型的时间,作为一套系统性能的测量标准。MLPerf 推理任务包括图像识别(ResNet50)、医学影像分割(3D-UNet)、目标物体检测(SSD-ResNet34)、语音识别(RNN-T)、自然语言理解(BERT)以及智能推荐(DLRM)。在 MLPerf V2.0 推理竞赛中,浪潮 AI 服务器基于 ImageNet 数据集在离线场景中运行 Resnet50,达到了 449,856 samples/s的计算性能,位居世界第一。本文将介绍浪潮在 MLPerf 推理竞赛中使用的卷积合并计算算法。

  Resnet 是残差网络(Residual Network)的缩写,该系列网络广泛用于目标分类等领域以及作为计算机视觉任务主干经典神经网络的一部分,典型的网络有 Resnet50、Resnet101 等。在 Resnet 神经网络中,主要计算算子是卷积计算层。Resnet50 神经网络具有 4 组残差结构,这 4 组残差结构包含 48 个卷积算子,通过设计卷积算子的计算算法,提高卷积算子的计算性能,可以减少 Resnet50 推理过程中的延迟。基于 A100 GPU 单卡的性能测试显示,在 BatchSize=2048 的情况下,优化后的卷积合并优化算法相比原算法可带来 14.6% 的性能提升。

  1. MLPerf Resnet50推理流程

  在 MLPerf V2.0 推理测试中,Resnet50 模型需要在 ImageNet2012 测试集上达到 FP32 精度(76.46%)的 99% 以上。数据中心赛道设置了离线(Offline)与在线(Server)两种模式,其中离线模式会产生一次推理时间大于 10 分钟的 samples 请求,可直接反映机器和算法的推理性能。

  Resnet50 推理流程如下。首先在 ImageNet2012 测试集中读取数据,并进行数据预处理,随后数据会加载到 TensorRT 中进行实际的推理测试。测试分为两方面,一是测试模型的精度;二是产生一次推理请求,TensorRT 会将请求中的图片全部推理完成得到总时间,根据计算时间得到每秒推理的样本数量,即为最终的成绩。

  2. 卷积合并计算算法

  2. 1 算法优化思路

  在 GPU 上运行 Resnet50 图像推理模型时,需要将每一个算子(卷积、池化、全连接等)放在 GPU 的 Kernel 中进行算子计算,由于 GPU 上运行 Kernel 时共享内存以及寄存器的资源有限(A100 的共享内存为 164KB),不可能将所有的计算过程数据放到 Kernel 中,而 GPU 的全局内存(A100 有 40G 或者 80G 全局内存)一般都很大,所以会将比较大的过程数据放在全局内存中。在进行推理时,根据 Kernel 的计算将数据按需从全局内存读取到 Kernel 中进行计算,每个算子在计算时会不可避免地产生 Kernel 与全局内存的数据交换,由于全局内存的读写访问延迟较大,会使算子计算性能下降。

  对于每个算子的 Kernel 计算,会产生两部分的全局内存访问,一部分是最开始的全局内存读取,另一部分是 Kernel 计算完成后的全局内存写回。为了降低全局内存访问带来的性能影响,有如下两种办法:

  一是采用算子合并的方式。默认的程序会将每个算子都放在单独的 Kernel 中进行计算,每个算子都会产生全局内存读和写两次访问。如果将两个算子放在一个 Kernel 中进行计算,对于连续的两个卷积计算,可以减少第一个卷积算子的写回以及第二个算子的读取;对于卷积与 Shortcut 的合并,可以减少一次全局内存的写回操作,通过减少全局内存的访问可以提高程序的计算性能。

  二是根据 GPU 不同架构的计算特性对 Kernel 的内部计算进行合理的优化设计。当不可避免地需要对全局内存进行访问时,做到全局内存进行连续线程的融合读取,充分利用向量化读取等加速对全局内存的访问,同时优化计算流程,通过 Double buffer 用计算来隐藏内存的访问延迟,对于需求较晚的全局内存数据,也可以通过 A100 的新特性-全局内存的异步复制来隐藏数据读取过程。

  本文主要针对 MLPerf 推理中 Resnet50 卷积神经网络的第二组残差结构中的部分算子进行计算合并,在充分考虑 GPU 计算特性前提下,进行合理的算法设计,提高 Resnet50 卷积神经网络的性能。

  2. 2 Resnet50合并计算算法

  在 Resnet50 神经网络中,第二组残差结构有 Res3.1、Res3.2、Res3.3、Res3.4,共四部分的卷积计算,其中 Res3.2、Res3.3、Res3.4 三部分计算结构一样,如下图所示:

  可以看到,Res3.1 的输出(input)作为 Res3.2 部分的输入,输入后会有两部分分支,在右部分的分支中,会先后计算 Conv1,Conv2,Conv3 三个卷积,其中 Conv1,Conv2 两个卷积后面都包含 Bn 和 Relu 过程,Conv3 后面会有 Bn 的计算过程;在右边分支计算完成后,会与 input 进行 Shortcut 操作,主要进行的是与输入数据 Sum 和 Relu 操作,两部分结果经过 Shortcut 操作后会得到 Res3.2 的输出完成这部分的计算。

  本文介绍的合并算法对图 2 虚线框中的计算进行合并,主要是对 Conv3 以及 Shortcut 的过程进行合并,包含 Conv3+Bn+Sum+Relu 过程。

  3. 卷积合并算法在GPU加速卡上的实现

  Res3.2 的计算参数主要如下:

  通过上表可以看到,Conv3 输入 Data 的H*W为 28*28,输入通道 Ic 为 128,输入的权重 Weight 的h*w为1*1,输入通道 Ic 为 128,输出通道 Oc 为 512;Shortcut 的输入同 Conv1,其中H*W为 28*28,输入通道 Ic 为 512;两部分计算合并之后的输出 Output 的H*W为 28*28,通道 Oc 为 512。(本文所有的算法都围绕 A100 GPU 进行介绍)

  3. 1关于dataweightoutputlayout变换

  本文采用的计算数据类型为 int8,因此下文介绍的所有内容都是基于 int8 开展的优化。

  算法对 data 以及 weight 进行了提前处理以适应 GPU 的计算特性,主要处理如下:

  对于 data,原始 layout 为[B, H, W, Ic]=[B, 28, 28, 128],算法将 Ic=128 以 32 为单位进行拆分为 4 组,形成[B, 4, H, W, Ic/32]=[B, 4(I1), 28, 28, 32(I2)]的 layout,这样做的目的是 32 个 int8 可以组成 16B 共 128 位数据的联合向量化读写,提高 GPU 中全局内存的通信速度。

  对于 weight,由于h*w=1*1,因此本文后续不再表示h*w,默认的 weight 的 layout 为[Ic, Oc]=[128, 512],算法将 Ic 以 32 为单位进行拆分为 4 组,将 4 放在左数第二维,将 32 放在左数第四维,这样做的目的也是为了程序在访问全局内存时做到 16B 共 128 位数据的联合向量化读写;算法将 Oc 以 128 为单位进行拆分为 4 组,将 4 放在左数第一维,将 128 放在左数第三维,这样做的目的是将 Oc 拆成了 4 组放在了不同的 block 中进行计算,这样在每个 block 进行计算的时候可以顺序的由全局内存加载 weight,不会产生数据内存位置的跳跃,这部分会在后面 block 的划分中进行介绍,这样就形成[O1, I1, O2, I2] =[4, 4, 128, 32]的 weight 的 layout。

  对于 output,原始 layout 为[B, H, W, Oc]=[B, 28, 28, 512],这部分数据类似于输入 data,将 Oc 以 32 为单位进行拆分为 16 组,形成 layout 为[B, 16, H, W, Oc/32]=[B, 16(O1), 28, 28, 32(O2)]。

  3. 2关于Grid以及Block的并行划分:

  对于 Grid 的划分,首先是x维度,由上文可知,对于 Conv3 的 Oc 为 512,本文将 Oc 划分为 4 组放到 Grid.x 维度,每组计算的 Oc 为 128;对于y维度,将H*W=28*28=784 分为 49 组放到 Grid.y 维度,每组计算的 HW 为 16;对于z维度,将B分为B/4 组放到 Grid.z 维度,每组计算B的数量是4。这样经过划分,Grid 的数量为[Grid.x, Grid.y, Grid.z]=[4, 49, B/4],即共有4*49*B/4 组计算同时并行进行。A100 GPU 上 SM 数量有 108 个,当B≥4 时,一个 kernel 共需要启动大于4*49*4/4=196 个 Block,完全满足 Grid 维度并行度的要求。

  对于 Block 中的划分,A100 GPU 中一个 SM 的 warp schedule 为 4 个,因此一个 block 中线程数量至少大于或等于 128,为了实现更好的并行度,算法选择一个 Block 中设置 8 个 warp 共 256 个线程。

  3. 3关于Block内部计算层次划分

  由上文可知,一个 Block 中划分的 Output 的计算 shape 为[B, H*W, Oc]=[4, 16, 128],由于在1*1 卷积计算中B维度与 HW 维度具有同等地位,因此将 BHW 合并为一个维度,此时本文用M表示 BHW 维度,即M=BHW=4*16=64,用N表示 Oc 维度,即N=Oc=128,此时一个 Block 中的计算维度变为[M,N]=[64, 128]。

  由前述 data 的 layout 的变换可知形成[B, 4, H, W, Ic/32]=[B, 4, 28, 28, 32]的 layout,由于 Grid.y 以及 Grid.x 对B维度以及 HW 维度进行了划分,此时一个 Block 中 data 的输入数据为[B, I1, HW, I2]=[4, 4, 16, 32],用上段所述 BHW 合并为 1 维用M表示,即M=B*HW=4*16=64,K表示 I1*I2 维度,即K=I1*I2=128,则此时一个 Block 中 data 的计算维度变为[M, K]=[64, 128]。

  由前述 weight 的 layout 的变换可知,weight 的 layout 为[O1, I1, O2, I2]=[4, 4, 128, 32],由于对 Oc 按照 32 为单位划分 4 组在 Grid.x 维度,因此每个 Block 中计算的 Oc 为 128,此时一个 Block 中的 weight 的计算数据为[I1, O2, I2] =[4, 128, 32],用N表示 O2 维度,即N=O2=128,用K表示 I1*I2 维度,即K=I1*I2=128,则此时一个 Block 中 weight 的计算维度变为[N, K]=[128, 128]。

  一个 Block 中实际要进行的计算就变为一个矩阵乘 data[M, K]点乘 weight[N, K]等于 output[M, N],即[64, 128]﹒[128,128]=[64, 128],共4*49*B/4 个 Block 并行完成所有整个卷积合并的计算,其中 data 的实际维度为[B, I1, HW, I2]=[4, 4, 16, 32],weight 的实际维度为[I1, O, I2]=[4, 128, 32]。

  经过前面的划分,一个 Thread Block 层次实际计算量为[64, 128]﹒[128,128]=[64, 128]。

  为了加速 int8 矩阵乘的计算,程序采用了 CUDA 中 mma 进行加速计算,其中 mma 的计算形状为[m ,n ,k]=[16, 8, 32],为了配合共享内存,寄存器以及 mma 形状的匹配,程序将内积方向的K维度 128 拆分为 2 组 64 进行计算,每组 64 进一步拆分为 2 组 32(k)进行计算,这样最基础的 Thread Block 层次进行的计算就变为图 3 中左上角虚线框中所示的[M, k]﹒[N, k]=[M ,N]即[64, 32]﹒[128, 32]=[64 ,128],由于一个 Block 中设置 warp 的数量为8,8 个 warp 会对 Thread Block 中的计算任务进行划分,每个 warp 计算任务为[32, 32]﹒[32, 32]=[32, 32]的矩阵乘,经过内积方向的 4 次 32 的循环,在 warp level 便可以将内积方向K=128 完全计算得[M, N]=[32, 32]的计算结果,则 8 个 warp 合并可得[M, N]=[64, 128]的计算结果。

  如上文所述,程序将内积方向的K维度 128 拆分为 2 组 64 进行计算,每组 64 进一步拆分为 2 组 32(k)进行计算。这么做的目的是将 data 以及 weight 的全局内存加载变成了 Double buffer 模式,即首先将第一组的数据由全局内存加载到共享内存,然后在利用第一组的数据进行计算前,便提交第二组数据由全局内存加载到共享内存的过程,这样可以利用第一组数据的计算过程时间来隐藏第二组数据的全局内存加载到共享内存的过程的时间,整体流程示意图如下:

  如前所述,每个 warp 计算任务为[32, 32]﹒[32, 32]=[32, 32]的矩阵乘,因此在 warp 的计算层次配合[m ,n ,k]=[16, 8, 32]的形状,需要进行 row=2,col=4,共 row*col=8 次 mma 的计算才可以得到 warp 层次的计算结果,在计算时配合 ldmatrix 的使用可以进一步提高程序的计算性能。

  对于 Mma 层次的计算,根据 mma 的形状,单次计算为[m , k]﹒[n, k]=[16, 32]﹒[8, 32]=[16, 8]。

  3. 4 Shoutcut的合并计算

  经过以上计算,每个 Block 程序会得到 Conv3 的[M, N]=[64, 128]的计算结果,由于程序对 Bn+Sum+Relu 进行了合并,因此需要对 Res3.1 输出的原始数据进行加载。根据 Grid[x , y, z]的划分,可以相应的得到 Shortcut 的数据偏移,为了隐藏这部分数据在全局内存加载到共享内存时通信延迟,程序利用了 A100 GPU 异步复制(pipeline_memcpy_async)的新特性,在程序的最开始便提交了这部分数据的加载,这样可以最大程度上利用计算的时间同时进行数据的加载以隐藏 Shortcut 的通信延迟,如图 4 所示。完成数据的加载后,会以 warp 为单位对每一个计算结果进行 Bn+Sum+Relu 的操作,最后将数据由寄存器写回共享内存,再写回全局内存完成整个卷积合并算法的计算。

  4. 性能提升效果

  根据上文介绍的卷积合并优化算法,在 TensorRT 中增加了关于卷积合并算法的 plugin 以替代原始算法,在 A100 GPU 单卡进行 Conv3+Bn+Sum+Relu 性能测试,在 BatchSize=2048 的情况下,原算法的性能为 123TOPS,经过优化后的卷积合并优化算法性能为 141TOPS,算子相比较原算法可以带来 14.6% 的性能提升。通过合并 Res3.2、Res3.3、Res3.4 三部分 Conv3+Bn+Sum+Relu 算子合并,可将 Resnet50 推理性能提升1%-2%。同样该算法合并思路可以用到其他残差结构中,通过合理的算法设计带来整体的程序性能提升。

  在 MLPerf V2.0 推理竞赛中,浪潮通过软件与硬件优化,基于 ImageNet 数据集 Resnet50 模型,在 Offline 场景中达到了 449,856 samples/s的计算性能,位居世界第一。

24小时阅读排行

    最新新闻

      编辑推荐

        相关新闻