Faster Inference & Training Roadmap

[jeromeku]

@danielhanchen

In the unsloth Gemma intro blogpost, you mention VRAM increase due to larger MLP size in Gemma compared to Llama and Mistral, and show a graph demonstrating decreased memory usage when running unsloth vs. HF and FA2:

• How does unsloth reduce memory usage?
• What are the model and runtime configs used to generate the HF vs FA2 vs unsloth graph? Is it inference or training?

Curious what optimizations are leading to memory decrease -- quantization, autograd efficiency, etc.

[danielhanchen]

@jeromeku I will get to reviewing GPTQ - sorry on the delay!!

• The VRAM reductions are from Unsloth's optims :) Ie Triton kernels, making memory copies go away, FA2, autograd tricks etc
• Oh training! All use 4bit quantization, the ga is set to 1, and bsz and sequence length moves dynamically so I can measure VRAM usage.

[jeromeku]

@danielhanchen

Thanks -- would be helpful to have a step-by-step breakdown of where the memory savings are coming from, i.e., an ablation study.

Is there interest in faster inference kernels, or is the focus primarily on the training side?

[danielhanchen]

@jeromeku For Mistral itself: https://unsloth.ai/blog/mistral-benchmark

Gemma's VRAM reduction should be similar to our breakdown for Mistral.

For inference for Gemma - I did make it 2x faster, but it's mainly cobbling up ideas from vLLM and other packages, so I only spent 1 week on it :) The goal is to merge GPT Fast and other ideas like EAGLE to make inference faster :)

[jeromeku]

@danielhanchen

I'd be interested in contributing on the inference front -- let's create a priority list of ideas for implementation?

[danielhanchen]

@jeromeku That'll be cool!! :) We can collab either via Github or async on our Discord - whatever suites you :)

[jeromeku]

@danielhanchen

Looking forward to it!

What's top of mind currently? Perhaps we can draw up a roadmap (if one doesn't already exist).

[danielhanchen]

@jeromeku Oh ye a roadmap would be nice - don't actually have one for inference specifically :)

[jeromeku]

@danielhanchen

You mentioned integrating ideas from fastGPT and EAGLE, what others did you have in mind?

What's on the roadmap for fine-tuning / training -- architectures, algorithms, etc.? Asking so I know what literature / code to review.

[danielhanchen]

@jeromeku In terms of inference specifically:

1. GPT Fast
2. Speculative Decoding (use a small model to generate tokens, then use a large model in 1 forward pass and see if the argmax of the logits match)
3. EAGLE (Speculative Decoding but only Word2Vec style ie lm_head -> embeddings)
4. All quant methods - HQQ, AWQ, Exllama etc
5. vLLM's Paged Attention
6. Full 1 singular Triton kernel fusion - ie can we write 1 forward pass in 1 humoungous Triton kernel? Very hard since there are synchronizations which have to be done
7. Using float8 like Fire Attention. cuDNN has float8 flash attention I think as well.
8. Rewriting matrix vector multiplication in Triton exactly (like what you were trying to do with GPTQ but not matmul, but matvec
9. Torch export

I might have more, but those are from the top of my head.

For training / finetuning:

1. Fast MoE matmul kernel DeepseekMoE support with Fused MoE kernel vllm-project/vllm#2453 but for training - much more complex than inference on batch sizes of 1. Mixtral selects the top 2 experts, which can easily be done in Triton. However, when you have bsz>1, we have issues. One has to do dynamic compressed packing then call torch.bmm. The backward pass is even more problematic, since it requires a reversed packing then calling torch.bmm, then deconstructing it. A nightmare.
2. Galore - extremely fascinating projecting gradients to a small (rank, rank) matrix, then using SVD to update the projectors. It's not Galore that I was fascinatined by, but rather Lomo, which does gradient updates dynamically, and this can save 20GB of VRAM during pretraining.
3. 1.58bit - I recently wrote on HN about how 1.58bit allows one to not to multiplications since (-1, 0, 1) becomes a simple sign flip then the mantissas are added after the exponents are flipped. Using 8bit floats, 1.58bit uses 2x less space than float8, which makes it possible to cram 2x transistors. Writing it in Triton can be more complex.

Just a brain dump!

[jeromeku]

@danielhanchen

Love it.

Inference:

• vLLM paged attention - happy to look into this as well as KV cache quantization.
• Triton GEMV - seems pretty straightforward -- can prototype such implementation in Triton -- I believe Torch compile already can generate such a kernel with proper inductor settings (effectively decomposes to a vectorized mul + add). Can also adapt existing GEMV CUDA kernels for quantized weights.
• Torch export - can look into it. I've done some work into decoupling Triton kernels from Triton runtime.

Training:

• I've been playing around with implementing a Cutlass kernel for MoE matmul which could help with larger batch sizes.
• Galore - top of my list of papers to read
• 1.58 bit - also been looking into Cutlass for optimizing custom quantized ops.

[danielhanchen]

• Oh ye KV cache quant is cool! On issue I have with it is dynamically quantizing the KV cache will cause overhead issues - a super fast method for quantization will have to be deployed.
• Triton GEMV: Ye the kernel is fine to create - one possibility is can we fold GEMVs and layernorms and everything into 1 large kernel
• CUTLASS is good :) Looking forward to it - my main view is we need to use as much Triton as possible for device agnostic purposes :)
• Ye Galore and 1.58bit :) 1.58bit actually can be very doable in Triton. Galore is very cool.

[jeromeku]

Let me know what I should prioritize.

Also, can you expand more on Triton GEMV? What kind of horizontal / vertical fusions to target?

[danielhanchen]

Oh so GEMV is generally OK I guess - the issue is the dequant step merged in (ie what you were doing with GPTQ, except its not matrix matrix mult but matrix vector mult) this allows different optimizations - ie is blocked mm better or is column or is row wise mv better? It depends on the cache footprint

But the goal is can we somehow merge X @ Wq, X @ Wk, X @ Wv together with RoPE and attention and everything into 1 large kernel