Repository navigation
Enhancement: Speculative decoding – load 2 models at the same time! #1207
Description
ggml-org#2926
ggml-org#3624
ggml-org#5625
Feature request
- Implement speculative sampling in koboldcpp
- Implement optional loading of 2 models which is required by speculative decoding
- Also allow switching the main model during normal generation: the client can choose, whether the current request should be:
- Generated speculatively using both main and draft models?
- Generated using only the main large model without drafting?
- Generated using only the small model as if it was a separate one?
- Context cache is independent for those two models when the speculative sampling is disabled in Lite, but gets synched on the next speculative request.
Background
I have been using MIQU model (https://huggingface.co/miqudev/miqu-1-70b) for quite a long time many months ago. It is 70b and of course it won't fit in my 3060 GPU with 12 Gb of VRAM.
That model was better than anything else I've ever tried! I wouldn't want to run it heavily quantized to 2 bits because I didn't want to sacrifice its quality, especially because I have 128 Gb of DDR4 RAM.
I could get, like, 1 token/second (or slightly more while the context is short) by activating CuBLAS with 0 offloaded layers.
But later, llama.cpp was updated with new quantization algorithm that hurts performance for older models (that have to be requantized, which is not the case for this stolen/unofficial miqu model). Anyway, it was not that bad even when I continued to play with miqu.
But recently a Mistral Large 2 model came out (https://huggingface.co/bartowski/Mistral-Large-Instruct-2407-GGUF) that has 123b parameters!
For me it was superior to miqu for every possible task, and it is even less censored.
Unfortunately, such huge model is running 0.3 tokens/second from the empty context and it gets even slower over time…
I tried different attempts to speed it up, but CuBLAS with 0 layers is still the best (and I cannot roll back to the older koboldcpp version because of GGUF format changes, to see if the previous CUDA kernel versions might be faster or not).
Q4 quant instead of Q5 gives a slight improvement: 0.4 tokens/seconds (+0.1 comparing to Q5).
After searching information about which model can be used as a draft model for speculative sampling for Mistral Large 2 I decided to try Mistral 7B Instruct v0.3 (https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF)
Strangely enough, llama.cpp has some redundant vocabulary checks (https://github.com/ggerganov/llama.cpp/blob/f018acba22095b8995bf6c5ef815b16a3ce4cf1b/examples/speculative/speculative.cpp#L119-L136). I had to recompile from source with those asserts commented out to make it accept Mistral 7B Instruct v0.3 (as Q5_K_M) as draft model for Mistral Large Instruct 2407 (as Q4_K_S).
Also I had to build with full CUDA support for a fair comparation.
The final speedup was huge! At 3-5 drafted tokes I got doubled speed of 0.85-1.0 token/second!
I think it worth to have it in koboldcpp as well.
What exactly I propose
- You need a separate model loader on a dedicated tab in config GUI. There a used can set the device and layers offload strategy. (I don't know technical details, for example how much of configuration have to be equal for both models for the sake of speculative sampling to work; but I imagine you want to get free control for anything that can be different for the draft model).
- All critical parameters should be read from the main model, like the context size (at worst the drafting degrades, but will not break the result provided the speculative sampling was implemented correctly).
- Add a new dropdown for Lite sampling tab: essentially asking what model to use: default/speculative, large/main, or small/draft. Default means "allow drafting if enabled", while other options would tell koboldcpp to not enable speculative decoding.
- When the client asked to not use speculative, koboldcpp proceeds with a normal generation, but with the chosen model (whether it is main or draft).
- Looks like you need to deal with two independent context caches, and synchronize them together only when the speculative drafting is requested. Meaning, if one "user" generates with only the main model, the other user can later generate with the draft model without destroying the first user's context cache.
- Default to speculative sampling for unknown clients who do not pass the new field over the API, so that they can benefit from the improved speed anyway.
- Think of some heuristics for the number of drafted tokens. For example, if the draft model heavily agrees with the main model during several steps, koboldcpp can increment the draft amount (up to the specified value in server config; default to e.g. 8); and otherwise decrement it if the last two drafted tokens were discarded enough times (down the other specified value, defaulting to minimal possible 2).
- Try to accept any model as a draft, even if it is not quite compatible. In case it is not possible to use a declared draft model also as a normal model in current llama.cpp – then some upstream changes would be necessary too.
- Probably, ContextShift would not work; maybe something else would not work too. Since this mode is completely optional, it won't hurt anyone who don't need it.
Sampling
Only greedy sampling (temp=0 or top_k=1) is straightforward to implement for speculative decoding. Though, some algorithms exist to allow stochastic sampling from several token probabilities (I'm not quite sure how it is implemented in llama.cpp: are they generate a most probable depth tree recursively? Are they just estimating output probabilities, sacrificing the authenticity of main model actual logits?)
Here I suggest to live with whatever is implemented in llama.cpp. Even if only greedy sampling would work correctly – this would be still a huge improvement, because:
- Large models are very confident in tokens: the latest update with logits list shows that Mistral Large (and miqu) most of the time returns 100% probabilities even with "sane" sampling parameters (top_p<=0.9, temp<=0.9), so de-facto it tends to behave as if we are already sampling greedily. This can be easily confirmed by multiple retries at any moment, and the model will basically say the same thing each time.
- Other than roleplay, sometimes a user might want to "have a local ChatGPT" and asking it questions (other than "come up with a story"). In those cases, it is pretty normal to sample with zero temperature, for many tasks like text translation, summarization, source code generation and answering factual questions.
- With a runtime setting to switch to the draft model, the user can just switch and retry many times whenever the story becomes boring! Since the draft model is both fast and smart (to be actually useful to be a drafter for the large one), its genuine answers would not be that bad.
Use cases!
- Speculative decoding for improving speed without compromising the quality of output texts, provided the user has extra VRAM/RAM. Basically, you add a correct drafter model and transparently get a noticeable speedup!
- Comparation of two unrelated models, ability to switch between them as if they are loaded in two instances of koboldcpp, but more conveniently. You write a story and then change the model to see how different it would be (especially useful to test finetuned versions by comparing their logits in the middle of a good story).
- Preserving the context cache when running two stories simultaneously in two Lite tabs, again as if they are loaded in two separated servers.
- Performing small tasks like memory summarization or image prompt expansion using the draft model while running a large story (with or without speculation), because the small model should be quick enough to reprocess the whole context, which is not the case for the large model.
I see another improvement that technically will be possible if everything is implemented: the ability to use two cache contexts while still running one model:
- The user selects something like "use the draft slot only as a separate context cache" when starting the server
- The speculative mechanism is not instantiated
- Instead, its separate context cache assigned to the same main model
- In the default drafting mode koboldcpp behaves as if the user choses to generate on the main model without speculation
- When Lite asks specifically to use the draft model – only its context cache slot is used, without destroying the main one
- This would allow to run two stories together with one model in one instance of koboldcpp, which is cheaper than running two instances on different ports
- If this idea becomes popular, you may turn in to arbitrary number of contexts instead, totally separating this logic from the speculative stuff
There are two things need to be done: 2 (or more) models and contexts at the same time, and speculative decoding using 2 models.
If you would implement several models – then you can rather easily add the speculative decoding too.
Otherwise, if you would want speculative decoding – you would have to implement loading of several models for this anyway.
Then, when having two models in memory – you can imagine something like "model offloading", or "switching on demand", where a model my be unloaded and replaced with another model at runtime.
But those are future possible improvements, while the speculative decoding is a useful thing by itself!
Activity
stepfunction83 commentedon Nov 25, 2024
Support for this was added to llama.cpp today:
ggml-org#10455
It shows 40-60% speedups in practice.
LostRuins commentedon Dec 1, 2024
Support is now added to KoboldCpp. Do give it a try!
aleksusklim commentedon Dec 1, 2024
How do I confirm whether it is working or not?
It prints
Attempting to load draft model for speculative decoding. It will be fully offloaded if possible. Vocab must match the main model.While I know that Mistal tokenizer is partially different between Large and 7B, there were no errors following.
It says it offloaded 33/33 layers (while I had 0 chosen for CuBLAS). Are you sure we don't need a cap configured?
During generation, I see 0% load on GPU, suddenly spiking to 40-60% periodically. Is that the draft model running?
Also, "Generating (X / Y tokens)" is incremented with more than 1 step.
But where is statistics of speculative sampling? For example, "drafted 40 tokens, accepted 29 tokens, in average 3.4 tokens for step"
So that I can judge, how well the drafting is going.
Still, can you add a flag to switch to the draft model at runtime? (Al least that: not actually use two models as independent, but simply force generation on a small model instead of running both speculatively)
This way it will be possible to quickly compare, what the drafting model "thinks" about your story in general, and how fast it goes. For example, if its output looks like nonsense – then something is wrong (too long context, unknown input language, foreign template format).
LostRuins commentedon Dec 1, 2024
You can run it in --debugmode. The draft model output will be displayed along with the validated result and whether it matches.
To verify drafting you can compare with two kinds of instructions:
"Please give me the first 100 positive integers" - draft model will do very well on this at low temperature. You will see large chunks of correctly guessed tokens being output
"Write a funny story about zebras" - draft model will do poorly on this as it's hard to speculate creative output.
aleksusklim commentedon Dec 1, 2024
Oh, Debug Mode, should have thought of that!
It's working, but it renders UTF-8 characters in system locale I think? For me it's win-1251.
Trying to generate some Russian texts, and it prints in console two-character sequences instead of Cyrillic letters.
I even tried to execute
chcp 65001before launching koboldcpp, but it gave no effect.Windows 10 console is capable of rendering Unicode characters (for example
chr(1044)prints'Д'and notР”in python 3.10) as long as the corresponding glyph is present in the chosen font.Or are you afraid of decoding Unicode at that point?
UPD: when the full generated text is printed on console at the end as a string, it renders correctly.
LostRuins commentedon Dec 2, 2024
Unicode characters are often made of multiple tokens. For example,
けis made of0xe3 0x81 0x91. In the case where the AI only generates the token for0xe3for example, it will not be rendered as what you expect.You should cross reference the token IDs with the byte representations in the vocab. It's more complicated than you think.
For example,
しけis made of the bytes0xe3 0x81 0x97 0xe3 0x81 0x91. When using Qwen2.5, this tokenizes to [0xe3 0x81 0x97 0xe3 0x81,0x91]. None of the subsequences make sense on their own.aleksusklim commentedon Dec 2, 2024
Oh, I got it. You print individual tokens up to the first failed guess but not more.
While actually you have all N drafted tokens that you got from the small model!
To be able to visually compare texts, it is enough to just utf8-decode whatever we have from the draft model so far, not truncating it on the very first wrong token. (Yeah, if the failed one was the very last – we get the same exact thing: nothing known after it)
If you would additionally print the full drafted string, it will be trivial to compare what was generated against "what the smaller model was trying to say"!
LostRuins commentedon Dec 2, 2024
sure. you can print the full drafted tokens by printing out the token IDs here:
https://github.com/LostRuins/koboldcpp/blob/concedo/gpttype_adapter.cpp#L694
I have a helper function
print_tok_vec_strwhich you can use, just pass in the draft ID array.aleksusklim commentedon Dec 4, 2024
@LostRuins, I have mistaken.
chcp 65001in the console before launchingkoboldcpp.exeactually solves the gibberish rendering!(I have no idea how my previous test didn't work out; now I've rebuilt from source inside
w64devkitshell, then noticed that from there I saw a different kind of gibberish, not cp1251; then tried to runpython koboldcpp.pyin CMD to compare, and then tried executingchcp 65001before which suddenly make it correctly render Unicode! Surprised, I tried that again with an official .exe – and it worked too, changing printed tokens to utf8)This is what I tried to add into the code, as you suggested:
Here are full logs of a run (Mistral-Large-Instruct-2407 + Mistral-7B-Instruct-v0.3):
CMD shell
The drafted part of
</s>столько я не знаюlooks very funny, since I've asked it to spell the Russian alphabet, and after the EOS it tried to say "I don't know that many" as if the model felt like it was forced to keep printing letters!So, two final questions:
\nin string though, probably it is not actually needed)LostRuins commentedon Dec 5, 2024
5106816
aleksusklim commentedon Dec 5, 2024
For 1: most of the time your .exe is run via Explorer and not from a terminal, and thus it is your process who creates and ultimately owns the console. But checking for this on Windows is even more inconvenient, so I agree that it is easy enough to just create a .bat file that would run chcp and then koboldcpp.
For 2: get_tok_vec_str prints a list of tokens and ids in parenthesis, not a "text"! The purpose was to easily read it…
aleksusklim commentedon Dec 21, 2024
Weird, now
chcp 65001is not working anymore… Neither with v1.80, nor with older.But… how!? I am 100% sure the last time it was just working as-is for me. Now I've tried running from .bat file, from simple CMD, from elevated admin prompt, under w64devkit, from Powershell, and even with
$OutputEncoding = [System.Console]::OutputEncoding = [System.Console]::InputEncoding = [System.Text.Encoding]::UTF8$PSDefaultParameterValues['*:Encoding'] = 'utf8'(https://stackoverflow.com/q/51933189)
I see the same utf8 byte-gibberish in console again.
But then I clicked "Extra → Unpack KoboldCpp To Folder" and then tried
chcp 65001+python koboldcpp.pythere.Now it works, printing correct characters!
LostRuins commentedon Dec 22, 2024
Like I said before, it's likely a setting from the terminal console itself and not within koboldcpp. Different terminals can have different settings for text encoding.