# graphpatch: a Python Library for Activation Patching

*This post is an announcement for a software library. It is likely only relevant to those working, or looking to start working, in* [*mechanistic interpretability*](https://www.lesswrong.com/posts/vnocLyeWXcAxtdDnP/a-comprehensive-mechanistic-interpretability-explainer-and)*.*

* * *

What is graphpatch?
===================

[graphpatch](https://graphpatch.dev) is a Python library for [activation patching](https://www.lesswrong.com/posts/xh85KbTFhbCz7taD4/how-to-think-about-activation-patching) on arbitrary PyTorch neural network models. It is designed to minimize the amount of boilerplate needed to run experiments making causal interventions on the intermediate activations in a model. It provides an intuitive [API](https://graphpatch.readthedocs.io/en/stable/api.html) based on the structure of a [torch.fx.Graph](https://pytorch.org/docs/stable/fx.html#torch.fx.Graph) representation compiled automatically from the original model. For a somewhat silly example, I can make Llama play [Taboo](https://www.lesswrong.com/posts/WBdvyyHLdxZSAMmoz/taboo-your-words) by [zero-ablating](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=IdovLFKgEm0lTfHT2YlzUnNT) its output for the token representing "Paris":

    with patchable_llama.patch(
      {"lm_head.output": ZeroPatch(slice=(slice(None), slice(None), 3681))}
    ):
      print(
        tokenizer.batch_decode(
          patchable_llama.generate(
            tokenizer(
              "The Eiffel Tower, located in",
              return_tensors="pt"
            ).input_ids,
            max_length=20,
            use_cache=False,
          )
        )
      )
    
    ["
 The Eiffel Tower, located in the heart of the French capital, is the most visited"]

Why is graphpatch?
==================

[graphpatch](https://graphpatch.dev) is a tool I wished had existed when I started my ~~descent into madness~~ entry into mechanistic interpretability with an [attempt](https://github.com/evan-lloyd/graphpatch/tree/main/demos) to replicate [ROME](https://rome.baulab.info/) on Llama. I hope that by reducing inconveniences ([trivial](https://www.lesswrong.com/posts/reitXJgJXFzKpdKyd/beware-trivial-inconveniences) and otherwise) I can both ease entry into the field and lower cognitive overhead for existing researchers. In particular, I want to make it easier to start running experiments on "off-the-shelf" models without the need to handle annoying setup—such as rewriting the model's Python code to expose intermediate activations—before even getting started. Thus, while [graphpatch](https://graphpatch.dev) should work equally well on custom-built research models, I focused on integration with the [Huggingface](https://huggingface.co/) ecosystem with:

*   Support for [PreTrainedModels](https://huggingface.co/docs/transformers/v4.41.3/en/main_classes/model#transformers.PreTrainedModel) loaded on [multiple devices or with CPU or disk offloading](https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/model#transformers.PreTrainedModel.from_pretrained.device_map)
*   Support for [8-bit quantization](https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/model#transformers.PreTrainedModel.from_pretrained.quantization_config) as provided by [bitsandbytes](https://huggingface.co/docs/bitsandbytes/main/en/index)
*   Support for [convenience](https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/text_generation#transformers.GenerationMixin) functions such as `generate()` in the above example

How do I graphpatch?
====================

[graphpatch](https://graphpatch.dev) is available on [PyPI](https://pypi.org/project/graphpatch/) and can be installed with `pip`:

    pip install graphpatch

You can read an [overview](https://github.com/evan-lloyd/graphpatch?tab=readme-ov-file#overview) on the [GitHub page](https://github.com/evan-lloyd/graphpatch) for the project. Full documentation is available on [Read the Docs](https://graphpatch.readthedocs.io/en/stable/index.html).

I have also provided a [Docker image](https://hub.docker.com/repository/docker/elloyd/graphpatch/general) that might be useful for starting your own mechanistic interpretability experiments on cloud hardware. See [this directory](https://github.com/evan-lloyd/graphpatch/tree/main/dev) for some scripts and notes on my development process, which may be adaptable to your own use case.