Some bugs on Gemma Branch

[ShenzheZhu]

When I run the get_circuit of gemma-2-2b, my program stuck on def get_circuit() with an infinite loop.
I trace the code and find the problem is in jvp()

Also, the circuit.py is lacking from attribution import patching_effect, jvp

[saprmarks]

I'm very sorry about this; the code is long due for some upkeep. I'm guessing that what's happening is that you're hitting the breakpoint() in the try/except block because—and really I need to document this—the code assumes that you are on the sfc branch of dictionary_learning. So try doing the following:

1. Add the missing from attribution import patching_effect, jvp import you noted.
2. In your dictionary_learning module, git pull and then git checkout sfc.

Let me know if that works. If not, I'm planning on doing some upkeep on the code this weekend and can try to fix this then.

[ShenzheZhu]

unfortunately, it doesn't work, and I am stucking in same position as I mentioned before

(base) root@autodl-container-6d0f499972-33f98415:~/feature-circuits# sh scripts/get_circuit.sh
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████| 3/3 [00:02<00:00, 1.35it/s]
computing circuit
Loading Gemma SAEs: 96%|████████████████████████████████████████████████████████████▋ | 26/27 [02:05<00:04, 4.83s/it]
Batches: 0%| | 0/10 [00:00<?, ?it/s]You're using a GemmaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the __call__ method is faster than using a method to encode the text followed by a call to the pad method to get a padded encoding.
The 'batch_size' argument of HybridCache is deprecated and will be removed in v4.49. Use the more precisely named 'max_batch_size' argument instead.
The 'batch_size' attribute of HybridCache is deprecated and will be removed in v4.49. Use the more precisely named 'self.max_batch_size' attribute instead.
finished patching_effect, get into N

[saprmarks]

As an update, it looks like this is because an nnsight update broke something here. I'm actively working on a fix (which might require downgrading nnsight).

[saprmarks]

I've just pushed a few changes that I think have gotten this branch working again. If you're in a rush to use this repo, then I think you should be able to use it by:

• Pulling the changes
• Rerunning pip install -r requirements.txt
• Making sure that your clone of dictionary_learning is on the sfc branch
• (If you want to use graphviz to visualize the circuits)sudo apt-get install graphviz
• Using a command like scripts/get_circuit.sh google/gemma-2-2b rc_train 0.1 0.01 none to compute circuits

I'll be making some further changes to clean the code up, remove the need to use the sfc branch of dictionary_learning and better document how to use the repo. If you're not pressed for time, it might be worth waiting until then.

[ShenzheZhu]

Thank you very much, I will check the code latter