Inconsistency in logit values between generation and direct model prediction
Author: lowlypalace
Pull request: no
Labels: name: Generation

### System Info

- torch  2.3.0+cu121
- transformers  4.41.1
- Python 3.10.12

### Who can help?

@gante
@ArthurZucker
@younesbelkada

### 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

1. **Initialization**: Use the `GPTNeoXForCausalLM` and `AutoTokenizer` from the `transformers` library, such as `EleutherAI/pythia-70m` model.

2. **Encoding and Decoding**:
   - Encode a simple input (e.g., end-of-string token).
   - Generate a sequence using the `generate` method with `output_logits=True`.

3. **Direct Prediction**:
   - Directly pass the generated text back into the model to get logits.

4. **Comparison**:
   - Compare logits from the `generate` method with those obtained from direct prediction.

5. **Precision Adjustment**:
   - Convert model to double precision and repeat the generation and prediction to compare logits to see that they still differ.

```
import torch
import os
from transformers import AutoTokenizer, GPTNeoXForCausalLM

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

torch.set_printoptions(precision=64)

# Set the random seed and deterministic algorithms for reproducibility
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
torch.use_deterministic_algorithms(True)

# Initialize tokenizer and model, set to evaluation mode, and move to CPU for reproducibility
tokenizer = AutoTokenizer.from_pretrained(f"EleutherAI/pythia-70m")
model = GPTNeoXForCausalLM.from_pretrained(f"EleutherAI/pythia-70m")

device = torch.device("cuda")

model.eval()
model.to(device)
model.double()

# Encode the input text
input_text = tokenizer.eos_token
input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)

# Generate sequences using the model
generate_outputs = model.generate(
    input_ids,
    max_length=len(input_ids[0]) + 5,
    do_sample=True,
    return_dict_in_generate=True,
    output_logits=True,
)

# Decode the generated sequence
generated_ids = generate_outputs.sequences
generated_text = tokenizer.decode(generated_ids[0])

# Re-encode and evaluate the generated text
inputs = tokenizer(generated_text, return_tensors="pt").input_ids.to(device)
outputs = model(inputs)

# Extract logits
# Note: Adjust indices based on where the new tokens start in the generated sequence
first_generated_token_logits_from_generate = generate_outputs.logits[1].squeeze()
first_generated_token_logits_from_model = outputs.logits[0, len(input_ids[0]):len(input_ids[0])+1].squeeze()

# Print out the logits for comparison
print("Logits from generate method for the first generated token:")
print(first_generated_token_logits_from_generate)
print("\nLogits from model for the first generated token:")
print(first_generated_token_logits_from_model)

```

Without calling `model.double()` (comment out the line):

```
Logits from generate method for the first generated token:
tensor([1071.1469726562500000000000000000000000000000000000000000000000000000,
         230.7740783691406250000000000000000000000000000000000000000000000000,
        1070.1348876953125000000000000000000000000000000000000000000000000000,
         ...,
         230.7809600830078125000000000000000000000000000000000000000000000000,
         230.7765350341796875000000000000000000000000000000000000000000000000,
         230.7785034179687500000000000000000000000000000000000000000000000000],
       device='cuda:0')

Logits from model for the first generated token:
tensor([1071.1472167968750000000000000000000000000000000000000000000000000000,
         230.7741088867187500000000000000000000000000000000000000000000000000,
        1070.1352539062500000000000000000000000000000000000000000000000000000,
         ...,
         230.7809143066406250000000000000000000000000000000000000000000000000,
         230.7764892578125000000000000000000000000000000000000000000000000000,
         230.7785339355468750000000000000000000000000000000000000000000000000],
       device='cuda:0', grad_fn=<SqueezeBackward0>)

```
       

With calling `model.double()`:

```
Logits from generate method for the first generated token:
tensor([1071.1467351228166080545634031295776367187500000000000000000000000000,
         230.7740661093434084705222630873322486877441406250000000000000000000,
        1070.1347990279657551582204177975654602050781250000000000000000000000,
         ...,
         230.7809335116805300458509009331464767456054687500000000000000000000,
         230.7765007713145166690082987770438194274902343750000000000000000000,
         230.7784833010678937625925755128264427185058593750000000000000000000],
       device='cuda:0', dtype=torch.float64)

Logits from model for the first generated token:
tensor([1071.1467351228168354282388463616371154785156250000000000000000000000,
         230.7740661093434368922316934913396835327148437500000000000000000000,
        1070.1347990279655277845449745655059814453125000000000000000000000000,
         ...,
         230.7809335116805016241414705291390419006347656250000000000000000000,
         230.7765007713144882472988683730363845825195312500000000000000000000,
         230.7784833010680358711397275328636169433593750000000000000000000000],
       device='cuda:0', dtype=torch.float64, grad_fn=<SqueezeBackward0>)
```

As you can see, even after calling `model.double()`, the logits differ slightly.

### Expected behavior

Logits from the `generate` method and direct model predictions should be identical.  I have not tried this with other models, but I belive this is not a model-specific issue.


Comments:
Hi @lowlypalace 
As long as the test pass with `torch.allclose()` there is nothing to worry about
```python
print(torch.allclose(first_generated_token_logits_from_model, first_generated_token_logits_from_generate, atol=1e-6, rtol=1e-6))
```
There might be some slicing operations within generate that are responsible of this slight difference - see: https://github.com/pytorch/pytorch/issues/76232 for more details

Hi @younesbelkada thank you for your reply.

My use case actually requires me to use `generate()` to create sequences with the specific top-k values. I then call the model directly and obtain the logits. Afterwards, I manually apply top-k filtering on these logits with `torch.topk`. The problem arises with the discrepancies between the logits (probabilities) generated from the model directly. As the logits don't match, I'm unable to consistently obtaining the same top-k tokens that were actually selected in generate method. This leads me to selecting the wrong top-k tokens (e.g. the ones that were actually sampled by generate but have zero probability in case of obtaining logits directly from model) and leads to miscalculations when it comes to the probability of a full sequence.

I've found that converting the model to double precision using `model.double()` partially addresses this inconsistency, but the issue still occasionally persists.

Is there anything I can do to make sure the logits remain identical when I call the model directly so that I can ensure the consistency between the selected top-k tokens?

I see, thanks for giving me more context to the problem ! how do you verify that logits match? Can't you use `torch.allclose(tensor1, tensor2)` with reasonable `atol` and `rtol` - e.g. `1e-3` ? would that still fail?

Hey! I think you can get identical logits without double precision if you disable caching by "use_cache=False" 🤔  

> Hey! I think you can get identical logits without double precision if you disable caching by "use_cache=False" 🤔

Hey, thank you so much for the hint! I've tried the code snippet above with single precision (without `model.double()` `use_cache=False`):

```
Logits from generate method for the first generated token:
tensor([1069.5090332031250000000000000000000000000000000000000000000000000000,
         228.6221313476562500000000000000000000000000000000000000000000000000,
        1071.3411865234375000000000000000000000000000000000000000000000000000,
         ...,
         228.6294097900390625000000000000000000000000000000000000000000000000,
         228.6256256103515625000000000000000000000000000000000000000000000000,
         228.6267395019531250000000000000000000000000000000000000000000000000],
       device='cuda:0')

Logits from model for the first generated token:
tensor([1069.5091552734375000000000000000000000000000000000000000000000000000,
         228.6221923828125000000000000000000000000000000000000000000000000000,
        1071.3410644531250000000000000000000000000000000000000000000000000000,
         ...,
         228.6295166015625000000000000000000000000000000000000000000000000000,
         228.6257019042968750000000000000000000000000000000000000000000000000,
         228.6268005371093750000000000000000000000000000000000000000000000000],
       device='cuda:0', grad_fn=<SqueezeBackward0>)
 ```
 
 
 However, double precision (with `model.double` and `use_cache=False`) seems to be stable now:
 
 ```
 tensor([1071.1467351228168354282388463616371154785156250000000000000000000000,
         230.7740661093434368922316934913396835327148437500000000000000000000,
        1070.1347990279650730371940881013870239257812500000000000000000000000,
         ...,
         230.7809335116805016241414705291390419006347656250000000000000000000,
         230.7765007713144882472988683730363845825195312500000000000000000000,
         230.7784833010680358711397275328636169433593750000000000000000000000],
       device='cuda:0', dtype=torch.float64)

Logits from model for the first generated token:
tensor([1071.1467351228168354282388463616371154785156250000000000000000000000,
         230.7740661093434368922316934913396835327148437500000000000000000000,
        1070.1347990279655277845449745655059814453125000000000000000000000000,
         ...,
         230.7809335116805016241414705291390419006347656250000000000000000000,
         230.7765007713144882472988683730363845825195312500000000000000000000,
         230.7784833010680358711397275328636169433593750000000000000000000000],
       device='cuda:0', dtype=torch.float64, grad_fn=<SqueezeBackward0>)
 ```
 
 Do you know if there is anything else I could do to make the logits match for single precision?
 

@lowlypalace I don't think there's anything else to make them identical. As @younesbelkada said, there will always be some small numerical precision errors. Disabling cache and recalculating keys/values every time was one straightforward way to reduce errors


Regarding your use case, you said

> (e.g. the ones that were actually sampled by generate but have zero probability in case of obtaining logits directly from model) 

I am wondering why they have zero probability if the logits are close enough when checked with `torch.allclose()`? I think applying top-k should give roughly same top-k tokens

Maybe this might help you. The `scores` are logits after applying top-k filtering, the ones used by generate to sample the next token :)

```python
# Generate sequences using the model
generate_outputs = model.generate(
    input_ids,
    max_length=len(input_ids[0]) + 5,
    do_sample=True,
    return_dict_in_generate=True,
    output_logits=True,
   output_scores=True,
)
scores = generate_outputs.scores
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.