[Feature][Spec Decode] Simplify the use of Eagle Spec Decode

[ShangmingCai]

This PR tries to implement #11943.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[ShangmingCai]

@LiuXiaoxuanPKU Hello, I saw the RFC you wrote, so I took some time to implement this PR.

According to the conversion script, we have to transform config.json to wrap it inside

{
    "model_type": "eagle",
    "model": {$origin_json_content}
}

so that it can be considered an EAGLEConfig. So, I made some changes inside vllm/config.py. First, I tried to implement it in the create_worker of SpecDecodeWorker, but it turned out that some config options were wrong, so we had to do this earlier and implement it in vllm/config.py.

Also, since we have to copy and load the lm_head.weight of the scorer model into the eagle ckpt so that it can be loaded in vllm/model_executor/models/eagle.py. I thought about it and felt it would be more appropriate to implement this in the init_device of SpecDecodeWorker.

Considering that the scorer model may enable tp, so I use tensor_model_parallel_gather to collect and concat weights first. I have run some experiments with different tp and verified the correctness. But I only test it with EAGLE-llama2-chat-7B and Llama-2-7b-chat-hf.

Do you think we have to add more tests? It appears that the conversion script is not enough for EAGLE-Qwen2-7B-Instruct, I am seeing this error:

ValueError("Found bias in the loaded weights but the model config doesn't have bias")

[ShangmingCai]

BTW, this PR is compatible with the old method with the conversion script. The converted config.json and weight binary can still be loaded successfully. After testing, the Draft acceptance rate remains consistent for both methods with EAGLE-llama2-chat-7B.

[LiuXiaoxuanPKU]

@LiuXiaoxuanPKU Hello, I saw the RFC you wrote, so I took some time to implement this PR.

According to the conversion script, we have to transform config.json to wrap it inside

{
    "model_type": "eagle",
    "model": {$origin_json_content}
}

so that it can be considered an EAGLEConfig. So, I made some changes inside vllm/config.py. First, I tried to implement it in the create_worker of SpecDecodeWorker, but it turned out that some config options were wrong, so we had to do this earlier and implement it in vllm/config.py.

Also, since we have to copy and load the lm_head.weight of the scorer model into the eagle ckpt so that it can be loaded in vllm/model_executor/models/eagle.py. I thought about it and felt it would be more appropriate to implement this in the init_device of SpecDecodeWorker.

Considering that the scorer model may enable tp, so I use tensor_model_parallel_gather to collect and concat weights first. I have run some experiments with different tp and verified the correctness. But I only test it with EAGLE-llama2-chat-7B and Llama-2-7b-chat-hf.

Do you think we have to add more tests? It appears that the conversion script is not enough for EAGLE-Qwen2-7B-Instruct, I am seeing this error:

ValueError("Found bias in the loaded weights but the model config doesn't have bias")

Thanks for the PR!! The place you load weights looks good to me. Yeah I agree it's more appropriate to implement this in the init_device of SpecDecodeWorker.

Some questions/comments here:
(1) Yeah, could you add simple tests here? A simplest test is just to test the weights of the proposer model's lm_head is the same as the weights of the scorer's lm_head. You can also add more e2e tests if you want.
(2) For Qwen2, I think that's the problem of the conversion script, which is a separate issue, we can fix it in the later PR.
(3) Could you also update the doc saying user does not need to convert the model themselves.

After (1) and (3), the PR is good to go! Thanks for the contribution!

[ShangmingCai]

(1) Yeah, could you add simple tests here? A simplest test is just to test the weights of the proposer model's lm_head is the same as the weights of the scorer's lm_head. You can also add more e2e tests if you want.

@LiuXiaoxuanPKU Sure. A simple weight assertion test is added in tests/spec_decode/test_spec_decode_worker.py.

But while I am adding some e2e tests, this error

ValueError("Found bias in the loaded weights but the model config doesn't have bias")

keeps popping up. Maybe the e2e tests can be added when this separate issue is fixed. You can ping me when it is done.

Also, I think we can print a warning and skip the bias load instead of throwing the ValueError for EAGLE. Maybe the ValueError should be thrown if and only if bias is required in the config but not found in the loaded weights.

(3) Could you also update the doc saying user does not need to convert the model themselves.

Done! However, I am not sure if I should remove the info about the existence of the conversion script. If users are using the earlier version of vLLM, maybe they need this script. Considering that using the script is an inelegant solution, I remove the info in the end and suggest the user update the version of vLLM and contact the community if problems still occur. Please inform me if I should add it back.

[LiuXiaoxuanPKU]

LGTM, thanks!

[LiuXiaoxuanPKU]

Suggested change

able to be loaded and used directly by vLLM after [PR 12304](https://github.com/vllm-project/vllm/pull/12304). If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.

able to be loaded and used directly by vLLM after [PR 12304](https://github.com/vllm-project/vllm/pull/12304). If you are using vllm version before [PR 12304](https://github.com/vllm-project/vllm/pull/12304), please use the [script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model, and specify `speculative_model="path/to/modified/eagle/model",`. If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.

[ShangmingCai]

Done.

[LiuXiaoxuanPKU]

I'm trying this PR, but it seems there are some minor errors:

code:

llm = LLM(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    num_speculative_tokens=3,
    speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
    speculative_draft_tensor_parallel_size=1,
)

error:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/data/lily/vllm-eagle-simplify/examples/offline_inference/example_eagle.py", line 7, in <module>
[rank0]:     llm = LLM(
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/utils.py", line 1051, in inner
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/entrypoints/llm.py", line 242, in __init__
[rank0]:     self.llm_engine = self.engine_class.from_engine_args(
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/engine/llm_engine.py", line 484, in from_engine_args
[rank0]:     engine = cls(
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/engine/llm_engine.py", line 273, in __init__
[rank0]:     self.model_executor = executor_class(vllm_config=vllm_config, )
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/executor/executor_base.py", line 51, in __init__
[rank0]:     self._init_executor()
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/executor/uniproc_executor.py", line 41, in _init_executor
[rank0]:     self.collective_rpc("init_device")
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/executor/uniproc_executor.py", line 51, in collective_rpc
[rank0]:     answer = run_method(self.driver_worker, method, args, kwargs)
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/utils.py", line 2220, in run_method
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/spec_decode/spec_decode_worker.py", line 328, in init_device
[rank0]:     self.proposer_worker.load_model()
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/worker/worker.py", line 183, in load_model
[rank0]:     self.model_runner.load_model()
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/worker/model_runner.py", line 1112, in load_model
[rank0]:     self.model = get_model(vllm_config=self.vllm_config)
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/model_executor/model_loader/__init__.py", line 14, in get_model
[rank0]:     return loader.load_model(vllm_config=vllm_config)
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/model_executor/model_loader/loader.py", line 386, in load_model
[rank0]:     loaded_weights = model.load_weights(
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/model_executor/models/llama.py", line 570, in load_weights
[rank0]:     return loader.load_weights(
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/model_executor/models/utils.py", line 235, in load_weights
[rank0]:     autoloaded_weights = set(self._load_module("", self.module, weights))
[rank0]:   File "/data/lily/vllm-eagle-simplify/vllm/model_executor/models/utils.py", line 224, in _load_module
[rank0]:     raise ValueError(msg)
[rank0]: ValueError: There is no module or parameter named 'embed_tokens' in LlamaForCausalLM

[ShangmingCai]

Sorry for triggering DCO misbehavior while trying to rebase the code to address an unknown doc build failure. It was a mistake. My bad.

[ShangmingCai]

I'm trying this PR, but it seems there are some minor errors:

@LiuXiaoxuanPKU Hello, thanks for the feedback, but I couldn't reproduce this embed_tokens error in my local environment. Therefore, I have added three e2e eagle tests and temporarily disabled the bias ValueError when loading eagle with a warning (I will revert these changes once these e2e tests validation is successful if this modification is not appropriate). All the tests have passed in my local environment, but it seems there exists some cleanup issues in the buildkit CI. Could you please manually retrigger these tests?

Additionally, some other tests weren't completed successfully due to issues with failed HuggingFace connection and cleanup issues too.

[ShangmingCai]

I'm trying this PR, but it seems there are some minor errors:

@LiuXiaoxuanPKU Hello, thanks for the feedback, but I couldn't reproduce this embed_tokens error in my local environment.

@LiuXiaoxuanPKU I think I just successfully reproduced the error regarding embed_tokens using the code from the main branch. Could you please check if the error is caused by the branch not being switched to the patched code?

If that’s the case, we can consider whether to revert the changes regarding the eagle bias and add the e2e tests when this issue is resolved.

[LiuXiaoxuanPKU]

Sorry for delaying the PR here. Yeah I double checked, it's the problem of the main branch. Your branch works as expected. I just re-triggered the CI, let's merge this PR after CI passes. Thanks!