[Bug Report] Mixtral generates nonsense

[joelburget]

Describe the bug

I followed the instructions in docs/source/content/special_cases.md as well as I could tell (ran the model in both full precision and with HookedTransformer.from_pretrained_no_processing), yet my model generations were nonsensical.

Code example

from transformer_lens import HookedTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = HookedTransformer.from_pretrained_no_processing(
    "mistralai/Mixtral-8x7B-v0.1",
    n_devices=4,
)

# test model actually works
for i in range(5):
    display(
        model.generate(
            "Once upon a time",
            verbose=False,
            max_new_tokens=50,
        )
    )

System Info
Describe the characteristic of your environment:

• Describe how transformer_lens was installed: pip install sae-lens
• What OS are you using? Ubuntu 22.04 (runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel-ubuntu22.04)
• Python version: 3.10

Additional context
Running on a 4x A100 SXM system on Runpod.

Checklist

• I have checked that there is no similar issue in the repo (required)

[jbloomAus]

Generate is not a good way to check a model is running properly. Can you run the following and share the results?

a_large_chunk_of_text = "Generate is not a good way to check a model is running properly. Can you run the following and share the results?"
loss = model(a_large_chunk_of_text, return_type="loss")
print(loss.item())

[joelburget]

Looks like the loss is ~6.82

[jbloomAus]

This could be related to this issue: #591 . If I were you, I'd pip install opt_einsum, then see if you get the same issue. opt_einsum should be automatically used by pytorch if it's installed and the other issue was solved by using it. Let me know if you still get garbage afterward.

[joelburget]

Still seems quite bad even when verifying that I have opt_einsum installed. For good measure I checked a few more times and got 5.5, 6.7, and 8.7. So note quite as bad 12.5 but still problematic.

[jbloomAus]

Are you adding or missing BOS tokens? When we accepted a PR for this, I assume it was working...

…

On Tue, May 21, 2024, 8:58 PM Joel Burget ***@***.***> wrote: Still seems quite bad even when verifying that I have opt_einsum installed. For good measure I checked a few more times and got 5.5, 6.7, and 8.7. So note quite as bad 12.5 but still problematic. Screenshot.2024-05-21.at.12.55.11.PM.png (view on web) <https://github.com/TransformerLensOrg/TransformerLens/assets/310981/c22250e0-eb7e-491c-97cc-5230db773bf4> — Reply to this email directly, view it on GitHub <#570 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AQPMYZ5T4T4JI74FA7UV7ZDZDORODAVCNFSM6AAAAABHGYQMD6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMRTGM2DINJXGM> . You are receiving this because you commented.Message ID: ***@***.***>

[joelburget]

I'm not using any special tokens at all. The a_large_chunk_of_text code block is what you pasted verbatim.

[jbloomAus]

Ahh sorry. Sometimes language models expect a special prepended token. You can often find examples of this in model cards on huggingface. I'm not sure but this could be the issue here.

…

On Tue, May 21, 2024, 9:42 PM Joel Burget ***@***.***> wrote: I'm not using any special tokens at all. The a_large_chunk_of_text code block is what you pasted verbatim. — Reply to this email directly, view it on GitHub <#570 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AQPMYZ36YLQK6DK7WXHTMZ3ZDOWT5AVCNFSM6AAAAABHGYQMD6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMRTGQYDQNJSGI> . You are receiving this because you commented.Message ID: ***@***.***>

[joelburget]

Unfortunately I don't think that's the issue. Here are the instructions from the Mixtral HF repo:

Running those in a Colab we can see that the only special token they prepend is BOS (1):

TransformerLens of course does this by default but I double-checked: adding prepend_bos doesn't change the behavior and (2) the tokenizer output is the same as with the official instructions.

[joelburget]

I also checked following the instructions on the HF repo without using TransformerLens. The generation looks reasonable. I'm not sure exactly how to interpret the loss or if it matters.

[bryce13950]

I wonder if this is a memory issue. The only difference between the config for TransformerLens and HuggingFace is n_ctx being capped at 2048 for memory concerns, where it is 32768 in HuggingFace. The generation isn't entirely nonsense, it's just mostly generating French, and every once in a while Spanish or English. What it is actually generating does for the most part make sense as something that would follow the idiom being passed through in the various languages. I could imagine that a lack of memory could cause something like languages to get mixed up, and I think it may be a good place to start investigating.

[neelnanda-io]

I don't see how memory could explain that? Low memory shouldn't silently lead to issues, it should lead to out of memory errors. Mathematically, TL and HF *should *be doing the same operations, and this discrepancy implies they are not, and it seems important to understand why. The obvious hypothesis is that it's because of the mixture of expert routing. This amplifies small errors because it's discrete, maybe. So it would be good to compare the routing on HuggingFace and the routing in TL and see how often they disagree in these generations.

…

On Tue, 4 Jun 2024 at 02:08, Bryce Meyer ***@***.***> wrote: I wonder if this is a memory issue. The only difference between the config for TransformerLens and HuggingFace is n_ctx being capped at 2048 for memory concerns, where it is 32768 in HuggingFace. The generation isn't entirely nonsense, it's just mostly generating French, and every once in a while Spanish or English. What it is actually generating does for the most part make sense as something that would follow the idiom being passed through in the various languages. I could imagine that a lack of memory could cause something like languages to get mixed up, and I think it may be a good place to start investigating. — Reply to this email directly, view it on GitHub <#570 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ASRPNKLPYFXFDHBHXPSPNG3ZFUHQVAVCNFSM6AAAAABHGYQMD6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNBWGM3TQNZTHE> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[bryce13950]

Well I think I have a pretty decent idea on how to start with this. I will work with @joelburget to try and isolate the error, and hopefully we can find it relatively quickly.

[joelburget]

I ran some experiments in this notebook: https://gist.github.com/joelburget/bae5ea4d997e804b2a65d02d5b61f5bc

1. all weights seem to match
2. when running a random tensor through block 1:
   a. 89 / 4096 mlp outputs match. Those that don't seem to be small differences like 0.38261502981185913 vs 0.38261523842811584
   b. 254 / 4096 attention outputs match.
   c. 3218 / 4096 block outputs match. Again we see small differences like -0.20334099233150482 vs -0.20334100723266602.

I'm surprised that there are more matching block outputs than mlp or attention, but whatever, I guess this is possible.

Ideas for where to go from here?