Lightning-AI / Lightning-AI/litgpt
Speculative decoding drops the first generated token
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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