AI架构设计3：GPU运算模型

GPU的运算模型

上一章节已经介绍了显卡的基本结构，由内存和计算单元SM组成。那么显卡的运算模型是什么样子？它如何将机器学习的任务转化为合适的运算指令？当然这里不针对Cuda编程展开细节的讨论，但是还是需要对它的运算模型有一定的了解。

计算模型中的线程（Thread）、线程块（Block）和线程网格（Grid）三个概念，分别对应着Cuda Core、SM和kernel三种不同维度的计算单元。一个内核（kernel）会启动多个线程块，这些线程块会在SM之间进行调度。分到任务的SM则将线程块的内部线程分配给具体的Cuda Core执行。

在计算模型中，所有线程都是并行执行的。一个线程块只能分配到一个SM上执行，同一个SM中的线程可以互相通讯。一个线程网格生成多个Block会被调度到多个SM执行，但是不同SM之间的Block则互不通讯。

如何将机器学习的任务转为高效的运算指令是很关键的，很多的优化都基于这个计算模型，Cuda编程讲究的是如何重构计算任务，有效的利用不同级别的缓存进而提高并行计算的效率。下图为不同计算单位对应的内存类型。

以A100为例，在大模型微调或者训练的时候，研发者更多的关注在80GHBM的优化以便确保能够完整的加载大模型进行相应的训练，而对于底层的缓存优化则交由Cuda工程师去操心。

计算任务转化为矩阵运算

上图则将通过具体的实例来说明如何将机器学习的任务如何转为矩阵运算，毕竟显卡最拿手的就是矩阵运算。彩色图像一般由N个通道组成，通道数目可以是1、2、3和4。常见是1和3通道，而2和4通道不常见。如何理解通道，这里不展开解释。简单的说每个像素点通过N的通道（数值）组合渲染出这个像素的颜色。以下图为例，在进行运算的时候可以将图像分割，进而形成每个分块计算。

若是以卷积运算为例，每个Block内部的线程则计算负责对应图像块的卷积核计算。上图左边为多个卷积核（这个卷积核是通过训练而得到，每个卷积核专注识别某个模式，假设核大小为C*K*K，C为通道数）。那么接下来的卷积运算，则是将这些核盖在图像的同等窗口进行点积计算。

有论文已经证明了卷积核其实也是可以通过特定的Transformer来实现

在实际的计算过程中，其实是将这些卷积核和图像都转化为对应的矩阵，然后进行矩阵相乘，而显卡最擅长的就是矩阵相乘。

将多有的卷积核重构为一个矩阵，则一次运算出所有的结果

为了配合卷积核的一次计算将图像的数据重构为对应的冗余矩阵

通过结构重构让GPU做最擅长的矩阵运算

这里留一道有趣的思考题：若将图像分成块，那么块块又不能互相通讯，如何计算块块边界的卷积？当然这不重要，重要的是通过这个例子，大致能够让架构师明白GPU的运算模型以及机器学习在做运算时候算子的工作原理，以便后续在设计算力调度的时候又更加立体的理解。