Pick some example evals from here (eg PIQA, TriviaQA, LAMBADA) and write code to run HookedTransformer on them: https://huggingface.co/docs/evaluate/index

A demo notebook doing this for specific benchmarks would be a good MVP, bonus is doing a generic function for any eval (or eg for multiple choice evals, vs other types).

There's an MVP demo of using checkpoints in the main demo, but it'd be nice to have something more substantial.

A good option is replicating some of the graphs in [the induction heads paper] for my interpretability-friendly models (or Pythia or Stanford CRFM). Prefix matching score and QK eigenvalue trace are easiest and should be fast, . Loss curves, in-context learning loss, loss by position and PCA of per-token loss should all be doable with approx the same code, the hard part is just going to be running a range of model checkpoints

You may want to add a PR to disable HuggingFace caching, and also to only load analyse eg every 10th checkpoint, to avoid blowing out your computer memory - by default, HuggingFace caches model weights to the hard drive, and this is pretty expensive if eg using 600 checkpoints of GPT-2 Medium!

Add a helper function (eg to ActivationCache) to calculate direct logit attribution easily. Crib from the code here: https://colab.research.google.com/github/neelnanda-io/Easy-Transformer/blob/main/Exploratory_Analysis_Demo.ipynb#scrollTo=Direct_Logit_Attribution

Hi,

When trying to use utils.test_prompt with an OPT model I am running into issues surrounding the prepending of BOS tokens. In particular, I think the call to PreTrainedTokenizerBase sets add_special_tokens to true, regardless of the prepend_base flag in model.to_str_tokens, which leads to at least one </s> being preprended (and if prepend_base=True, then two are prepended).

For example:

prompt: "Two young, White males are outside near many bushes"
# Split into prompt = "Two young .... many", answer = "bushes"
Tokenized prompt: ['</s>', '</s>', 'Two', ' young', ',', ' White', ' males', ' are', ' outside', ' near', ' many']
Tokenized answer: ['</s>', ' bushes', '.', '\n']

I am not sure if this is the desired behaviour for OPT?

Even if it is, I think the utils.test_prompt needs to be adjusted to expect the </s> at the beginning of the tokenized answer, in this case.

Happy to make these changes once you confirm what the expected behaviour is.

Thanks!

Direct path patching is like activation patching, but rather than patching in the output of component A, it acts on pairs of components A and B (in a layer after A). And we only patch in the output of A into the input of B, and all other components see the old output of A. I want to add a section to Exploratory Analysis Demo demonstrating this for all pairs of heads.

Eg to do direct path patching on the query of head B, we'd add a hook saying patched_B_query = original_B_query + (clean_A_output - corrupted_A_output) @ W_Q / layer_norm_scale

For reference, an old PR to add it an early version of the library #49

Line: 14
https://github.com/neelnanda-io/Easy-Transformer/blob/f2e85755e987a37f2968ce7eab1ad9215c0c8da9//easy_transformer/EasyTransformer.py#L11-L21

The QK and OV matrices can be accessed from model.QK and model.OV attributes, but not from the individual attention layers. Would be convenient to be able to do that, and not require computing all QK/OV matrices for all layers to access just one.

This project uses Patrick Kidger's torch typing to give tensor types including shapes. He recommends using JAXTyping for newer projects (not JAX specific, better maintained, more compatible with type checkers, etc) And it'd be good to update TL to use it.

@dkamm might be up your alley? Sorry if this invalidates your enum work, but might be a more elegant solution!

Add my interpretability-friendly models to HuggingFace (documented here: https://docs.google.com/document/d/1WONBzNqfKIxERejrrPlQMyKqg7jSFW92x5UMXNrMdPo/edit#heading=h.chq47zvs9cii )

This probably looks like adding the HookedTransformer + HookedTransformerConfig class to HuggingFace AutoModel, I'm not super sure how it works. Ideally this would be able to slot into code using HuggingFace models to eg evaluate them or generate text.

(Not actually a TransformerLens issue per se, but useful!)

Line: 355
https://github.com/neelnanda-io/Easy-Transformer/blob/f2e85755e987a37f2968ce7eab1ad9215c0c8da9//easy_transformer/EasyTransformer.py#L352-L362

I really need to document this one better lol - basically a wrapper around Python slice, but None -> slice(None) (equivalent to [:], doesn't add a dummy axis), and n -> [n] (for an integer n), which reduces the number of axes, which Python slices don't support.

Hi @neelnanda-io , do you want to have pre-commit in the requirements? If so, I could add that
or is the pre-commit-config.yml obsolete?

Currently, running run_with_hooks or run_with_cache ends by running model.reset_hooks(), which deletes all hooks. This means that if I, eg, want to create a model without positional embeddings, I can't add a hook that just sets pos_embed to zero, without breaking run_with_hooks or run_with_cache. The underlying problem is that PyTorch hooks are global state, but I want run_with_* to present them as local state to the user.

I'd like to add an add_perma_hook method to a HookPoint, which tracks hooks separately from add_hook, so that reset_hooks only deletes the normal hooks and not the perma-hooks (essentially creating a separate class of global state hooks and local state hooks)

As discussed here #161 , currently the testing

i) does not provide signal on which tests are failing
ii) seems to be incorrect? They pass locally for me

I'm a bit busy to do the good honest engineering required to improve this ATM :(

Currently from_pretrained has a bunch of Boolean flags for various simplifications to the transformer weights, and many default to True. I want these to be on by default, but it makes life a pain if you want to load a large model, or study a model exactly as the makers intended (since you need to set 5 ish Boolean flags to False, and it's not robust to new flags being added). I'd like there to be a from_pretrained_no_processing method with the same API as from_pretrained, which acts as a wrapper but sets all Boolean flags to False.

model = EasyTransformer.from_pretrained("gpt2").to(device)
returns None

This was working yesterday, so might be the new update.

Add tests that the tokenization methods work (to_tokens, to_string, to_str_tokens, get_token_position)

Go through the documentation and clarify things that are unclear (this is hard for me to do, so even just having someone new to the library flag confusions is helpful!) The behaviour of prepend_bos is the main confusion. Docs can be copied from https://colab.research.google.com/github/neelnanda-io/TransformerLens/blob/v2/Main_Demo.ipynb#scrollTo=GUSyRfQuKmHU

Add a helper function to implement activation patching (probably in utils.py). Crib from the code here: https://colab.research.google.com/github/neelnanda-io/Easy-Transformer/blob/main/Exploratory_Analysis_Demo.ipynb#scrollTo=JAyQI8Mlt3W7

I'd guess a good API is to input the clean cache, and corrupted prompt/tokens, specify which activation to patch and whether to iterate over positions, layers, head index, etc.

Redwood's IOI codebase is a good example of how you might implement this: https://github.com/redwoodresearch/Easy-Transformer/blob/main/easy_transformer/experiments.py

I was running into this issue where I couldn't commit anything. Turns out that because we previously had a pre-commit rule about using nbdev then removed it, because I'd previously installed pre-commit things broke. I fixed this by running rm -rf .git/hooks in the project root (the project doesn't have any other git hooks needed)

Line: 140
https://github.com/neelnanda-io/Easy-Transformer/blob/f2e85755e987a37f2968ce7eab1ad9215c0c8da9//easy_transformer/EasyTransformer.py#L137-L147

I'm trying to install the project via poetry install for python 3.9.

I get this error:

  PoetryException

  Failed to install /Users/rusheb/Library/Caches/pypoetry/artifacts/be/5d/9c/38ed00c38e66f11b3f1295c0b4fa2565c954b8e0c8d63deac26e996efa/y_py-0.5.5.tar.gz

  at /opt/homebrew/lib/python3.10/site-packages/poetry/utils/pip.py:58 in pip_install
       54│
       55│     try:
       56│         return environment.run_pip(*args)
       57│     except EnvCommandError as e:
    →  58│         raise PoetryException(f"Failed to install {path.as_posix()}") from e
       59│

  • Installing yarl (1.8.2)
Warning: The file chosen for install of y-py 0.5.5 (y_py-0.5.5.tar.gz) is yanked. Reason for being yanked: Inconsistent wheels

Output of poetry env info:

❯ poetry env info

Virtualenv
Python:         3.9.16
Implementation: CPython
Path:           /Users/rusheb/code/TransformerLens/.venv
Executable:     /Users/rusheb/code/TransformerLens/.venv/bin/python
Valid:          True

Also, I saw this warning at the top of the poetry install output:

Warning: poetry.lock is not consistent with pyproject.toml. You may be getting improper dependencies. Run `poetry lock [--no-update]` to fix it.

I think running poetry lock should fix it. I'll try to raise a fix now.

I was just merging this in, and this:

makes me a bit sad, as I have to think about 4 different activation functions in the otherwise really simple and clean file. Can all the activation functions be made modules instead? Happy to work on this if there's agreement

Originally posted by @ArthurConmy in #8 (comment)

I've received complaints that it needs better documentation! I would want to add clear info re the recommended name for each activation. The code is messy because it's designed to be robust to different names for things, to work for different sub-layers and for layers that aren't part of a block, etc.

Source: 170b0a0#diff-860fcdbbede66b9c1c31227cdfaf5c1cfef77faf1e9e0a376bc83633412ed47fR313

Hey,

If one unembeds from intermediate layers of the residual stream, the lack of normalization (coming from layer_norm_pre) leads to different (intuitively, less likely to be sensible) results.

I don't know if there is a warning about this / whether we might want to add a include_normalization flag to model.unembed(...)

Thanks!

Alex

Make this table better and cover key info for model architecture - whether it uses parallel attn & MLPs, and what positional embedding it is.

Add text at the bottom documenting the models more qualitatively, can basically copy this glossary: https://docs.google.com/document/d/1WONBzNqfKIxERejrrPlQMyKqg7jSFW92x5UMXNrMdPo/edit#heading=h.chq47zvs9cii

I'd want to add a separate table with training info: include training dataset, number of tokens, whether they were trained with dropout, whether they have checkpoints, whether trained with weight decay.

torchtyping used to give a function type signature is extremely verbose when printing the type, to somewhat ridiculous extents. Eg, here is model.run_with_cache?:

Signature:
model.run_with_cache(
    *model_args,
    return_cache_object=True,
    remove_batch_dim=False,
    **kwargs,
) -> Tuple[Union[NoneType, typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'pos', 'd_vocab',), 'cls_name': 'TensorType'}], typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ((),), 'cls_name': 'TensorType'}], typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'position - 1',), 'cls_name': 'TensorType'}], Tuple[typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'pos', 'd_vocab',), 'cls_name': 'TensorType'}], Union[typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ((),), 'cls_name': 'TensorType'}], typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'position - 1',), 'cls_name': 'TensorType'}]]]], Union[transformer_lens.ActivationCache.ActivationCache, Dict[str, torch.Tensor]]]
Docstring: Wrapper around run_with_cache in HookedRootModule. If return_cache_object is True, this will return an ActivationCache object, with a bunch of useful HookedTransformer specific methods, otherwise it will return a dictionary of activations as in HookedRootModule.

When torch.set_grad_enabled(False), you can overwrite the model's positional embeddings:

#%%
from transformer_lens.HookedTransformer import HookedTransformer
import torch
torch.set_grad_enabled(False)

#%%
model = HookedTransformer.from_pretrained("gpt2")

#%%
assert not torch.allclose(
    torch.norm(model.W_pos[0]), torch.tensor(0.0),
)

# %%
def sketchy_remove_pos_embed(z, hook):
    z[:] = 0.0
    return z

_ = model.run_with_hooks(
    "Hello, world",
    fwd_hooks = [("hook_pos_embed", sketchy_remove_pos_embed)],
)

model.reset_hooks()

# %%
assert torch.allclose(
    torch.norm(model.W_pos[0]), torch.tensor(0.0),
)

datasets, make deterministic

Line: 229
https://github.com/neelnanda-io/Easy-Transformer/blob/cf3b7b244dd79ce133a03cfdfa4cdcff074ab7c5//easy_transformer/experiments.py#L226-L236

Not sure exactly what the best way to cache is, I'd try to copy how HuggingFace transformers does it. The easiest way might be to just create a separate model on HuggingFace with the post processing weights and to pull the weights from that

Probably best to cache a version with all the default flags to from_pretrained and just use that by default, otherwise use the existing loading code.

I'd do this for small-ish and important models, eg all 4 GPT-2s and my interpretability friendly models.

By default, lines such as

logits, cache = model.run_with_cache( "Hello, Arthur", names_filter=lambda name: name in ["hook_pos_embed"] )

error since the following line requires hook_embed to be cached. Is this intended? It's a bit annoying editing all my NameFilters.

https://github.com/neelnanda-io/Easy-Transformer/blob/57ee0b8eee09e3d283170702afc052b5ad8643c6/easy_transformer/activation_cache.py#L34

good integration testing would run test_model on a GPU

We should be able to use the python script here: https://github.com/nojvek/vscode-ipynb-py-converter to

i) write notebooks as .py files with #%%
ii) have these automatically converted to .ipynb files on push

This means that a) we can easily test notebooks (since they are .py files, see discussion here) and b) have automatically updated URLs like https://colab.research.google.com/github/neelnanda-io/Easy-Transformer/blob/main/EasyTransformer_Demo.ipynb to display demos.

@alan-cooney this may be of interest since this would work well as another "on push" action like those in #68

In Main_Demo.ipynb, plotly's imshow does not correctly render in jupyter lab (for me) when running locally in jupyter lab.
An example is this line in cell 20

imshow(ioi_patching_result, x=token_labels, xaxis="Position", yaxis="Layer", title="Normalized Logit Difference After Patching Residual Stream on the IOI Task")

For me adding these two lines in the Setup section fixes it

But this also requires that DEVELOPMENT_MODE be True (which does not happen on its own).
Before contributing the fix I therefore wanted to check what's the intention behind DEVELOPMENT_MODE vs IN_COLAB?

Currently, these are not always inverse: running locally without changing anything gives DEVELOPMENT_MODE=False and IN_COLAB=False. Should we just have IN_COLAB and get rid of DEVELOPMENT_MODE?

I've been having issues setting up my dev environment.

OS: MacOS Monterey
Model Name: MacBook Air
Model Identifier: Mac14,2
Chip: Apple M2

What I was doing

• I cloned the repo
• ran poetry config virtualenvs.in-project true and poetry install --with dev
• after the last command, I got
  

What I tried to fix it

1. Upgrading pytorch

• Change pytorch version to 1.13.1
• This produced a new error, stack trace

2. Installing in a docker container

FROM ubuntu:latest

RUN apt-get -yqq update
RUN apt-get -yqq install git

RUN apt-get -yqq install python3    
RUN apt-get -yqq install python3-pip
RUN apt-get -yqq update

RUN git clone https://github.com/neelnanda-io/TransformerLens.git

WORKDIR /TransformerLens

RUN pip3 install poetry
RUN poetry config virtualenvs.in-project true
# RUN poetry install --with dev

This produced the same nvidia-cudnn-cu11 errors as before, which did not change even after I tried bumping the pytorch version.

Add support for having a model with layers split across several GPUs.

Make sure the layers (and its HookPoints) know what device they're on, so that hooks can ensure that they aren't needlessly moving information between GPUs. ActivationCache is a dictionary and should work by default.

The MVP would be doing 2 GPUs: putting the embed + first half of layers on GPU 1 and the second half + unembed on GPU 2. This is probably the most that's needed to support eg NeoX?

I'm not sure of the most elegant way of doing this, or how to do this without making the code really messy. I lean towards either adding a method which edits the model to move layers between devices, or making a separate ParallelHookedTransformer class

Line: 228
https://github.com/neelnanda-io/Easy-Transformer/blob/cf3b7b244dd79ce133a03cfdfa4cdcff074ab7c5//easy_transformer/experiments.py#L225-L235

Add a docstring at the top of each file in transformer_lens/ briefly summarising what it's about (so this shows up in VSCode on hover when importing it)

Add to evals.py support for checking how good a model is at tasks like Indirect Object Identification. Notably, where the eval involves generating a synthetic dataset of names, as they do in the paper (can copy from their codebase), running the model on them, and returning the average logit diff and accuracy.

Bonus:

• Add support for prompts of different token length
• Support multi-token answers.
• Write this in a generic way, where any generated dataset can be subbed in
• Give an option to pair up prompts, so "John and Mary went to the store, John gave the bag to" is followed by "John and Mary went to the store, Mary gave the bag to", to avoid biases where the model just favours common names

Yesterday @epurdy and I were working on an implementation of causal tracing from ROME. We ended up getting stuck for a while because when we ran the model it unexpectedly prepended a BOS token (in contrast to running the tokenizer by itself or with model.to_str_tokens, which doesn't). Looking through the codebase I noticed that of the six functions with a prepend_bos option, they're evenly split on the default. I think these defaults make sense in isolation but may be confusing when used together. Not sure what the right fix is here (or if it needs to be changed at all) but thought I'd at least raise this for discussion.

https://github.com/neelnanda-io/Easy-Transformer/blob/e5790469df3ebe26547c2016f899389e9f7bec30/easy_transformer/EasyTransformer.py#L146

https://github.com/neelnanda-io/Easy-Transformer/blob/e5790469df3ebe26547c2016f899389e9f7bec30/easy_transformer/EasyTransformer.py#L299

https://github.com/neelnanda-io/Easy-Transformer/blob/e5790469df3ebe26547c2016f899389e9f7bec30/easy_transformer/EasyTransformer.py#L345

https://github.com/neelnanda-io/Easy-Transformer/blob/e5790469df3ebe26547c2016f899389e9f7bec30/easy_transformer/EasyTransformer.py#L395

https://github.com/neelnanda-io/Easy-Transformer/blob/e5790469df3ebe26547c2016f899389e9f7bec30/easy_transformer/EasyTransformer.py#L935

https://github.com/neelnanda-io/Easy-Transformer/blob/e5790469df3ebe26547c2016f899389e9f7bec30/easy_transformer/utils.py#L440

Google have a very cool-looking tool for (mostly non-MI) interpretability of language models, called LIT. It seems designed to be framework agnostic, and to be able to take a wrapper around many kinds of models, with functions to enable various LIT functions. I want to add a wrapper to HookedTransformer such that it can integrate with LIT, ideally for as many LIT functions as possible.

The MVP in mind here would just be a Colab which gets LIT to work with TransformerLens, and maybe showing some things you can do with it. I'm not sure whether this kind of integration should actually be merged into the library, but I'd love for a small demo to exist!

Add tests that the methods in the ActivationCache class work correctly.

Go through the documentation and clarify things that are unclear (this is hard for me to do, so even just having someone new to the library flag confusions is helpful!)

Add tests that the FactoredMatrix class works (essentially that each of its methods correctly mimics the result for the actual matrix product).

Go through the documentation and clarify things that are unclear (this is hard for me to do, so even just having someone new to the library flag confusions is helpful!)

These lines test for equality but the error message is about divisibility:

https://github.com/neelnanda-io/TransformerLens/blob/c0f1c0e473f6cd1a4ca189af54818d6b64342b7c/transformer_lens/HookedTransformerConfig.py#L162-L165

This can lead to the misleading error message "d_model=16 is not divisible by n_heads=2 * d_head=4".

I'm guessing the check should change but I'm not sure.

Add the option to load models in bfloat16 and float16. Esp important for large models like GPT-J and GPT-NeoX.

Ideally, load from HuggingFace in this low precision, do weight processing on the CPU, and then move the processed model weights to the GPU. Might be easiest to do the weight processing once and caching to HF (see #103 )

Currently mypy type checks are not passing - with a large number of easy to resolve errors.

poetry run mypy transformer_lens

Resolving these would help make the codebase more robust, and the corresponding checks in /.github/workflows/checks.yml can then be enabled (they're commented out).

https://github.com/neelnanda-io/TransformerLens/blob/72b9be00c12b5be37baaa16b57e5df5e2aff3aad/.github/workflows/checks.yml#L58

Often you want to look at vectors over the vocabulary (eg the logits at a specific position). This is >50,000 dimensions and this is hard to interpret! I want there to be nice utils to visualize a vector like this.

An MVP would be a function mapping this to a pandas dataframe, with the token index, token string value, logit, log prob and probability. Either for just the top K, or for the entire vocab.

But I expect there's many ways to make something nice here! One option is to imitate nostalgebraist's graphing style for plot_logit_lens in `transformer_utils link. This takes a layer x position x d_vocab tensor, and visualises it as a layer x position heatmap, printing the string value of the top token in each cell, and colouring by the top token value.

Add something like Read The Docs to auto-generate docs

Bonus: Clean up docstrings so things aren't a mess

Model off of the BERT PR #82 - make a separate class for it with its own config.

Will involve adding support for cross-attention.

Code: https://github.com/huggingface/transformers/blob/v4.25.1/src/transformers/models/t5/modeling_t5.py#L1466
Paper: https://arxiv.org/abs/1910.10683