[Model] MLPSpeculator speculative decoding support

[JRosenkranz]

This is a PR to include MLPSpeculator from fms. This feature will improve inference performance significantly with minimal effort from the user by leveraging the existing speculative decoding framework with an improved speculator architecture implementation. The architecture itself is similar to medusa, with a key difference that each head conditions its outputs on the prior head. This allows it to produce better formed n-grams and increase acceptance rate. A more comprehensive description of the speculator can be found in the following technical report.

Current Speculators:

• llama3-8b-acclerator
• codellama-13b-accelerator
• llama-13b-accelerator
• granite-7b-instruct-accelerator
• granite-20b-instruct-accelerator

Design:

MLPSpeculator: A speculator model in modeling code which each forward pass corresponds to a single layer of the MLPSpeculator -> implement forward via single steps of the speculator.
HiddenStatesWorker: A special form of Worker which will save the hidden states -> will contain an instance of MLPSpeculator as well as the base model. Overrides the Worker execute_model method and sets the previous hidden states in the MLPSpeculator
SingleStepSpeculativeModelRunner: A model runner which will run a single step of the speculator while also only computing the inputs a single time.

This PR will leverage the work already done in spec_decode_worker. For this speculator -

proposer_worker - MultiStepWorker which is backed by an MLPSpeculator
scorer_worker - HiddenStatesWorker

This PR had a few main challenges that were addressed in the following way:

A. Implement forward in speculator such that it knows which head to be using

Because we wanted to re-use the multi_step_worker, we opted to make the MLPSpeculator model its own model in modeling code which would only return one hidden states per iteration in the multi_step_worker. To do this, we had to do the following:

1. Each time a new hidden state comes in to the HiddenStatesWorker set a counter (corresponding to the current head in the speculator) to 0
2. Each forward pass of the speculator within the multi_step_worker will increase that counter by 1 and call the proper head
3. Forward will return the hidden states and save them for the next head (previous_hidden_states)

B. Pass the hidden states to the MLPSpeculator from the scorer_worker

Our speculator requires the base model hidden states as input, which is not currently supported in vLLM. For this, we created a new HiddenStatesWorker which acts similarly to a regular worker, but instead also saves the previous hidden states within the MLPSpeculator for later use. This is done in the spec_decode_worker when score_proposals is called (in turn calling execute_model). Aside from the previous_hidden_states, this will also update any other information required for MLPSpeculator such as the context length per request id as well as the current head counter described in A.1

C. The hidden states must be pruned for next get_spec_proposals call

Because the hidden states from the base model come to the speculator unpruned, we must prune them prior to generating our new candidate tokens. Pruning will happen by getting the number of accepted tokens per request using the prior forward’s context length and the current context length (saved as instance variables in MLPSpeculator)

Using the offline example:

batch_size=1

Non-Speculative: 100%|██████████| 1/1 [00:00<00:00,  1.62it/s, Generation Speed: 86.03 toks/s]
Speculative: 100%|██████████| 1/1 [00:00<00:00,  2.52it/s, Generation Speed: 133.83 toks/s]

[cadedaniel]

This is awesome! I will take a look today or tomorrow.

[cadedaniel]

Thanks. I am still reviewing. Initial thoughts:

• We want to eventually support chunked prefill + spec decode + MLPSpeculator. This means that how we determine if we're in a prompt run or not will change, we should make sure the design fits it (right now I think it's missing something).
• Do we need to re-use the multi step worker? I would support having one specialized to MLPSpeculator. The reason why is there is a lot of overhead in each step of the multi step worker (prepare inputs). If we have a custom MLPSpeculatorWorker (or similar), we can skip that for subsequent steps, and potentially even cuda graph them in the future.
• TODO see if there's a better way to get hidden states from worker via alternative state tracking.

[tdoublep]

@cadedaniel I did a pass through and refactored things a bit to try a reduce the overhead - I now introduced a standalone MLPSpeculatorWorker so we don't need to use MultiStepWorker at all here.

Here are some performance results for llama-2-13b:

Without speculation (5 repetitions, batch size 1):

Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.73s/it, Generation Speed: 42.32 toks/s]
0.023635914325714113
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.73s/it, Generation Speed: 42.30 toks/s]
0.023645806312561034
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.73s/it, Generation Speed: 42.29 toks/s]
0.02365394473075867
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.73s/it, Generation Speed: 42.27 toks/s]
0.023662294149398803
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.74s/it, Generation Speed: 42.21 toks/s]
0.023697214126586916

We see ~23ms per token latency and 42 toks/s throughput.

With speculation (5 repetitions, batch size 1, before refactor @ 6ba9a1e):

Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.53s/it, Generation Speed: 78.93 toks/s]
0.01267344355583191
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.54s/it, Generation Speed: 78.90 toks/s]
0.01267833948135376
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.54s/it, Generation Speed: 78.82 toks/s]
0.01269181728363037
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.54s/it, Generation Speed: 78.78 toks/s]
0.012697651386260986
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.55s/it, Generation Speed: 78.54 toks/s]
0.012735981941223145

We see ~12.7 ms per token and ~78.8 tok/sec throughput.

With speculation (5 repetitions, batch size 1, after refactor @ bf2f102):

Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.50s/it, Generation Speed: 80.15 toks/s]
0.012481423616409302
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.50s/it, Generation Speed: 80.10 toks/s]
0.01248981237411499
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.49s/it, Generation Speed: 80.20 toks/s]
0.012472513914108276
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.50s/it, Generation Speed: 80.12 toks/s]
0.012485114336013793
Processed prompts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.50s/it, Generation Speed: 80.05 toks/s]
0.01249562382698059

We see ~12.4 ms per token latency and 80.1 toks/s throughput.

• We see a small but non-neglible improvement from refactoring via MLPSpeculatorWorker.
• These speed-ups of nearly 2x over without speculation are pretty consistent with what we've seen from this speculator in our past experiments.

[LiuXiaoxuanPKU]

• TODO see if there's a better way to get hidden states from worker via alternative state tracking.

I have the same concern about this. I feel synchronizing the HiddenStatesWorker and Worker's execute model logic might introduce some overhead. Is it possible to reuse the Worker, maybe add some variables to track if we need the hidden states?

[LiuXiaoxuanPKU]

Performance numbers are really promising!
Just went a first pass, will take another pass after merging the conflict.
On the high level, why do we need to introduce the MLPSpeculatorModelRunner?

[JRosenkranz]

a few more review comments to consider

[JRosenkranz]

Without this, we are assuming the set of top_k_tokens_per_head are each 1 for every head. We may want to see how this affects performance (acceptance rate) as typically we create a tree and traverse the tree to find the optimal candidate.

[Alvant]

So, n_candidates is also not used? (it is assumed to be equal 1?)

[JRosenkranz]

@LiuXiaoxuanPKU

On the high level, why do we need to introduce the MLPSpeculatorModelRunner?

The MLPSpeculator takes different input than just a normal model in VLLM. It doesn't require any of the metadata required for paged attention, and also calls a different method from the MLPSpeculator. This was created to reduce the overhead of requiring all the parameters be computed for forward

We may be able to just merge this all into the MLPSpeculatorWorker though and remove the class

[cadedaniel]

Going to give another pass tomorrow.

[njhill]

Thanks @cadedaniel! I'll be pushing some updates shortly, best to wait for that before looking again.

[cadedaniel]

let me know when it's ready @njhill

[LiuXiaoxuanPKU]

Will give it a review tonight/tmr.

[LiuXiaoxuanPKU]

LGTM, thanks!

[LiuXiaoxuanPKU]

Impressive! It seems MLPSpeculator + w/o cudagraph is still much faster than Original + cudagraph.

[jamestwhedbee]

FWIW, I just tested llama3-8b-acclerator out on an AMD MI100 using the openai entrypoint. At first, I was confused and disappointed looking at the throughput metrics (token/sec looked worse) until I found #5708.

Turns out chunks/second was being measured and looked worse but tokens per second were ~1.67x!

Do you all have any plans for a Llama 3 70B accelerator? I would test this on MI300X.

[cadedaniel]

@jamestwhedbee this is great. Do you know what batch size you used and the category of prompt?

[sahilsuneja1]

@jamestwhedbee We have a preliminary version of the llama3-70b speculator, achieving a ~2x logical speedup (would be lesser in practice). We are working on improving it 🤞

[jamestwhedbee]

@jamestwhedbee this is great. Do you know what batch size you used and the category of prompt?

@cadedaniel so at first I was using https://github.com/ray-project/llmperf at various batch sizes with the sonnet dataset. But I now wonder if I was getting incorrect tokens/second data from this approach (or maybe it was correct but as you will see the speculative metrics in the vLLM logs were also surprisingly low here). I don't think it is representative of the type of prompts I'd care about anyways.

Example command

python token_benchmark_ray.py --model "meta-llama/Meta-Llama-3-8B-Instruct" --mean-input-tokens 550 --stddev-input-tokens 150 --mean-output-tokens 150 --stddev-output-tokens 10 --max-num-completed-requests 8 --timeout 600 --num-concurrent-requests 1 --llm-api openai

And a representative log I saw doing this

Speculative metrics: Draft acceptance rate: 0.174, System efficiency: 0.286, Number of speculative tokens: 4, Number of accepted tokens: 130, Number of draft tokens tokens: 748, Number of emitted tokens tokens: 267.

I didn't understand what was going on yet so I also played around with the number of speculative tokens

 python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-model ibm-fms/llama3-8b-accelerator --num-speculative-tokens 1 --use-v2-block-manager --enforce-eager

which improved acceptance and "tokens/second" slightly

Speculative metrics: Draft acceptance rate: 0.347, System efficiency: 0.674, Number of speculative tokens: 1, Number of accepted tokens: 194, Number of draft tokens tokens: 559, Number of emitted tokens tokens: 753.

I switched to manually prompting single batch for detailed descriptions of historical events and started seeing better acceptance rates and this is when I noticed that even though "tokens/second" was lower, the actual stream was faster.

Speculative metrics: Draft acceptance rate: 0.528, System efficiency: 0.486, Number of speculative tokens: 4, Number of accepted tokens: 412, Number of draft tokens tokens: 780, Number of emitted tokens tokens: 474.

I'm very interested in understanding per-user tokens per second in higher batch sizes as well. I am worried about measurement errors though. When #5742 lands, I'd be happy to run some more robust testing.

[cadedaniel]

@jamestwhedbee thanks for all the details. I am suspicious of these results as well. I expect the MLPSpeculator to have better acceptance rates than 0.5 based on my experience training draft models. Could you run on a sharegpt dataset?

cc @njhill there may be a correctness issue (draft acceptance rate seems too low)

also there is a spec decode vLLM slack channel. if you want to join you can share results there.

[cadedaniel]

talked offline with @njhill . we are missing e2e tests for this PR. we'll follow up to see if we can quickly fix or if we have to revert.

[jamestwhedbee]

@cadedaniel sure do you have a link handy? I can't find it.

[cadedaniel]

@jamestwhedbee send me an email at cade@anyscale.com

[WangErXiao]

How can I get the accelerator model for a existed model?

[haichuan1221]

@jamestwhedbee We have a preliminary version of the llama3-70b speculator, achieving a ~2x logical speedup (would be lesser in practice). We are working on improving it 🤞

Can we test the llama 3 70b speculator now?

[Alvant]

How can I get the accelerator model for a existed model?

According to the paper, here is the source code for accelerator training. However, as far as I understood (from the training params docs here and code here), only Llama models are supported for acceleration...

[JRosenkranz]

How can I get the accelerator model for a existed model?

According to the paper, here is the source code for accelerator training. However, as far as I understood (from the training params docs here and code here), only Llama models are supported for acceleration...

There is no limitation on model architecture. We have an experimental branch where we have done training with GPTBigCode as well as Mixtral here. Feel free to take a look, however it is still experimental.

For the above speculators, granite-20b-code-instruct-accelerator is a speculator for a base model with GPTBigCode architecture.

[JRosenkranz]

@jamestwhedbee We have a preliminary version of the llama3-70b speculator, achieving a ~2x logical speedup (would be lesser in practice). We are working on improving it 🤞

Can we test the llama 3 70b speculator now?

@haichuan1221 We are working on releasing it. Will keep you updated as to when it is available. We still need end-to-end verifications, but results look promising!

[haichuan1221]

@jamestwhedbee We have a preliminary version of the llama3-70b speculator, achieving a ~2x logical speedup (would be lesser in practice). We are working on improving it 🤞

Can we test the llama 3 70b speculator now?

@haichuan1221 We are working on releasing it. Will keep you updated as to when it is available. We still need end-to-end verifications, but results look promising!

I am looking forward for it!

[sahilsuneja1]

@haichuan1221: you can find the speculator for llama3-70b here: https://huggingface.co/ibm-fms/llama3-70b-accelerator

[hustxiayang]

@haichuan1221: you can find the speculator for llama3-70b here: https://huggingface.co/ibm-fms/llama3-70b-accelerator

hi, it seems that "model.safetensors " is not complete and just one shard.

[sahilsuneja1]

hi, it seems that "model.safetensors " is not complete and just one shard.

@hustxiayang, that's strange. I tested the model.safetensors version using this, and it worked ok.
Does, pytorch_model.bin ckpt version work for you?

@njhill I tested using FMS setup, maybe vLLM setup doesn't like auto-generated model.safetensors? Could you please confirm when you have time?

[njhill]

@hustxiayang how long are the prompts you’re using? There was an issue where the spec decoding gets disabled when the num of tokens is > 2k. This is now fixed and will be in the upcoming 0.5.5 release.