Memory Wall in Neural Network Inference

Memory Wall in Neural Network Inference

神经网络推理的瓶颈在于访存带宽，通常无法发挥出加速器的全部算力。本文总结了目前常用的推理加速器及其设计，并分析了常用神经网络的访存瓶颈。文章大部分内容参考自Computer Architecture: A Quantitative Approach。

1 Compute centric accelerators

1.1 CPU

一般来说，CPU擅长于做控制逻辑复杂的workload。但是，CPU里面也存在SIMD的加速器设计。

• Vector Architecture in RISC

  • Basic idea

    • Read sets of data elements into “vector registers”
    • Operate on those registers
    • Disperse the results back into memory

  • Y = a × X + Y

    左边使用Loop的原始汇编代码，右边是使用Vector Architecture优化后的汇编代码。

    

    • 如果 X 和 Y 均为32维向量，1 lane

      那么使用Vector Architecture加速，需要32×3个cycle

  • Optimization

    可以优化的地方很多，讲两个有意思的优化。

    • Multiple Lanes

      如果是4 lane的话，上述例子仅需要32×3/4个cycle就可以完成。

      

    • Scatter-Gather

      读取的向量可以不存储在连续的内存中。

      例如，执行这段代码：

      for (int i=0; i<n; i++){
          A[K[i]]=A[K[i]]+C[M[i]];
      }
      

      使用Scatter-Gather优化后的汇编代码如下：

      

      这种优化技术用在稀疏矩阵场景很有用。

• SIMD extensions for x86

  简单来讲，就是利用长寄存器，一次处理多个数据。比如寄存器长度是256 bit，一个数据的长度是64 bit，这时候，就可以一次处理4个数据。

  使用SIMD extensions优化Y = a × X + Y 后的汇编指令：

  

  • 这种优化要求向量必须存储在连续的内存空间中

• Vector Arch. vs. SIMD Extensions

  看这张图，对于神经网络推理加速场景，x86完全干不过RISC（哈哈，连高性能计算都干不过RISC，x86是不是要凉啊）。

  

1.2 GPU

简单来讲，GPU的架构就是堆很多的寄存器和ALU，获取高算力。对应到下图，你可以看到很多的SIMD Lane 和 DP Unit。

但是，GPU里面的处理器算力很强，访存带宽跟不上怎么办？首先，它使用了HBM，内存带宽很高；另外，其中的SIMT (Single Instruction, Multiple Thread) 编程模型利用多线程来掩盖访存时延。

如下图，简单来讲，当thread 8运行的时候，其他的线程（比如thread 1, 3）也没闲着，它们正在取数据。这样的话，GPU里面的处理器时刻处于满负载状态，不需要傻傻地等到数据加载到寄存器才运行。（熟悉GPU的同学应该发现了，图里面的thread改成warp更合适，懒得改了）

1.3 TPU

听过好多人吹google的TPU多么牛逼，其实也就那样，就是一个专门用来加速GEMM的设备。下图是TPU的架构图，右边的Matrix multiply unit 就是TPU的核心部件。

通用的加速器GPU可以用来加速向量加法，甚至有人整活，直接在GPU上构建了一个数据库。然而，TPU是一个专用的加速器，只能用来加速GEMM算子，其他的活统统干不了。

不过，TPU里面加速GEMM的硬件设计很有意思。专业的名字叫做Systolic Execution，不过，我喜欢叫它脉冲阵列。如果矩阵的维数是N×N，TPU使用2N-1个cycle就可以把GEMM算子计算完。

2 Neural network and Memory wall

在当前以计算为中心的编程模型中，需要先从内存中把数据加载到cache，然后ALU进行计算任务，最后将结果保存至内存中。

目前，ALU算的很快，把内存中数据加载到cache的带宽有限，这时候，ALU的算力就会被浪费，性能瓶颈卡会在内存带宽上，这就是“内存墙”。如果我们读入cache一个数据，对它进行N次计算操作后才保存结果，N定义为Operational Density。如果Operational Density，就没有“内存墙”的问题。如果Operational Density比较低，程序就会卡在“内存墙”上。

接下来，我们分析一下MLP、CNN、LSTM和Transformer的Operational Density，判断它们是否会卡在内存墙上。

2.1 MLP

MLPs (Multi-Layer Perceptions) ，多层感知机，大家应该都挺熟悉的，里面的算子是GEMV。

• Operational Density为2×batch size每一个矩阵参数乘和加操作各一次。

2.2 CNN

CNN 里面的基础的算子是卷积算子。convolution kernel 复用很多次，Operational Density很高。

• Operational Density为 2 x DimFM[i]^2

2.3 LSTM

LSTM里面主要也是GEMV算子。

• Operational Density为~2

2.4 Transformer

Transformer中的参数Operational Density与sequence len有关。假设输入token数量为N，generate token 数量为S。主要是GEMM和GEMV算子。

• 参数Operational Density为 (N×2+S×2)/(1+S)

在推理过程中，不仅仅需要访问参数，还需要访问KV cache (>30%)

• KV cache Operational Density为 2
• pagedAttention
• 参数Operational Density随batch递增，但是，KV cache Operational Density仍然是2。

3 Roofline

• 除了CNN，其他Neural network 的 Operational Density比较低，推理的时候卡在内存带宽，发挥不出硬件算力。
• 增大batch从而增大Operational Density，是缓解这个问题的一种方法。