Megatron fused CUDA kernels to improve Hugging Face model classes' scalability

[g-karthik]

🚀 Feature request

Support for custom fused CUDA kernels with HF model classes.

Motivation

It appears that Hugging Face model classes do not scale very well as-is unlike Megatron-LM, even when the latter is configured with a degree of model-parallelization = 1 for a "fair" performance comparison.

One of the presumed reasons for this is that Megatron-LM leverages custom fused CUDA kernels written by NVIDIA, specifically these.

Could we get variants of existing HF classes (perhaps for GPT2Model, GPT2LMHeadModel, etc.) such that the variants leverage some/all of these fused CUDA kernels? All this while still ensuring that one can load the original pre-trained weights into these variant classes.

Any guidance/low-level thoughts towards making this happen would also be greatly useful!

@thomwolf @patrickvonplaten @LysandreJik @stas00

[stas00]

I think the biggest barrier to using custom CUDA kernel is that it'd require transformers to move from a python-only package, to a compilation-required type of package (even if JIT), which in my experience is the type of a package that is far from trivial to use and often raises a barrier to entry.

If I'm not mistaken some fused kernels have been pushed upstream into the pytorch-core, so if you know of any that we could receive precompiled via pytorch, then we can definitely use those.

And if they aren't and you have some resources to initiate the conversation - it'd definitely help to request that such kernels will be added to pytorch-core. Definitely tag me if I do start such a thread at pytorch Issues.

---

I love your spirit of proposing various performance optimizations, @g-karthik and I'd love to work on all of those you have been proposing here and at Deepspeed issues, but so far I find no free resources to do so and all my time is spent on making things work.