如何在CUDA中为Transformer编写一个PyTorch自定义层

如今,深度学习模型处于持续的演进中,它们正变得庞大而复杂。研究者们通常通过组合现有的 TensorFlow 或 PyTorch 操作符来发现新的架构。然而,有时候,我们可能需要通过自定义的操作符来实现更多的优化。随着深度学习模型规模不断增长,为实际生产和可扩展训练设计专门优化的操作符将会变得更加重要。因此,本文作者学习了如何在 CUDA 中为 Transformer 编写一个 PyTorch 自定义层。

性能分析

首先,我们需要对一种深度学习模型很熟悉,这样我们就可以找到其性能瓶颈,并查看在我们进行了优化之后有多大的提升。我们可以使用内置的 PyTorch 分析器,也可以使用通用的 python 分析器。我们将同时考察这两种方法。

torch.autograd.profiler

PyTorch 提供了一个名为「torch.autograd.profiler」的 API。我们可以通过如下方式使用该 API:

with torch.autograd.profiler.profile(use_cuda=True) as prof:
    # Execute ops here
print(prof)

接着,PyTorch 会自动找到每个操作符并衡量他们的性能。性能分析结果如下:

分析器显示出了每个操作符在 CPU 和 GPU 上花费的时间。分析结果很直观,并且看起来似乎很精确,但是我们很难分辨出每个操作符并将它们与我的源代码匹配起来。例如,上面的输出结果显示出了三个不同的「unsqueeze」操作符,但是我们并不知道它们是在哪里被调用的。因此,我转而使用其它的分析器来寻找性能的瓶颈点

逐行分析器

因为 PyTorch 是基于 python 编写的,所以我们也可以使用通用的 python 分析器。我找来了一个逐行分析器(https://github.com/rkern/line_profiler),它可以逐行分析一个 python 应用程序。在要分析的函数的顶部添加「@profiler」装饰器之后,我们可以在命令行中用「kernprof」替换「python」来运行分析器。此外,在 CUDA 的环境下,我们必须设置一个环境变量「CUDA_LAUNCH_BLOCKING」来同步对 CUDA 调用。

运行一个 epoch 的后分析多头注意力机制前馈函数的结果如上图所示。结果显示了测量每一行所花费的时间,因此我们可以很容易地找到需要优化的目标代码。我们将重点关注第 85、87 和 88 行中的掩码操作。它组合了多个操作符来模拟「掩码处理后的 softmax」操作:为 softmax 的掩码输入填充负无穷数,从而使 softmax 忽略它们。在本文中,我将尝试优化这些操作。请注意,它当前花费了函数执行时间的 19.1%(7.2 + 5.9 + 6.0),而第 86 行花费了 15.2% 的执行时间。让我们使用这个值作为对比基线。

还有另一个适合优化的地方:第 86 行和第 90 行中的矩阵乘法,因为它们的输入或输出都填满了许多 0。本文不会对此进行深入探讨。

掩码处理后的 Softmax

首先,我认为我们可以通过将运算过程封装进一个操作符中来优化掩码处理后的 softmax,因为执行多个操作符本身就会产生开销。每次调用每个独立的操作符时,对 CUDA 核函数的调用会产生开销,而主机和 GPU 之间的数据传输也需要时间。

我们将使用一个名为「MaskedSoftmax」的自定义 CUDA 操作符。我们将其直接简略地定义如下:

x 是一个softmax 函数数的输入张量,m 代表一个掩膜张量,s 是一个用于归一化的标量值。该方程与 softmax 类似,只是掩码处理后值被规定为零,并乘以归一化系数。下图显示了掩码处理后的 Softmax 的一个示例。掩码处理后的位置变为零,并且使用 softmax 计算出其余位置上的值。


第一版

我首先写了一个简单版的 Masked Softmax。它由三个与 softmax 具有相同计算流程的遍历组成:(1)找到一个输入的最大值,(2)计算指数运算的值的和,以及(3)将每个值作为输入计算出指数运算的值,用它们分别除以指数运算的值的和。与 softmax 的不同之处在于,它还会加载掩码值,如果掩码值为 1,则将每个对应位置上的输入值转换为零。

template <typename scalar_t>
__global__ void __launch_bounds__(32) masked_softmax_cuda_forward_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ mask,
    scalar_t* __restrict__ output,
    unsigned int hidden_size,
    unsigned int m0, // size of mask dimension 0
    unsigned int m1, // size of mask dimension 1
    unsigned int m2, // size of mask dimension 2
    scalar_t scale) {

  // This threadIdx.x is a number between 0 and 31 because we only launched 32 threads.
  const int tid = threadIdx.x;
  // blockIdx.x, y, z are offsets of 0th, 1st, 2nd dimensions of input tensor.
  const unsigned int ibase = blockIdx.x * gridDim.y * gridDim.z * hidden_size +
                             blockIdx.y * gridDim.z * hidden_size +
                             blockIdx.z * hidden_size;
  const unsigned int mbase = blockIdx.x * (m0 > 1 ? m1 * m2 * hidden_size : 0) +
                             blockIdx.y * (m1 > 1 ? m2 * hidden_size : 0) +
                             blockIdx.z * (m2 > 1 ? hidden_size : 0);
  unsigned shfl_mask = __ballot_sync(0xffffffff, threadIdx.x < hidden_size);

  // Find a maximum input.
  scalar_t max_x = -FLT_MAX;
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    max_x = fmaxf(max_x, m == 0 ? input[ibase + i] * scale : -FLT_MAX);
  }
  // Reduce values in threads to find a global maximum number.
  for (unsigned int i = 16; i > 0; i >>= 1) {
    max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
  }

  // Find a sum of exponential inputs.
  scalar_t exp_sum = 0;
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    exp_sum += m == 0 ? std::exp(input[ibase + i] * scale - max_x) : 0;
  }
  // Reduce values in threads to find a global summation of exponential inputs.
  for (unsigned int i = 16; i > 0; i >>= 1) {
    exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);
  }

  // Calculate outputs and save to global memory.
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    output[ibase + i] = m == 0 ? std::exp(input[ibase + i] * scale - max_x) / exp_sum : 0;
  }
}

CUDA 中有「warp」和「block」的概念。Warp 是一组 32 个线程,而一个 block 则包含多个 warp。每个 block 有一个共享的内存,任何线程都可以访问一个全局内存。每个线程使用不同的线程和 block 的 id 执行相同的核函数代码,因此每个核函数使用全局内存中的 id 查找和读取相关输入,并将每个输出保存到全局内存中。由于计算是分布式的,如果有需要,我们可能需要减少不同 block 或线程中的值。

在这个 softmax 的实现中,我们需要一个约简来获得值的和或最大值。由于访问全局/共享内存是 CUDA 核函数中常见的瓶颈,所以我试图绕开它。为此,我为每个 block 创建了一个 warp,并使用了「shuffle」函数。它使用寄存器进行 warp 内的通信,因此线程可以在不访问共享内存的情况下交换值。

for (unsigned int i = 16; i > 0; i >>= 1) {
  max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
}

通过这个自定义的操作符,掩码处理后的 softmax 占用执行时间的比例降至了 15%。这并不是一个巨大的提升,但无论如何也比之前要快一些了。

现在,内置的 PyTorch 分析器也显示出了这个自定义操作符的性能提升。因此,由于逐行的分析器需要用太长的时间进行性能分析,我将这个第一版的掩码处理后的 softmax 用作进行进一步优化的对比基线。

进一步的优化

正如我所提到的,对于全局内存的访问是一个主要的瓶颈。在一些假设条件下,我们可以最小化内存访问的次数。前面的第一版现在可以从全局内存中读取两种类型的值(掩码和输入)。用于归一化后的点乘注意力机制的掩码通常有如下所示的形式。

从最左或最右开始,它们是连续的,而基本的 transformer 只有从最左开始的三种形式。因此,我们不需要为每个输入加载掩码值。在读取每一行之前,加载一个表示掩码长度的值就足够了。

我们可以使用下面的代码直接将掩码转化为一种新的形式:

mask = mask.size(2) - mask.sum(dim=2, dtype=torch.int32)

接着,我们只需要首先加载掩码长度,将每个循环迭代与掩码长度相同的次数,并将其余输出设置为零。

// Load a mask length.
  const unsigned int mask_offset = blockIdx.x * (m0 > 1 ? m1 : 0) +
                                   blockIdx.z * (m1 > 1 ? 1 : 0);
  unsigned int mask_size = min(static_cast<unsigned int>(mask[mask_offset]),
                               hidden_size);

  unsigned shfl_mask = __ballot_sync(0xffffffff, threadIdx.x < mask_size);

  scalar_t max_x = -FLT_MAX;
  // Iterate loop as much as the mask length.
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    max_x = fmaxf(max_x, input[ibase + i] * scale);
  }
  for (unsigned int i = 16; i > 0; i >>= 1) {
    max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
  }

  scalar_t exp_sum = 0;
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    exp_sum += std::exp(input[ibase + i] * scale - max_x);
  }
  for (unsigned int i = 16; i > 0; i >>= 1) {
    exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);
  }

  // We initialized "output" to zero, so remaining outputs will be zero.
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    output[ibase + i] = std::exp(input[ibase + i] * scale - max_x) / exp_sum;
  }

这样一来,这个操作就变得快多了。它现在只占用了执行时间的 9%。

掩码处理后的 Softmax(MaskedSoftmax)的执行时间现在比第一版快 2.5 倍。

我还检查了这种优化在多大程度上提高了整个训练的速度。我在 lm1b 数据集上训练了语言模型,并且测量了运行每个(碎片)epoch 的平均时间。第一个 CUDA 的版本比单纯组合 PyTorch 操作符的方法快了约 0.8%,第二个版本比原始版本快了约 1.8%。

结语

我在 CUDA 中编写了一个自定义的操作符并使 Transformer 的训练快了约 2%。我首先希望仅仅在 CUDA 中重写一个操作符来得到巨大的性能提升,但事与愿违。影响性能的因素有很多,但是我不可能找到每一个因素。此外,由于我对 CUDA 并不熟悉,我也遇到了很多 bug。代码越多,bug 越多。这使得我写了很多意想不到的测试代码。这是在提升模型性能和用于写代码的时间之间的一种折中。

编写一个自定义的操作符并没有我想象的那么简单,但是我可以从中学到许多关于 CUDA 如何工作的知识,以及诸如 block、线程、核函数、内存、同步、缓存这样的概念。我希望本文能够对那些想要入门 CUDA 性能优化的人有所帮助。

  • 完整代码:https://github.com/tunz/tcop-pytorch 

  • 使用场景:https://github.com/tunz/transformer-pytorch.

原文链接:https://tunz.kr/post/5

工程PyTorchTransformerCUDAsoftmax
3
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

核函数技术

核函数包括线性核函数、多项式核函数、高斯核函数等,其中高斯核函数最常用,可以将数据映射到无穷维,也叫做径向基函数(Radial Basis Function 简称 RBF),是某种沿径向对称的标量函数。最常应用于SVM支持向量机中

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集(或特征)的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的,这让「神经网络在执行预测任务时可以更多关注输入中的相关部分,更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时,源句子中仅有少部分是相关的;因此,可以应用一个基于内容的注意力机制来根据源句子动态地生成一个(加权的)语境向量(context vector), 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

暂无评论
暂无评论~