Lightning-AI / Lightning-AI/litgpt

Speculative decoding drops the first generated token

Open Beginner friendly
#2,307 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
13.7k
Forks
1.5k
Avg merge
15h 37m
Merged PRs (30d)
1

Description

# Speculative decoding drops the first generated token

## Bug description

`generate()` samples the target model's first new token during prefill:

```python
token = next_token(
target_model,
input_pos,
prompt.view(1, -1),
...
)
```

That token is then passed to `speculative_decoding()` as context, but the output accumulator is
initialized as an empty list and only records the later tokens returned by speculative decoding.
The prefill token is therefore never returned.

This changes the generated text and also makes the result one token shorter than
`max_returned_tokens`. The existing `test_generate` currently codifies that short output with
`T + max_new_tokens - 1`.

The same gap means a stop token sampled during target prefill is not checked before another round
of generation.

## Minimal CPU reproduction

This does not require a GPU or model download:

```python
from unittest.mock import patch

import torch
from torch import nn

import litgpt.generate.speculative_decoding as generation

model = nn.Module()
model.max_seq_length = 10
prompt = torch.tensor([1, 2])

with (
patch.object(
generation,
"next_token",
side_effect=[torch.tensor([9]), torch.tensor([7])],
),
patch.object(
generation,
"speculative_decoding",
side_effect=[torch.tensor([8]), torch.tensor([8])],
),
):
output, _ = generation.generate(model, model, prompt, 5, speculative_k=1)

print(output.tolist())
assert output.tolist() == [1, 2, 7, 8, 8]
```

Current output:

```text
[1, 2, 8, 8]
```

## Expected behavior

The first target token sampled after the prompt should be the first returned completion token. It
should count toward `max_returned_tokens` and should go through the same stop-token handling as
tokens produced by later speculative decoding rounds.

## Proposed fix

Initialize the generated-token accumulator with the target prefill token when it is not a stop
token, skip further generation when it is a stop token, and handle the resulting empty completion
without calling `torch.stack([])`. Add focused CPU regression coverage and update the existing
length assertion.

## Environment

- Operating system: Linux
- LitGPT: `main` at `7bf2960dfb26bae8e815c9a16a22732974824ac1`
- PyTorch: 2.13.0

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in litgpt/generate/speculative_decoding.py, especially generate() and its next_token/speculative_decoding calls. Reproduce the issue with the supplied CPU mocks and inspect test_generate, then add focused coverage for the prefill token and stop-token case. Done means the prefill token is returned and counted, stop handling is consistent, and the existing length assertion passes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
84/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.