Lightning-AI / Lightning-AI/litgpt
Stop sequence matching mishandles overlapping prefixes and buffered tokens
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 1.5k
- Avg merge
- 15h 37m
- Merged PRs (30d)
- 1
Description
## Bug description
`generate_fn()` and `batched_generate_fn()` keep one prefix length for each stop sequence and reset that length to zero after any mismatch. This misses a valid match when the mismatching token is also the beginning of a new occurrence of the same sequence.
For example, with stop sequence `[1, 2]` and generated tokens `[1, 1, 2]`, the second `1` should restart the match. Instead, the current matcher resets to zero and returns all three tokens rather than stopping before the final `[1, 2]`.
There is a related buffering problem when multiple stop sequences are active. With stop sequences `[1, 3]` and `[2]`, generated tokens `[1, 2]` should return `[1]`: the first token is temporarily buffered as a possible prefix, then becomes safe when `[2]` completes. The current implementation returns immediately on the completed sequence and drops the buffered `1` as well.
## Minimal reproduction
```python
from unittest.mock import MagicMock, patch
import torch
import litgpt.generate.base as generation
from litgpt.chat.base import generate
def run(generated, stop_tokens):
model = MagicMock()
model.max_seq_length = 100
samples = iter(generated)
with patch.object(
generation,
"multinomial_num_samples_1",
lambda *args, **kwargs: torch.tensor([next(samples)]),
):
output = list(
generate(
model,
torch.tensor([5, 3]),
len(generated) + 2,
stop_tokens=stop_tokens,
)
)
return torch.cat(output).tolist() if output else []
print(run([1, 1, 2], ([1, 2],))) # actual: [1, 1, 2], expected: [1]
print(run([1, 2], ([1, 3], [2]))) # actual: [], expected: [1]
```
## Expected behavior
Stop matching should retain the longest suffix that is also a prefix of each stop sequence. When a sequence completes, tokens buffered only for other partial matches should be emitted before generation stops. The result should not depend on the ordering of `stop_tokens`.
This affects both the single-prompt and batched streaming generators because they use the same reset-on-mismatch logic.
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/base.py and trace the shared stop-matching logic used by generate_fn() and batched_generate_fn(), then use the reproduction in the issue through litgpt/chat/base.py. Done means overlapping prefixes are retained, buffered tokens from other partial matches are emitted before stopping, and results do not depend on stop_tokens ordering.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100