Repository navigation
Investigate custom Triton kernels for depthwise-separable convolutions in Conformer #2733
Description
Describe the bug
A decent chunk of time in the Conformer model at training time is spent in the convolution module. Of that, a decent chunk is in the depthwise convolution, which sets groups to the input size. Suspiciously much, if anything.
On ROCm, at least on my 7900XT, it appeared that streaming: True shredded performance and ~doubles training time (which i didn't observe on CUDA), so I tried giving optimizing it a shot.
I didn't observe improvements from using torch.compile, even with max-autotune.
I tried AMD's Composable Kernels and the relevant torch.compile backend, but I will decline elaborating on my experience getting it working lest this become an unhinged rant at the ROCm ecosystem.
So, I wrote up a prototype using Triton (OpenAI's compiler, not the inference server) with no prior experience for a forward pass, reusing realistic shapes:
This prompted an instant massive WTF. I am not 100% sure my benchmark is correct (I reused Triton's tutorials'), but I couldn't find anything wrong and the output tensors matched.
Even more strangely, similar performance benefits could be observed on A100.
So, this issue will track my progress on trying to apply this in practice. Currently only doing this on my free time.
Expected behaviour
speed
To Reproduce
No response
Environment Details
No response
Relevant Log Output
No response
Additional Context
No response
Activity
asumagic commentedon Oct 27, 2024
Noting that I wrote this kernel to process the input tensor in a
(batch, steps, channels)format.In comparison, PyTorch kernels use a
(batch, channels, steps)format. It remains a fair comparison, since there is no way around this. PyTorch has some support for NWHC, but only forConv2D.In both cases, the native representation was fed and thus no permutation affects the benchmark.
The performance differential might originate from this (maybe? not sure) and the existing kernels not being specialized enough? Haven't looked into PyTorch/CuDNN to understand why.
asumagic commentedon Oct 30, 2024
I spent a bit more time on this recently in my spare time and implemented the backward pass. It is probably not fully correct yet, especially when Dynamic Chunk Convolutions are involved (getting gradcheck failures and not quite sure how to resolve them yet).
I observed (very) significant performance gains on ROCm, especially with DCC, but a degradation on Nvidia (and an A100 definitely shouldn't be slower than my 7900XT).
I didn't try to inspect performance further, but it might have to do with the way I implemented the for loop over the kernel size that might be way suboptimal. Will have to check if trying to remove away that loop to express it as a single load and dot product would work.
I also hacked a dirty proof of concept to add it to a Conformer RNN-T model. It helped training speed a fair bit, and didn't fully break the model (although I didn't let it run nearly long enough to check if WER was correct).
Another potential candidate would be the downsampling CNN with Fast Conformer (which replaces some Conv2D by a depthwise-separable convolution), which is also super slow on ROCm, but that requires implementing stride.