3 ms·
Slightly tangential, but I had intended to start playing around with LLaMA and building some agents. I got the 4-bit versions up and running on my 3090 before
by fpgaminer 4y ago
Slightly tangential, but I had intended to start playing around with LLaMA and building some agents. I got the 4-bit versions up and running on my 3090 before I was quickly nerd snipped by a performance problem...
The popular repo for quantizing and running LLaMA is the GPTQ-for-llama repo on github, which mostly copies from the GPTQ authors. The CUDA kernels are needed to support the specific kind of quantization that GPTQ does.
Problem is, while those CUDA kernels are great at short prompt lengths, they fall apart at long prompt lengths. You could see people complaining about this, seeing their inference speeds slowly tanking as their chats/prompts/etc got longer.
So off I went, spending the last week or so re-writing the kernels in Triton. I've now got my kernels running faster than the CUDA kernels at all sizes [0]. And I'm busily optimizing and fusing other areas. The latest MLP fusion kernels gave another couple percentage boost in performance.
Yet I still haven't actually played with LLaMA and made those agents I wanted... sigh And now I'm debating diving into the Triton source code, because they removed integer unpacking instructions during one of their recent rewrites. So I had to use a hack in my kernels which causes them to use more bandwidth than they otherwise should. Think of the performance they could have with those! ... (someone please stop me...)
[0] https://github.com/fpgaminer/GPTQ-triton/ https://github.com/fpgaminer/GPTQ-triton/
- Aeolun 4y agoAh well, at least you are spending your time productively.
- Nimitz14 4y agoAny recommendations for material to know to do exactly this sort of optimization work (involving triton)? I guess it's a mix of knowing computer architecture and compilers?
- fpgaminer 4y agoTriton itself is fairly "easy", at least as far as "low level optimization languages" go. It's just (restricted) python. If you know PyTorch, you can muddle your way through Triton. They have a few tutorials. Reading up on nvidia architectures, PTX, and CUDA are likely to improve your skill at Triton.
- Nimitz14 4y agoYeah I've been able to kinda muddle my way through but progress is frustrating as I frequently get errors I don't understand. Was thinking maybe I should study compilers a bit and then I will be able to understand the triton source code better which will help me understand why I get errors.
- aljungberg 4y agoYour triton code is great, nice work. Wouldn’t feel too bad about spending your time that way! As it happens I was also thinking it might be worthwhile to dive into the Triton sources but for another reason: half2 arithmetic. That’s one thing that the Triton branch lost that the (faster) CUDA kernels had and I think it made a difference. In theory with compatible hardware you can retire twice as many ops per second when processing float16 data which we are in this case. Can’t see anyone having tried to get half2 to work with Triton though.