Lookahead in forward

[Dionysour]

System Info

transformers '4.47.0', bitsanbytes '0.45.0', torch '2.5.1+cu124', NVIDIA RTX A4000

Reproduction

During assisted greedy generation with the verifier loaded in 8-bit, I have noticed that it diverges from vanilla greedy generation. I have investigated this problem and discovered lookahead during forward on a particular sequence. Here's a snippet of the notebook without unnecessary code. You can run it and see the problem, but you need a Llama2 chat token because it's the model I was using.

https://gist.github.com/Dionysour/24b352bb685f7d4a8ffd18896455700d

Expected behavior

change of [n] token doesn't affect logits of [:n] tokens

[matthewdouglas]

Thanks for the report. I was unable to reproduce this with Llama 3.1 8B in an environment with torch 2.5.1+cu124, transformers 4.47.0, bitsandbytes 0.45.0, on an RTX 4090.

I'm downloading Llama2 now and will give that a try.

[Dionysour]

I think this behaviour with llama 3.1 is expected. The bug is very model and sequence dependent. I saw different predictions for different models too. So to reproduce it model must be "meta-llama/Llama-2-7b-chat-hf"

[poedator]

I also have a reproduction of the issue found by @Dionysour in this notebook.
https://nbviewer.org/gist/poedator/52440da0cdab808a4c0ff929f009ae31

tested on A100 with
bitsandbytes 0.44.1
torch 2.4.1
transformers 4.45.2
python 3.9

Apparently something is wrong with how bnb-8bit model deals with attention.