Repository: pytorch
Issue: [RFC] Integrate DeepSpeed CUDA kernels
Labels: feature,triaged

**Motivation**

On May 27th DeepSpeed [achieved the fastest BERT training time.](https://www.deepspeed.ai/news/2020/05/27/fastest-bert-training.html) with a claimed speedup of 28% over [NVIDIA BERT](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/modeling.py), which is similar in its implementation to nn.Transformer (which in turn is very similar to fairseq). As part of this they shipped [a set of CUDA kernels](https://github.com/microsoft/DeepSpeed/tree/master/csrc/transformer), which are then [launched by a C++ BERT implementation](https://github.com/microsoft/DeepSpeed/blob/master/csrc/transformer/ds_transformer_cuda.cpp#L177-L288). The blog post credits about 10% of these improvements to improved CPU-GPU computation overlap via an improved DataLoader and via  a sparse CrossEntroyLoss, so we should be able to get an optimistic 15% just through the transformer kernels. They have both a stochastic and a non-stochastic version of BERT. The non-stochastic version is closest to our nn.Transformer (which is also standardized on the [BERT paper](https://arxiv.org/abs/1706.03762)) and I propose we port the CUDA kernels involved so we can offer DeepSpeed-level performance by default in PyTorch.

**Proposal**

Similar to nn.Transformer the [DeepSpeed BERT implementation decomposes into a TransformerLayer,](https://www.deepspeed.ai/tutorials/transformer_kernel/#integrate-transformer-kernel) which is where the optimized kernels kick in. So we should hopefully be able to expect the promised speedups by porting this TransformerLayer to PyTorch

Here are the [major layers / operators](https://github.com/microsoft/DeepSpeed/tree/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/includes) [involved](https://github.com/microsoft/DeepSpeed/blob/31f46feee2d491d58a13404e354440551de9d5bf/op_builder/transformer.py#L5) in the TransformerLayer implementation and their corresponding CUDA kernels

[Normalize_Layer](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/includes/normalize_layer.h#L12) - [fused_bias_residual_layer_norm](https://github.com/microsoft/DeepSpeed/blob/c78c29f9384ec8a3addb049504a4167508fe1e4d/csrc/transformer/normalize_kernels.cu#L16)

[FeedForward](https://github.com/microsoft/DeepSpeed/blob/8353c5949d395d27c898c78fd3ae72ccfe878c26/csrc/includes/feed_forward.h#L10) - cublas_gemm_ex, column_sum_reduce (fused transpose + add)

[StridedBatchGemm](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/includes/strided_batch_gemm.h#L9) - cublas_strided_batched_gemm

[Softmax](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/includes/softmax.h#L13) - [attn_softmax](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/transformer/softmax_kernels.cu#L9) (takes mask argument)

[Dropout](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/includes/dropout.h#L8) - [dropout_kernel](https://github.com/microsoft/DeepSpeed/blob/c78c29f9384ec8a3addb049504a4167508fe1e4d/csrc/transformer/dropout_kernels.cu#L465) (takes mask argument)

[Gelu](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/includes/gelu.h#L16) - [fused_bias_gelu](https://github.com/microsoft/DeepSpeed/blob/f0f2a7026877d3fd2f6f5565a67cdffc89750cf0/csrc/transformer/gelu_kernels.cu#L98) (gelu fused with bias addition)

[launch_bias_add_transform_0213](https://github.com/microsoft/DeepSpeed/blob/95575579b3304f3cdd6bc1f9715535d424768b87/csrc/transformer/ds_transformer_cuda.cpp#L206) - [bias_add_transform_0213](https://github.com/microsoft/DeepSpeed/blob/c78c29f9384ec8a3addb049504a4167508fe1e4d/csrc/transformer/transform_kernels.cu#L161)

[launch_transform4d_0213](https://github.com/microsoft/DeepSpeed/blob/95575579b3304f3cdd6bc1f9715535d424768b87/csrc/transformer/ds_transformer_cuda.cpp#L223) - [transform4d_0213](https://github.com/microsoft/DeepSpeed/blob/c78c29f9384ec8a3addb049504a4167508fe1e4d/csrc/transformer/transform_kernels.cu#L402) (MHA output concatenation)

Possible Next steps



1. Setup benchmark[ based on fairseq reference pipeline](https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md) 
2. Integrate deepspeed via [Python API](https://github.com/microsoft/DeepSpeed/blob/31f46feee2d491d58a13404e354440551de9d5bf/deepspeed/ops/transformer/transformer.py#L421) to verify performance gains
3. Write TransformerEncoderLayer as private native function backed by DeepSpeed kernels
4. Integrate fused nn.TransformerEncoderLayer into [fairseq TransformerEncoderLayer](https://github.com/pytorch/fairseq/blob/4817a9142f49793ec2eedbd71fe5bd872e58e7b5/fairseq/modules/transformer_layer.py#L17) or such

Further goals / variants

We can incrementally pull apart the TransformerEncoderLayer and interface with the individual efficient kernels with potentially various degrees of performance improvement. Some additional outcomes that can aid Transformers in a wider sense could include



*   Fused nn.MHA
*   Masked variants for MHA experimentation (softmax and dropout)
*   Invertible variants of LayerNorm and Softmax
*   Efficient kernel for torchtext MultiheadAttentionContainer output reshaping

**Related work**

There are still many other ongoing efforts related to parts of DeepSpeed around dataloading, distributed training and optimizers, but none of them seem to cover porting these CUDA kernels that make up the efficient Transformer implementation itself. Part of the reason for this RFC is also to confirm this.

- DataLoader (Vitaly)
- Porting ZerO optimizers (Shen Li)


cc @ngimel @VitalyFedyunin @mrshenli