`generate()` produces incoherent output when `inputs_embeds` has length 1

[tyarkoni]

System Info

• transformers version: 4.57.1
• Platform: Linux-6.8.0-51-generic-x86_64-with-glibc2.35
• Python version: 3.10.18
• Huggingface_hub version: 0.35.3
• Safetensors version: 0.6.2
• Accelerate version: 1.10.1
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA H100 PCIe

Who can help?

@xadupre (original author of code in question), @zucchini-nlp

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Steps to Reproduce

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load any causal LM
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model.eval()

# Create a single embedding (simulating a prefix like a style token)
single_embedding = torch.randn(1, 1, 768)  # [batch=1, seq_len=1, hidden_dim=768]

# Generate with length-1 inputs_embeds
with torch.no_grad():
    outputs = model.generate(
        inputs_embeds=single_embedding,
        max_length=20,
        do_sample=True,
        temperature=1.0,
        pad_token_id=tokenizer.eos_token_id,
    )

# Decode and observe gibberish
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("Generated text:", generated_text)
# Output will be incoherent repetitive tokens like "the the the I if the..."

Comparison with working case (length ≥ 2):

# Add a second embedding (e.g., BOS token embedding)
bos_embedding = model.get_input_embeddings()(torch.tensor([[tokenizer.bos_token_id]]))
two_embeddings = torch.cat([single_embedding, bos_embedding], dim=1)  # [1, 2, 768]

# Generate with length-2 inputs_embeds
with torch.no_grad():
    outputs = model.generate(
        inputs_embeds=two_embeddings,
        max_length=20,
        do_sample=True,
        temperature=1.0,
        pad_token_id=tokenizer.eos_token_id,
    )

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("Generated text:", generated_text)
# Output will be coherent

Expected behavior

Expected Behavior

generate() should produce coherent text regardless of whether inputs_embeds has length 1 or length > 1, as long as the embeddings are valid.

Actual Behavior

With length-1 inputs_embeds, generate() produces incoherent, repetitive gibberish that appears to be high-frequency tokens without proper conditioning on previous context.

Suggested Fix

The _cache_dependant_input_preparation method needs to properly handle the transition from inputs_embeds mode to input_ids mode after the first generation step. Specifically:

1. After the first token is generated from inputs_embeds, set inputs_embeds = None for subsequent iterations
2. Or, maintain proper bookkeeping so that the embeddings prefix is correctly tracked throughout the autoregressive loop
3. Or, ensure Exception 4 logic properly handles the case where we've transitioned from embeddings to token IDs

[tyarkoni]

Further information (note: generated by Claude Sonnet 4.5, but I manually verified this):

Bug Location

The issue is in the _cache_dependant_input_preparation method in transformers/generation/utils.py, specifically around lines 472-481.

Root Cause

The method has special handling for inputs_embeds when input_ids.shape[1] == 0 (Exception 4):

if inputs_embeds is not None and input_ids.shape[1] == 0:  # Exception 4
    inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]
elif (
    inputs_embeds is not None  # Exception 1
    or (cache_position[-1] >= input_ids.shape[1])  # Exception 3
):
    input_ids = input_ids[:, -cache_position.shape[0] :]

The problem:

1. On the first iteration with inputs_embeds of length 1:

   • input_ids.shape[1] == 0 (empty tensor)
   • The code correctly takes Exception 4 path
   • First token is generated successfully

2. On subsequent iterations:

   • input_ids now contains the generated token, so input_ids.shape[1] > 0
   • Exception 4 condition is FALSE
   • The code falls through to Exception 1: inputs_embeds is not None
   • But inputs_embeds is still the original embedding and doesn't correspond to the newly generated input_ids
   • The autoregressive loop loses track of the generated context

When inputs_embeds has length ≥ 2, the logic works because there's proper coordination between the embedding prefix and subsequent token generation.

[i3hz]

def _cache_dependant_input_preparation(
        self,
        input_ids: torch.LongTensor,
        inputs_embeds: Optional[torch.FloatTensor],
        cache_position: Optional[torch.LongTensor],
    ) -> tuple[torch.FloatTensor, torch.LongTensor]:
        """
        Generic cache-dependent input preparation
        The code is put in a separate function to allow granular unit testing
        as it needs a different implementation to be exportable.

        If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        - Exception 1: when passing input_embeds, input_ids may be missing entries
        - Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        - Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
        - Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and
          generate the first token for each sequence. Later use the generated Input ids for continuation.

        The current implementation does not rely on ``self`` and could be
        a class method. It is left as a standard method to be easily rewritten.
        """
        if inputs_embeds is not None and input_ids.shape[1] == 0:  # Exception 4
            # Initial embeddings case first forward pass with only embeddings
            inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]
        elif inputs_embeds is not None and cache_position[0] > 0:
            # After first iteration, stop using embeddings and switch to input_ids
            inputs_embeds = None
            input_ids = input_ids[:, -cache_position.shape[0] :]
        elif (
            inputs_embeds is not None  # Exception 1
            or (cache_position[-1] >= input_ids.shape[1])  # Exception 3
        ):
            # Original logic for other cases
            input_ids = input_ids[:, -cache_position.shape[0] :]
        
        return inputs_embeds, input_ids

Is this the correct fix?
When I passed hello to my testing script as input this is the result
I (HelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHello

After modifying the function to this I got
that was an amazing day," he muttered, "We made good calls. Got to do another

If this is the correct fix I would love submit a PR for it. @zucchini-nlp @xadupre

[i3hz]

@vasqu can you please look into this , idk who else to ping . Sorry for the ping

[vasqu]

a new generation edge case 😓 cc @gante

[i3hz]

Thanks a lot :)

[zucchini-nlp]

The root cause is in the way we try to infer prefill vs decoding stages in this line

transformers/src/transformers/generation/utils.py

Line 623 in 5689dd6

if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:

The logic for inferring prefill vs decoding has always been imperfect and we need to split out prefill to a separate logic. The PR was just merged for it in #40652 so the fix will be quite straightforward. I will make a PR shortly

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.