Lightning-AI / Lightning-AI/litgpt
attention mask is incorrect when generate with softcapping
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
When using `litgpt generate` on models with softcapping, `build_mask_cache` creates mask as `torch.bool`
https://github.com/Lightning-AI/litgpt/blob/ef9647cfa7cd73e03b0e29126bfe8b42cae509eb/litgpt/model.py#L465
and then it's added to scores.
https://github.com/Lightning-AI/litgpt/blob/ef9647cfa7cd73e03b0e29126bfe8b42cae509eb/litgpt/model.py#L309
Therefore, attention mask is not accounted for and it does +0 or +1 on scores.
### What operating system are you using?
Unknown
### LitGPT Version
```
```
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/model.py at build_mask_cache around line 465 and the score calculation around line 309. Reproduce generation with a model using softcapping, then inspect how the mask is represented when it is added to scores. Done means the attention mask is correctly accounted for during softcapped generation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100