Custom Triton GPU kernels maintain wall-clock training efficiency.
Alexia Jolicoeur-Martineau#1122
New paper 🧵
We show that dynamic short convolutions consistently improve Transformers across scales. We make these gains practical with an efficient parameterization and custom Triton GPU kernels.
The improvements carry over to MoEs and linear attention variants (Mamba-2/GDN).
Users congratulated the release of Dynamic Short Convolutions for improving Transformers because it extends canon_layers with dynamic kernel sizes.
This work looks interesting, basically like @ZeyuanAllenZhu Canon layer but the conv weights are input dependent.
Which sounds horrible for perf, but they focus a lot on this and include triton kernels as well as consider wall clock time in their scaling experiments.
New paper 🧵
We show that dynamic short convolutions consistently improve Transformers across scales. We make these gains practical with an efficient parameterization and custom Triton GPU kernels.
The improvements carry over to MoEs and linear attention variants (Mamba-2/GDN).
Triton Kernels: https://github.com/OliverSieberling/dynamic-conv1d Paper: https://arxiv.org/abs/2606.03825
We motivate dynamic short convolutions by the need for context-dependent local composition in language. E.g. "the old can opener" vs. "the old can swim"
In particular, when applied to QKV, dynamic short convolutions can compose local context into multi-token keys for retrieval.
If we place dynamic short convolutions after every linear layer (instead of just after the qkv-projection), the gains are even larger.
Fitting scaling laws for this variants gives an approximate 1.60x compute advantage over compute-matched transformers.
Naive PyTorch implementations of dynamic convolutions are veryyy slow. The main bottleneck is memory traffic from materializing/moving intermediate tensors.
We address this with custom Triton GPU kernels that perform competitively with CUDA kernels for static short convolutions.
We provide two kernels:
- General/Head-wise: shares filters across groups of channels, which shrinks convolution weight tensor (main I/O bottleneck) and improves latency - Low-rank: fuses an up-projection and produces the dynamic convolution weights on-chip, without materializing in HBM.
We integrate our Triton kernels into lm-engine and measure end-to-end training throughput on an H100.
Adding dynamic convolutions on QKV is only ~7% slower, and therefore the 1.33x compute advantage translates into a significant wall-clock time advantage.
A static short convolution is a depthwise-separable causal convolution with small kernel width (typically W=4). We make it dynamic by letting the filter depend on the current time step. Each token therefore selects its own local convolution filter.
We train language models of various scales (150M-2B params) and apply dynamic short convolutions to Q, K, and V before the attention. We find that this significantly improves language modeling across scales.
Fitting scaling laws suggests an approximate 1.33x compute advantage.
Our recipe transfers beyond standard transformers:
Modern linear RNNs already use static short convolutions on the queries, keys, and values. Replacing them with dynamic short convolutions substantially improves language modeling performance for both Mamba-2 and Gated DeltaNet.
@osieberling Congrats on the release! So this is basically canon_layers with dynamic kernel size right?
Joint work with @bharatrunwal2 @rpanda89 @yoonrkim, supported by @MITIBMLab.
@Ali_NT99 not quite, the convolution kernel itself is input-dependent, so each token performs a different (learned) convolution. Through this you could potentially learn a dynamic kernel size, but the approach is more general/powerful than just this.
Custom Triton GPU kernels maintain wall-clock training efficiency.
Alexia Jolicoeur-Martineau#1122New paper 🧵
We show that dynamic short convolutions consistently improve Transformers across scales. We make these gains practical with an efficient parameterization and custom Triton GPU kernels.
The improvements carry over to MoEs and linear attention variants (Mamba-2/GDN).