[Core] [Bugfix] Add Input Embeddings

[qthequartermasterman]

Note

This PR is just #11684, but rebased onto main and then with pre-commit errors fixed, since it has been some time since @Bryce1010 has updated that PR.

Adds support for passing prompt_embeds to LLM.generate as

llm.generate({"prompt_embeds": input_embeds}, sampling_params)

or

llm.generate(
    [{"prompt_embeds": input_embeds} for input_embeds in inputs_embeds], sampling_params
)

this enables use cases when only the embedding layer is finetuned, and have the same model backend support multiple custom tuned embedding layers

FIX #416
FIX #8323
FIX #14621

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

[liangwythu]

our project need this feature，Would love to see this one merged!

[DarkLight1337]

Can you also update this PR with the unit tests in #6869 to see whether this solution works correctly?

[DarkLight1337]

Oops I accidentally closed the PR, reopened it now

[yukang2017]

Hi,

I set inputs_embeds as (num_tokens, embed_dim) and get the following issues. Are there any advice? Thanks.

[lzl-mt]

Hi,

I set inputs_embeds as (num_tokens, embed_dim) and get the following issues. Are there any advice? Thanks.

@DarkLight1337 I also meet this issue. When my input is (T, V), I get the error:
*** IndexError: index 14 is out of bounds for axis 0 with size 14

When my input is (B, T, V), I get the error:
RuntimeError: query, key, and positions must have the same number of tokens.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @qthequartermasterman.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[qthequartermasterman]

@DarkLight1337

There are three tests still failing in CI.

Async Engine Inputs: test_mp_crash_detection has a subtle bug in it. It is supposed to be testing how long error recovery takes after LLMEngine failes to initialize. But since LLMEngine is being initialized in a different process, the mock which raises an exception is not applied properly. So the actual initialization process is what is being timed here. Because there are more cudagraphs to compile, this takes longer than the 60 seconds defined in the test. The "easy" answer is to bump up the threshold in the test. Another "easy" option is to enforce eager mode to avoid compilation. But the "correct" answer is to refactor the test to cause a failure within the LLMEngine so it actually tests what it claims to be testing. Using mocks directly would be difficult because of the fact that LLMEngine is being initialized in another process. I'm not 100% certain how to do it cleanly, to be honest, without a refactor of the source code for the MQLLMEngine and/or the build_async_engine_client_from_engine_args to be able to pass in a function that raises an error (maybe a fake executor_class?)

How would you like to proceed with this?

Both the speculative decoding and v1 tests are failing on main locally for me. 🤷 I see the same v1-test is failing on main an hour ago: https://buildkite.com/vllm/ci/builds/19179#01968e82-8601-44f0-8825-816c8f31a236. The speculative decoding tests seem not to have run in that build, so I don't know for certain if they're failing on main in CI. For the life of me I cannot figure out how those speculative decoding tests would have been affected by any of the changes in this PR.

[DarkLight1337]

I can help force merge if the test failures are unrelated

[KyleMylonakisProtopia]

Super happy to see this in. Thank you everyone for all the hard work!

[qthequartermasterman]

Thanks @DarkLight1337 for the patient review and iteration process!

[njhill]

I'm not sure we want to double the graph compilation time during startup in all cases, even when inputs_embeds aren't being used. This is the reason for at least one of the test failures and likely others too.

[DarkLight1337]

My bad for merging this. I have opened #17607 to fix this problem by disabling input embeddings by default.

[gshtras]

Should be elif? The way it is now it'll raise on mllama with empty encoder input