LightSeq的技术亮点主要在于定制的Operation、动态GPU内存复用和层级的解码优化,本文对这三方面分别做个粗浅的介绍。
论文链接:https://arxiv.org/pdf/2010.13887.pdf
1 定制的Operation
大多数深度学习框架通过调用方法库中的kernel function来实现编码器-解码器计算。这些kernel function一般是细粒度的,通常需要调用多个kernel function来实现一个功能模块。
以tensorflow实现的layer normalization为例,即使基于compilation optimization techniques (automatically merging broadcast and element-wise operations),仍需要三个kf和两个中间结果(mean and variance)的读写operation。
基于CUDA,我们可以为layer normalization写一个定制的kf,并且将两个中间结果写入到寄存器中,从而只需一次kernel launch且不需要将中间结果写入到gpu内存中,大幅减少了计算耗时。
基于这个idea,LightSeq结合cuBLAS的GEMM和定制kf来实现Transformer layers: