Support mistral 7 b

[Felhof]

Description

This PR closes #387 by adding support for Mistral and implementing Grouped Query Attention.

Mistral 7B

The weight's from Huggingface's Mistral-7B-v0.1 can now be loaded into a HookedTransformer using

# these parameters are necessary as Mistral uses RMS norm and not layernorm.
tl_mistral = HookedTransformer.from_pretrained(
    "mistral-7B",
    fold_ln=False,
    center_writing_weights=False,
    center_unembed=False
)

Note that Mistral is only supported in transformers >= 4.34 and hence Python >= 3.8 is required to use it.
The demo notebook Mistral.ipynb features a comparison of Huggingface's Mistral with the HookedTransformer implementation. I tested it on the same prompts that were used in the Llama demo. The differences in the resulting logits were around 0.01 which is slightly higher than what Llama2 gets. This may be due to a known issue with rotary embeddings which also affects LLama2 and Pythia.

Grouped Query Attention

Mistral utilizes Grouped Query Attention (GPA) which is not used by any other model supported by TransformerLens and had to be implemented. I added a new class GroupedQueryAttention and an abstract base class AbstractAttention which now features the common functionality of both attention classes. The difference between Attention and GroupedQueryAttention is how they handle the key and value projections, as in GPA groups of queries share the same keys and values (see image below). This mostly affects the internal workings of the class. To avoid confusion, to not break existing code interacting with attention, and to make the design of future code easier, public attributes such as W_K and W_V have the same shape for both classes. This is because in the case of GPA, the underlying parameters are hidden behind a property that expands them using torch.repeat_interleave. A GPA block should behave the same as a regular Attention block whose weights are the result of applying torch.repeat_interleave to the GPA block's weight. There is a unit test that confirms this.

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[neelnanda-io]

Thanks a lot for doing this! I left some small comments, but this isn't a full review (I haven't looked at exactly how you implemented grouped query attention), so I'll leave that for someone else to finish off.

[alan-cooney]

Just to flag I'll review this once the previous PR goes through on attention, as we'll need to adjust for a few conflicts in it

[alan-cooney]

Hey @Felhof nice work on this, and the abstract attention class approach is great (we're planning on doing this for a bunch of the components).

Would you be able to fix the merge conflicts first and then I can do a full review? There's a few things there as the last PR to go through had some attention component changes.

[neelnanda-io]

Hey! I'm curious what the status of this PR is? A few of my MATS scholars want to use Mistral. Can they just check out this PR?

[ojh31]

Hey! I'm curious what the status of this PR is? A few of my MATS scholars want to use Mistral. Can they just check out this PR?

That worked for me, modulo the comment I made above about using the right version of transformers (and the tokenizer being disgusting)

[abdurraheemali]

@alan-cooney @neelnanda-io @Felhof I want to run an activation patching experiment on llama-70b, so I'm going to check out the support-mistral-7b branch and report (in an hour-ish?) whether the grouped query attention implementation works for me or not

Edit: I couldn't get it to work with some other dependencies I had, but I'll try again later this week.

[neelnanda-io]

@alan-cooney @Felhof What's the status of this PR? I'm not sure which of you it's blocked on. Either way, I hear that people have been able to checkout this branch and get Mistral working, so thanks a lot for the work up to that point!

[Felhof]

The branch has been working for a while but needs approval to be merged into main :) I made sure it's up-to-date with main again.

[wesg52]

@alan-cooney Ping on review :)

[alan-cooney]

Sorry folks!! I'm on it now

[alan-cooney]

Nice! Thanks for adding this and the approach is great.

A general point (don't feel the need to change, but to be aware) is that we're trying to type all methods and also add docstrings to them (Google style). If you have a chance it would be good to do this as well to e.g. explain the interleave approach that is used.

[Felhof]

Thanks for the review Alan! I have added better typing and documentation and removed the demo

[alan-cooney]

Thanks! And thanks for adding this!

[neelnanda-io]

Woot! Really glad this got merged in! Thanks for adding it @Felhof and sorry for the long delay