Speculative decoding potential for running big LLMs on consumer grade GPUs efficiently #10466
Replies: 4 comments 17 replies
-
|
Would there be any benefit in pruning down a 0.5B model to be even smaller? From your examples above it looks like the speculative models' size reduction has the biggest effect? You could prune the later layers like this: https://arxiv.org/abs/2403.17887 but with a calibration dataset you could probably prune down the width of the MLP hidden state quite significantly too... The I think you could even apply L1-regularisation during fine-tuning to spasify the weights and then remove all those close to zero, but the effectiveness of this would depend on whether the induced sparseness was evenly distributed for the corresponding tensors in each layer (which from the paper above; I doubt is the case). It would be interesting to see where the balance point is between "tiny and fast/dumb" vs "small but slower/less-dumb" actually is. If using greedy speculation then it won't make any difference, but if you have to actually apply the softmax (instead of just finding the maximum logit), then for stuff like coding using only English; it would be perfectly valid to remove a lot (most) of the tokens and prune down the |
Beta Was this translation helpful? Give feedback.
-
@jukofyork I tried the AMD-Llama-135m on Llama2 and AMD-Llama-135m-code on CodeLlama and didn't see as much benefit. Just based on pure hunch I think the "magic" of speculation is exposed with distilled models using the full cross entropy loss training metric on a large data set so the model doesn't lose any particular part of its knowledge. It gives a "fuzzy" version of the bigger model tending to follow the same patterns as the big model. Thats why I think Qwen2.5 series and Llama3.1 paired with Llama3.2 work so well.... I think (hypothesize) all are using distilling to make the smaller versions of the models. In theory any created model could be distilled but its compute intensive (similar to training a model from scratch as far as I understand) so need to rely or ask the model creators to make distilled versions of the models they release. I think 0.5B is a reasonable size but could not explore anything smaller since thats where Qwen 2.5 stopped. I thought more about dynamic layer offload speedup potential too and the numbers I am pushing around seem like the idea can be viable on a 4070 with gen4 PCI (I think a single 4070 might be able to push 5tps gen on Llama 70B which is close to its compute bound potential), but most likely not on a 4090 as the ratio of compute to PCI BW is about 2X on the 4090, it would need gen5 PCI to be viable. Still 2X reduction from peak is a lot better than 10X when trying to compute on CPU with its x10 less or more memory BW. I was thinking about feeding some ggml and the vulkan backend into Qwen 2.5 32B and see if it can suggest how to wedge in dynamic offload concept efficiently. |
Beta Was this translation helpful? Give feedback.
-
Interesting! I think you're right that "distilled" training is likely to be pretty hard to do as you'll need the full outputs to use cross-entropy loss on (I assume they do it this way?). If you are only interested in greedy sampling then one interesting alternative would be to use (smoothed) hinge loss on the "all vs one" targets:
Hinge loss has other potential benefits too:
The only downside is it can sometimes be a bit tricky to get working with gradient descent. |
Beta Was this translation helpful? Give feedback.
-
I dont have backround in ML but my basic understanding is teacher student approaches to build models do this. Particularly Llama 3.1 I believe they trained 405B then distilled the smaller models with subsets of data corpus off 405B (pretty sure about that). I am speculating Qwen did the same with their 2.5 series and they wound up with quite performant small models. So its essentially a full train from nothing using cross entropy loss to minize the error between target and teacher on a desired training corpus. Way beyond the compute potential of users must be done by model creators on GPU farms.
I dont know enough about ML to understand the implications here, though based on my comm theory background it seems analagous to a hard decision decode vs a soft decision decode which can often results in a fairly significant entropy loss in communication systems.
Still reserved about possible entropy loss in the training of the draft....
Thanks for interesting comments! Hopefully more model creators will be creating distilled edge versions of their bigger models which can be easily leveraged to run the bigger models far more effiiciently. It seems like such a no brainer but I dont think it has been done intentionally to date (llama 3.2 series were created for edge apps, not drafting. I also believe smaller Qwen 2.5 were also targeting edge devices and drafting bigger models on heavy lift platforms is just a nice side benefit). |
Beta Was this translation helpful? Give feedback.
-
Based on a quick search: https://www.gaussianwaves.com/2009/12/hard-and-soft-decision-decoding-2/ This would actually be analogous to "0-1 loss" which isn't used in machine learning due not being differentiable. It's the black "square" loss on this graph:
They are all really a surrogate for 0-1 loss, but with different properties. The main appeal of cross-entropy loss is that when it's paired with the binary or multinomial logistic function (aka "softmax"), produces well-calibrated probability estimates. Nobody would consider using least-squares loss for classification now, but back in the 80s and 90s it was very common. The use of hinge-loss was only popularised in the early-mid 2000s for use with Support Vector Machines (and usually solved via the dual quadratic programming problem and using a non-linear kernel). But there isn't any reason you can't solve the primal problem without a kernel (but as I said above you would likely want to use one of the smoothed/Huberised variants due to it being trickier to solve via gradient descent - it tends to get "stuck" easily due to the zero-derivatives in a similar way to stock ReLU can get "dead nodes" [if not initialsed properly] and why ML tends to use smoothed versions of ReLU now instead...). The key difference is if you don't care about the probability estimates then you can use hinge-loss directly on the dot-products of the
The "margin" coming from the little triangle formed by the blue line between 0 and 1 on the previous diagram.
Again, looking back at the first diagram, you can see that cross-entropy loss naturally wants to drive the output corresponding to the true class further and further away (or looking at the second diagram: once it's crossed the margin - it wants to keep pushing it away). What you can't see in the first diagram (due to the reduced x-axis) is the use of Neither of these properties are beneficial if all you care about is predicting the most likely next token. In actual fact, if it were possible the loss we would really like to use here is... The 0-1 loss! This would be the loss that most agrees with our goal of predicting the correct output most often, but since this isn't possible; the use of hinge-loss is likely to be the next best choice! Hopefully this makes sense, and I'm fairly sure I could easily adapt the Unsloth CE-loss kernel to calculate smoothed/Huberised hinge-loss: https://github.com/unslothai/unsloth/blob/main/unsloth/kernels/cross_entropy_loss.py Then it would just be a case of running the larger model over a corpus of stuff you want to predict, but saving the greedy prediction of the larger model at every step as the new ("one-hot") training data for the tiny-models' fine-tuning dataset (I think - not thought about this part too hard). EDIT: This would potentially open up the possibility of training these tiny models on just a subset of the possible tasks you want to perform, eg: just on coding data or coding related data, etc. Likely meaning an even smaller/faster/more accurate tiny model that specialises in just predicting this subset of the larger models' abilities. |
Beta Was this translation helpful? Give feedback.
-
I think it would be interesting to explore both hinge loss with greedy and cross entropy loss with full distribution match error criteria. Small draft models may be within fine-tune training viability even on consumer grade hardware as they can be loaded F16 or F32 so back prop gradient training can work. The spec decode server already gives the needed platform for teacher-student setup. Needed steps might be
It would be an interesting experiment to start with Gemma 2 2B it checkpoint which can't speculate Gemma 2 9B right now and see if it can be improved for use as a speculator. |
Beta Was this translation helpful? Give feedback.
-
|
Just thinking about this some more and wondered how feasible it would be:
I'm thinking along the lines of using the draft model to create a tree (with probabilies on the edges and tokens in the nodes), and then use it to decide on a set of batches for the larger model to generate in parallel. If we constrain the branching factor to a fixed k, then we can again use Hinge Loss to try to pick the top-k using k-vs-all. I don't have a good idea of how the cost of batch processing grows though and it all depends on this. |
Beta Was this translation helpful? Give feedback.
-
The k sequence wavefronts in beam search need to be time aligned. At any time t each iteration of a width k beam search will have k beams all of which must be validated by the target. You could in theory independently advance the k beams forward a block of N sample time steps each, N>1, using the draft, but none of this advance has been validated by the target. Once you start validating each advance with the target (in parallel blocks of N), the first miss in each beam will define how far forward you can advance at that iteration. If one of the beams misses at the first token, which becomes more probable as k increases, the advance for all candidate sequences is limited to one token for all k beams. If this happens on the last beam evaluated, you are forced to throw out all the computation done on the previous beams to result in advancing one step for all beams at that iteration, including all the computation for the length N drafts of k beams and all the parallel validation of that advance by the target on all previous beams. The algorithm I am talking about (and which I use) is O(k*n). At each time step it computes only the k most probable sequences in the search. I am not familiar with any O(k^n) beam search algorithm. |
Beta Was this translation helpful? Give feedback.
-
Yeah, you're correct: I was thinking about k-child breath-first search for decision tree induction which would be O(k^n), but if you have enough threads can be performed for all nodes of the same depth in parallel (so actually O(n) regardless of k). I've probably done a horrible job explaining it, but I was thinking about a much simpler method and not really necessarily needing to use Beam Search:
You could use Beam Search, heuristic truncation as #3624, Best-first Search, or whatever to generate the tree in (1). The only reason for wanting to do k-child breath-first search is that it would be just as easy to train a tiny model using hinge loss for predicting top-k as top-1, and in theory each depth / sequence length could be done in parallel (so O(n) if you can do all sequences at a certain depth in a batch on the smaller model), but my post kinda mixed up these two and I can see where the confusion comes from :) |
Beta Was this translation helpful? Give feedback.
-
|
This does a much better job of explaining what I'm thinking: https://catboost.ai/news/catboost-enables-fast-gradient-boosting-on-decision-trees-using-gpus I don't really know the scaling laws about increasing the batch size for the tiny model like this, but if it were free (which it obviously isn't) then you could compute a tree of all k^n sequences for the same cost as a single auto-regressive generation of n tokens...
Once you have chosen your set of candidate sequences, it's really just a case of seeing if batch processing more than 1 sequence for the larger model in parallel actually gains anything (and using the code in #3624 might make some of the overlapping sequences much cheaper to compute?). |
Beta Was this translation helpful? Give feedback.
-
Yeah thats a different beast than a straightforward beam search for inference. I don't understand all the details of what they are talking about but they do seem to need exponential gradient tree expansion (no leaf pruning) which is a beast. I just feel the target is going to blow out every one of those drafted trees as invalid though for even small values of depth "n" in speculation, particularly on a non pruned exponential tree. Hunch is spec is a complete bust for this type of algo. You got me thinking though. I think it is possible to leverage spec in either 1) adaptive beam search (where decode stays single beam until it hits a low prob token, then transitions into beam search until the most likely path gets a token with prob above adapt threshold at which point it continues single beam again) or 2) front load beam search, where only N tokens at the start of inference are done with k beams, followed by selection of the best beam and single beam from thenon. In both of those case spec can be used effectively any time the algorithm is in single beam mode. Adaptive beam search might be able to assist in training by keeping the target on a higher probability path so the training updates are done with improved data. |
Beta Was this translation helpful? Give feedback.
-
|
Yeah, there seem to be endless possibilities, but it all comes down to the relative cost of batching. |
Beta Was this translation helpful? Give feedback.
-
|
Testing my server rebase for regressions after all the recent changes along with a few new "LRM" (Marco-o1 and QwQ) models and RPC mode also. The spec algo I implemented is greedy match with fixed size draft block and no probs computes. Hardware: RTX4070 GOLDCOIN:
HUMANEVAL 1ST PROBLEM
|
Beta Was this translation helpful? Give feedback.
-
|
How does your implementation compare to the spec approach on master? You can do |
Beta Was this translation helpful? Give feedback.
-
|
Unfortunately I can't run the master spec algorithm due to a large number of other changes I require in the server to support my shell based model loading and inference platform. I expect I will be slightly faster as I don't need any probs computed. Also I made sure to leverage sampling the next token from the target for free if I get hits on all the drafted tokens as the deepmind guys suggested, so I think I am at max theoretical performance for a simple fixed block length spec algorithm. Draft quality along with generation content is going to dominate the spec performance in practice. Just a simple example of how huge this is showing 3X speedup on predicable content and no speedup on unpredictable content. I don't think unpredictable content will be typical in practice though unless forced to be pathologically unpredictable with the prompt. Llama 3.1 8B drafted by Llama 3.2 1B: speceasy.txt: generate the integers from 0 to 100 separated by spaces TIME=1 NDRAFT=8 lm speceasy.txt 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 TIME=1 NDRAFT=0 lm speceasy.txt 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 spechard.txt: generate fifty random integers from 0 to 100 separated by spaces TIME=1 NDRAFT=8 lm spechard.txt 14 73 28 42 91 19 67 85 31 46 13 59 75 22 88 49 62 11 98 35 76 29 43 81 54 17 93 24 69 58 41 95 27 52 79 18 65 39 82 47 21 94 33 60 71 25 89 38 56 97 48 23 64 44 90 16 72 50 86 32 55 83 26 68 37 92 20 57 74 40 63 51 80 34 61 78 30 96 15 99 6 8 4 2 0 1 5 7 3 9 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40 42 44 46 48 50 52 54 56 58 60 62 64 66 68 70 72 74 76 78 80 82 84 86 88 90 92 94 96 98 100 TIME=1 NDRAFT=0 lm spechard.txt 14 73 28 42 91 19 67 85 31 46 13 59 75 22 88 49 62 11 98 35 76 29 43 81 55 24 93 18 69 52 97 38 65 16 82 41 58 94 25 72 48 21 89 33 60 79 15 44 90 27 63 86 51 39 95 20 74 57 83 32 66 47 17 92 40 61 80 26 54 68 36 84 10 45 96 23 71 50 78 34 87 56 64 12 70 53 77 30 92 8 99 6 58 98 4 92 5 91 3 97 2 94 1 96 7 93 9 |
Beta Was this translation helpful? Give feedback.
-
Not sure about the Goldcoin and Humaneval tests, but you can run the integer prompts on # start the server without SD
./bin/llama-server \
-m ../models/qwen2.5-7b-coder-instruct/ggml-model-q8_0.gguf \
--port 8011 --ctx-size 8192 -ngl 99 -fa
# start the server with SD
./bin/llama-server \
-m ../models/qwen2.5-7b-coder-instruct/ggml-model-q8_0.gguf \
-md ../models/qwen2.5-0.5b-coder-instruct/ggml-model-q4_0.gguf \
--port 8011 --ctx-size 8192 -ngl 99 -ngld 99 -fa \
--draft-max 4 --draft-min 0 --draft-p-min 0.0
# send request
curl --request POST --url http://localhost:8011/v1/chat/completions -H "Content-Type: application/json" -H "Authorization: Bearer no-key" -d '{ "messages": [{ "role": "system", "content": "You are a helpful assistant." }, { "role": "user", "content": "generate the integers from 0 to 100 separated by spaces" }], "top_k": 1, "samplers": ["top-k"]}' | jq
Yes, I agree. |
Beta Was this translation helpful? Give feedback.
-
|
@steampunque any update on this? |
Beta Was this translation helpful? Give feedback.
-
I havent had a chance to dig into it yet but I still think its viable particularly for the new 5000 series from NV with PCI gen 5 combined with a supporting motherboard (though I think gen4 x16 is also viable, i.e. I think 4070 or any other x16 gen 4 card can potentially work). Originally I planned to prototype idea in vulkan due to simpler backend and I have also done some graphics related vulkan API programming but I found spec decode does not speed up vulkan at all. In order for the idea to work, decode a batch of 4 needs to be just slightly more time than decode a batch of 1 to net out a max possible speed up factor of ~3. I still do plan to dig into it but decided to hold off due in part to large amount of churn in ggml backend over last many months which I dont feel like rebasing patches to every other day. Since the layer decoding is a serial process by nature all that is really needed is enough VRAM to hold one active compute layer and one transfer layer which are ping ponged. Hidden state and KV for each layer stay also in VRAM. Output layer may want to stay in VRAM also. Should completely obsolete need for RPC on models which can be speculated as long as host has enough main memory. Llama and Qwen both speculate extremely well but Gemma is an example which I found cannot speculate well so it is not universally applicable. Deepseek R1 spec with latest version looks good so I'm still encouraged. I'm pretty sure the speculated Deepseek R1 32B is inferencing correctly since I recently benched it at 93.0 AVG on Hendryks math levels 1 to 5 using my spec algo also with the R1 1.5B draft as compared to the 94.3 published by Deepseek which is most likely not quantized at all. Spec algo : greedy match with fixed size draft block and no probs computes. Hardware: RTX4070 GOLDCOIN:
|
Beta Was this translation helpful? Give feedback.
-
I recently added an efficient greedy-only spec decode to my downstream server patch (a completely different implementation than the current spec decode PR). I then evaluated tg performance for two cases : 1) Solve the first humaneval problem with coding model and 2) solve the goldcoin problem with general model. I used Qwen 14B for the target and 0.5B, 1.5B, and 3B for the drafts. I evaluated tg vs. draft token length on a 4070 fully offloaded with the target and draft weights where target is IQ4_XS quant and draft is Q6_K quant.
HUMANEVAL first problem:
TARGET Qwen2.5-Coder-14B-Instruct
DRAFTS Qwen2.5-Coder-0.5B-Instruct, Qwen2.5-Coder-1.5B-Instruct, Qwen2.5-Coder-3B-Instruct
TPS vs draft tokens:
GOLDCOIN
I have 10 apples. I find 3 gold coins in the bottom of a river. The river runs near a big city that has something to do with what I can spend the coins on. I then lose 4 apples but gain a gold coin. Three birds run into my path and drop 6 apples each. I play an online game and win 6 gold coins but I have to share them equally with my 2 teammates. I buy apples for all the coins I have. The price of an apple is 0.5 coins. How many apples do I have? And where is the river? Use step-by-step reasoning to solve this problem.
TARGET Qwen2.5-14B-Instruct
DRAFTS Qwen2.5-0.5B-Instruct, Qwen2.5-1.5B-Instruct, Qwen2.5-3B-Instruct
TPS vs draft tokens:
TARGET Llama 3.1 8B Instruct
DRAFT Llama 3.2 1B Instruct
TPS vs draft tokens:
TARGET Gemma 2 9B it IQ4_XS
DRAFT Gemma 2 2B it IQ4_XS
TPS vs draft tokens:
Results Summary:
Coding shows a max speedup of 2.5x tg at 10 draft tokens speculated using 0.5B model. At 1.5B draft the max speedup is 1.63x at 4 draft tokens. At 3B draft the max speedup is 1.33 at 4 draft tokens. The efficiency crossover (where draft+target is the same as no draft) is >32 draft tokens for 0.5B, >16 draft tokens for 1.5B, and 11 draft tokens for 3B.
Goldcoin shows a max speedup of 1.4x tg at 4 draft tokens speculated using 0.5B model. at 1.5 draft the max speedup is 1.17x at 4 draft tokens. At 3B draft the max speedup is 1.08 at 1 draft token. The efficiency crossover (where draft+target is the same as no draft) is 12 tokens for 0.5B, 6 tokens for 1.5B, and 3 tokens for 3B.
With Llama 3.18B instruct drafted by Llama 3.2 1B instruct a speedup in token gen of 1.83x is found at draft tokens of 5.
With Gemma2 9B it drafted by Gemma2 2B it there is never any speculative decoding speedup. Guess 2B not distilled from 9B at all but was trained on a completey different data set.
Conclusions and potential for running big LLMs on consumer grade GPUs:
Small draft model is needed (sine qua non). 0.5B size seems to work well. Any model in the range of 8G or above can benefit by distilling a 0.5B draft and speculating the model. Returns fall off rapidly as draft gets bigger, already questionable at 1.5B and not really useful at 3B draft. Coding is far more efficient than general text gen with speculation. Qwen 2.5 series is perfect for exploiting the potential of speculation.
For running big LLMs on consumer grade GPUs with limited memory it is desired to avoid the need to store all model weights and output layer in VRAM because there is not enough room. Most of the model weights are sitting there doing nothing most of the time, i.e. a 32 layer model has 31 dead weights sitting there occupying VRAM doing nothing 31/32 of the time. To get around this problem it is necessary to dynamically swap the layers into VRAM as they are needed from CPU RAM which is normally much higher capacity. If the draft size at the efficiency crossover is big enough, there may be (emphasis on may, it needs to be investigated for feasibility) enough time to compute the target batch (say 8 to 10 samples) and simultaneously transfer the next layer into the GPU. The GPU capacity needs 1 working layer allocation and one transfer allocation (two total model layers which are ping ponged between compute and transfer) + a fully offloaded speculator. KV for speculator and target should also both be in GPU mem. Even if it is needed to go above the efficiency crossover, it can still be more efficient to do dynamic layer loading to GPU because offloading to CPU is an immediate 10X or higher slowdown due to memory BW limits.
Beta Was this translation helpful? Give feedback.
All reactions