fix: add inputs_embeds support to prepare_inputs_for_generation

[majiayu000]

Summary

Add inputs_embeds parameter support to _fast_prepare_inputs_for_generation to enable generation with embeddings directly.

Problem

When calling model.generate(inputs_embeds=...) on Unsloth-patched models, users get:

ValueError: You passed `inputs_embeds` to `.generate()`, but the model class LlamaForCausalLM doesn't have its forwarding implemented.

This happens because Unsloth's _fast_prepare_inputs_for_generation function didn't include inputs_embeds in its signature, causing transformers to think the model doesn't support embedding-based generation.

Solution

1. Added inputs_embeds = None parameter to the function signature
2. Modified the return dict to include either inputs_embeds or input_ids based on what was provided

Test Plan

• Python syntax check passes

Fixes #3779

[gemini-code-assist]

Summary of Changes

Hello @majiayu000, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces crucial support for inputs_embeds within the _fast_prepare_inputs_for_generation function, addressing a ValueError that previously prevented Unsloth-patched models from performing generation directly from embeddings. The change enhances the model's flexibility by allowing users to provide pre-computed embeddings as input, streamlining the generation process for advanced use cases.

Highlights

• Inputs Embeds Support: Added inputs_embeds parameter to the _fast_prepare_inputs_for_generation function, enabling direct embedding-based generation for Unsloth-patched models.
• Bug Fix: Resolved a ValueError that occurred when attempting to call model.generate(inputs_embeds=...) due to the absence of inputs_embeds in the function signature.
• Generation Logic Update: Modified the return dictionary within _fast_prepare_inputs_for_generation to conditionally include either inputs_embeds or input_ids based on what was provided, ensuring correct input handling for generation.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request aims to add support for inputs_embeds in _fast_prepare_inputs_for_generation to enable direct generation from embeddings. The change to the function signature is correct. However, the logic to decide whether to use inputs_embeds or input_ids appears to be flawed. I've provided a critical review comment with a suggested code change to ensure inputs_embeds is correctly prioritized when provided.

[gemini-code-assist]

The condition if inputs_embeds is not None and input_ids is None: is likely incorrect. When model.generate(inputs_embeds=...) is called, transformers often creates a dummy input_ids tensor, so input_ids will not be None. This causes the else branch to be taken, ignoring the provided inputs_embeds.

To correctly prioritize inputs_embeds when it's available, the condition should be simplified.

Suggested change

	if inputs_embeds is not None and input_ids is None:
	result["inputs_embeds"] = inputs_embeds
	else:
	result["input_ids"] = input_ids
	if inputs_embeds is not None:
	result["inputs_embeds"] = inputs_embeds
	else:
	result["input_ids"] = input_ids

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Use inputs_embeds on first step even if input_ids is set

If generate() supplies a placeholder input_ids tensor alongside inputs_embeds for the first step (a common pattern; see how other models in this repo gate on “empty past_kv” rather than input_ids is None), this condition will ignore inputs_embeds and feed the dummy input_ids instead. That means callers using inputs_embeds=... can silently get outputs that ignore their embeddings whenever input_ids is non-None (even empty), defeating the new feature. Consider selecting inputs_embeds whenever past key values are empty, not only when input_ids is None.

Useful? React with 👍 / 👎.

[majiayu000]

Thanks @gemini-code-assist and @chatgpt-codex-connector for catching this bug!

I've fixed the condition to prioritize inputs_embeds when available:

# Before (buggy)
if inputs_embeds is not None and input_ids is None:

# After (fixed)
if inputs_embeds is not None:

This ensures that when model.generate(inputs_embeds=...) is called, the provided embeddings are used even if transformers creates a dummy input_ids tensor.

Related issue: #3798

[chatgpt-codex-connector]

To use Codex here, create a Codex account and connect to github.

[gemini-code-assist]

You're welcome, @majiayu000! I'm glad to see the fix for prioritizing inputs_embeds has been implemented. This indeed correctly addresses the scenario where transformers might provide a dummy input_ids even when inputs_embeds are intended for use. Great work!