撰文 | 郭冉、姚迟、郑泽康、柳俊丞
2020年末,OneFlow 发布了《OneFlow 性能优化分享:如何实现一个高效的 Softmax CUDA kernel?》 ,其中介绍了OneFlow深度优化后的Softmax,尤其对很多框架没有考虑的 half 类型做了充分优化,使得性能大幅超过了 cuDNN 的实现。
今天,奉上另一个重要算子 LayerNorm 的性能优化实践技术分享。
此外,OneFlow 还带上了可以独立使用的 OneFlow Softmax(具体见文末说明),欢迎大家试用、提建议。
目录
OneFlow 深度优化 LayerNorm CUDA Kernel 的技巧
欢迎下载体验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,最后按上述公式计算得到方