Future Plan of Transformer Kernel

[hxbloom]

I'm using Megatron-LM example to train GPT-2 on my cluster. I've also tested the DeepSpeed Transformer Kernel in the bing_bert example, it is really helpful, much faster than the original PyTorch version with less memory consumption.

I would like to know if you have any future plan to extend the transformer kernel, for example, support more models like GPT-2, or integrate model parallel into the kernel for large model training?

Thanks!

[RezaYazdaniAminabadi]

Hi Dong,

Thanks for pointing out the use-cases for Transformer kernel. There us a plan for supporting other types of transformer networks. For the GPT2, we have the kernels modified to support that, and we obtained about 30% and 40% speedup in the forward and backward pass. We are going to release them very soon. Please stay tuned!

Thanks,
Reza

[hxbloom]

Look forward to your next release!

Dong