CUDA优化之LayerNorm性能优化实践

撰文 | 郭冉、姚迟、郑泽康、柳俊丞

2020年末，OneFlow 发布了《OneFlow 性能优化分享：如何实现一个高效的 Softmax CUDA kernel？》 ，其中介绍了OneFlow深度优化后的Softmax，尤其对很多框架没有考虑的 half 类型做了充分优化，使得性能大幅超过了 cuDNN 的实现。

今天，奉上另一个重要算子 LayerNorm 的性能优化实践技术分享。

此外，OneFlow 还带上了可以独立使用的 OneFlow Softmax（具体见文末说明），欢迎大家试用、提建议。

目录

OneFlow 性能优化后的测试结果

与 NVIDIA Apex 的对比结果

与 PyTorch 的对比结果

LayerNorm 性能优化

LayerNorm 计算方法

LayerNorm 中求方差的方法

OneFlow 深度优化 LayerNorm CUDA Kernel 的技巧

1.num_cols <= 1024的情况

2.num_cols > 1024的情况

3.num_cols 较大时，不使用 Shared Memory 的情况

OneFlow Softmax 库

欢迎下载体验OneFlow新一代开源深度学习框架：https://github.com/Oneflow-Inc/oneflow/

---

OneFlow 性能优化后的测试结果

OneFlow 优化后的 LayerNorm 分别与 NVIDIA Apex、PyTorch 做了性能对比，测试结果显示，OneFlow LayerNorm 有明显的性能优势。

与 NVIDIA Apex 的对比结果

NVIDIA Apex 中实现了高效的 fused LayerNorm Kernel 来扩展 PyTorch 算子，我们对 OneFlow 优化后的 LayerNorm Kernel 和 NVIDIA Apex 进行了对比测试，测试结果如下：

横轴为 num_cols 大小，纵轴为 Kernel 执行需要的时间(越低越好)：

我们将时间换算成访存带宽，结果如下，纵轴为 Kernel 达到的有效带宽(越高越好)：

其中测试环境为 NVIDIA A100-PCIE-40GB GPU，数据类型为 half， Shape =(49152， num_cols)，我们将最后一维动态变化，测试了从32到32768不同大小的 LayerNorm Kernel，可以看到在所有情况下，OneFlow 的 Kernel 执行时间和有效访存带宽都优于 Apex 的实现。

与 PyTorch 的对比结果

PyTorch 的 LayerNorm 暂时不支持 half 类型，因此我们用 float类型做了一组对照，需要注意的是PyTorch中LayerNorm是分两个CUDA Kernel(RowwiseMomentsCUDAKernel和LayerNormForwardCUDAKernel)做的，所以看起来性能比较差。

横轴为 num_cols 大小，纵轴为 Kernel 执行需要的时间(越低越好)：

可以看到，在各组对比实验中，OneFlow 的性能也是最优的。

LayerNorm 性能优化

LayerNorm 是语言模型中常用的操作之一，其 CUDA Kernel 实现的高效性会影响很多网络最终的训练速度，Softmax 这种优化方法也适用于 LayerNorm，LayerNorm 的数据也可以表示为 (num_rows, num_cols)，计算过程中对每一行的元素做 Reduce 操作求均值方差。因此我们使用了和 Softmax 同样的优化方法来优化 LayerNorm 操作，本文以 LayerNorm 前向计算为例进行介绍。

LayerNorm 计算方法

以 PyTorch 为例，LayerNorm 的接口为:

torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None)

其中 input 形状为：[∗, normalized_shape[0], normalized_shape[1], …,normalized_shape[−1]]

第一个参数 normalized_shape 只能是输入 x_shape 的后几维，例如 x_shape 为 (N, C, H, W), normalized_shape 可以是 (W)， (H, W) ，(C, H, W) 或 (N, C, H, W)。输入 x 在 normalized_shape 这几维上求均值和方差。

第三个参数 elementwise_affine 代表是否要对 normalize 的结果做变换，即 normalize 的结果乘 gamma，加 beta。若 elementwise_affine=True，就多了两个模型参数 gamma 和beta，形状为 normalized_shape。

例如对于输入 x 形状为 (N, C, H, W)， normalized_shape 为 (H, W) 的情况，可以理解为输入 x 为 (N*C, H*W)，在 N*C 个行上，每行有 H*W 个元素，对每行的元素求均值和方差，得到 N*C 个 mean 和 inv_variance，再对输入按如下 LayerNorm 的计算公式计算得到 y。若 elementwise_affine=True ，则有 H*W 个 gamma 和 beta，对每行 H*W 个的元素做变换。

LayerNorm 中求方差的方法

常见的求方差的方法有 two pass 方法、naive 方法、和 Welford 算法，本文摘录一些关键的公式和结论，详细的介绍和推导可参考：Wiki: Algorithms for calculating variance ，和 GiantPandaCV: 用Welford算法实现LN的方差更新

• two-pass方法

使用的公式是：

two-pass 是指这种方法需要遍历两遍数据，第一遍累加 x 得到均值，第二遍用上面公式计算得到方差。这种方法在 n 比较小时仍然是数值稳定的。

• naive方法

使用的公式是：

这种方法是一种 single pass 方法，在计算方差时只需要遍历一遍数据累加 x 的平方及累加 x，最后按上述公式计算得到方