Skip to content

Files

Latest commit

Janos Kramarcopybara-github
Janos Kramar
and
Jan 17, 2025
2ed8201 · Jan 17, 2025

History

History
48 lines (32 loc) · 2.64 KB

README.md

File metadata and controls

48 lines (32 loc) · 2.64 KB

Mishax

Introduction

Mishax is a utility library for mechanistic interpretability research, with its motivations explained in this blog post. It enables users to do 2 things:

mishax.ast_patcher enables running code from some other library (e.g. a deep learning codebase) with some source-level code modifications applied. For mechanistic interpretability this can be used to stick probes in the model and intervene at arbitrary locations. This otherwise requires forking the code that’s being modified, but that comes with more maintenance requirements.

mishax.safe_greenlet, given a complicated function f that allows running arbitrary callbacks somewhere deep inside (e.g. using Flax’s intercept_methods), enables transforming it into an ordinary-looking Python for loop that iterates over internal values and allows them to be replaced with other values. Behind the scenes, this will run f in a kind of separate “thread” –- but the user can mostly ignore that, and use the loop to read and write representations into the model during a forward pass, in a way that interoperates well with the rest of JAX.

In mishax.examples.gemma you can find an example of instrumenting an LLM codebase; it's pointed at the https://github.com/google-deepmind/gemma reference implementation of Gemma.

Note

ast_patcher relies on code transformations of the target code, which violates some usual abstractions. Careless use may reduce codebase maintainability -- AST patching is best deployed in moderation and with care. For more details, see the ModuleASTPatcher docstring.

Setup

python3 -m venv $HOME/mishax-venv
source $HOME/mishax-venv/bin/activate
python3 -m pip install git+git://github.com/google-deepmind/mishax.git

With deps for the Gemma example:

python3 -m venv $HOME/mishax-venv
source $HOME/mishax-venv/bin/activate
python3 -m pip install git+git://github.com/google-deepmind/mishax.git[gemma]

To deactivate the virtual environment, run deactivate.

Run tests

source $HOME/mishax-venv/bin/activate
python3 -m pip install git+git://github.com/google-deepmind/mishax.git[dev]
python3 -m mishax.ast_patcher_test
python3 -m mishax.safe_greenlet_test
python3 -m mishax.examples.gemma_test

Colab Tutorial

A colab notebook demonstrating how to instrument Gemma internals with mishax is available here: colab link