Dose `num_logits_to_keep` in `model.generate()` really work?
Author: bwnjnOEI
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.47.1
- Platform: Linux-6.9.3-76060903-generic-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.27.0
- Safetensors version: 0.5.0
- Accelerate version: 1.2.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+cu124 (True)
- 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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA GeForce RTX 4090

### Who can help?

@zucchini-nlp @gante @ArthurZucker 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I want to customize the `num_logits_to_keep` in `LlamaForCausalLM`'s `forward` function by setting it through generate, but I failed or don't know how to check the logits of tokens produced by `num_logits_to_keep`.

a code snippet
```python
outputs = model.generate( 
    ids['input_ids'],
    max_new_tokens=50,
    pad_token_id=tokenizer.eos_token_id,
    output_scores=True,
    return_dict_in_generate=True,
    num_logits_to_keep=3,
)

or

model.generation_config.num_logits_to_keep = 3
```

### Expected behavior

outputs.logits != None ?

Comments:
Hey @bwnjnOEI !

The `num_logits_to_keep` is working. You have to access `out.scores` since you are requesting to `output_scores` in the code snippet. For logits try to `output_logits=True`

Closing as complete as there is no bug :)

> Hey @bwnjnOEI !
> 
> The `num_logits_to_keep` is working. You have to access `out.scores` since you are requesting to `output_scores` in the code snippet. For logits try to `output_logits=True`
> 
> Closing as complete as there is no bug :)

Thank you @zucchini-nlp for your answer, but I still have doubts. When the `LlamaForCausalLM.forward`, I want to keep the logits of the last `num_logits_to_keep` tokens at each generating step. According to your hint, I have set `num_logits_to_keep` and `output_logits=True` in the `.generate`. However, the shape of logits printed directly from `LlamaForCausalLM.forward` at each step is inconsistent with the shape of logits returned by `.generate` in the dictionary; it seems to only have the logits of the last token. Have I missed something?

### result 
```bash
input ids len:  5
max_new_tokens:  10
num_logits_to_keep:  3

The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])
The shape of logits in generating step in .forward():  torch.Size([1, 3, 128256])

The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
The shape of logits at each index of .generate().logits torch.Size([1, 128256])
```
I apologize for asking such a naive question, but there are really very few related documents or examples, and the `GenerationMixin` codebase is quite large. Thank you again here.

@bwnjnOEI right, since `generate()` usually return not the whole logits as per model output, but only the current logit for the token being generated. So we end up with one logit for each new token that the model generated. The main idea of `num_logits_to_keep` here is to be more memory efficient and get rid of logits that are not contributing to next token prediction, thus by default it has a value of `1` when not specified :)

If you want full logits, I'd recommend to use a `forward()` pass

hello, I have a same transformers version: 4.47.1
but when I use `num_logits_to_keep` in `model.generate` function,
```
 output = model.generate(input_ids.repeat(10, 1), num_logits_to_keep=1, do_sample=True, max_new_tokens=100, use_cache=True, )
```
 there is an error:
```
Traceback (most recent call last):
  File "/home/wufang/PSCache/eval/efficiency/efficiency.py", line 237, in <module>
    output = model.generate(input_ids.repeat(10, 1), num_logits_to_keep=1, do_sample=True, max_new_tokens=100, use_cache=True, )
  File "/mnt/wufang/anaconda3/envs/fix-prcache/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/wufang/anaconda3/envs/fix-prcache/lib/python3.10/site-packages/transformers/generation/utils.py", line 2009, in generate
    self._validate_model_kwargs(model_kwargs.copy())
  File "/mnt/wufang/anaconda3/envs/fix-prcache/lib/python3.10/site-packages/transformers/generation/utils.py", line 1388, in _validate_model_kwargs
    raise ValueError(
ValueError: The following `model_kwargs` are not used by the model: ['num_logits_to_keep'] (note: typos in the generate arguments will also show up in this list)
```
I really don't know why

@wfloveiu which model is that? Unfortunately `num_logits_to_keep` is not present on older models like T5, it was added only for LLMs which use much more memory to generate

If you want this argument to be used with other models, feel free to submit a PR 🤗