Fit random directions

curt-tigges · Mar 25, 2024 · 1a8f882 · 1a8f882
1 parent 037cbc9
commit 1a8f882
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -32,4 +32,5 @@ RUN apt install wget -y
 RUN wget https://github.com/wkhtmltopdf/packaging/releases/download/0.12.6-1/wkhtmltox_0.12.6-1.focal_amd64.deb
 RUN apt install -f ./wkhtmltox_0.12.6-1.focal_amd64.deb -y
 RUN pip install imgkit
+RUN pip install pytest transformers_stream_generator tiktoken
 
diff --git a/TransformerLens b/TransformerLens
diff --git a/random_directions.py b/random_directions.py
@@ -8,7 +8,9 @@
 # %%
 def generate_random_directions(d_model, n_layers, model_name):
     torch.random.manual_seed(42)
-    for layer in range(n_layers + 1):
+    bar = tqdm(range(n_layers + 1), total=n_layers + 1)
+    for layer in bar:
+        bar.set_description(f"Layer {layer}")
         random_direction = torch.randn(d_model)
         random_direction /= random_direction.norm()
         save_array(random_direction, f"random_direction_layer{layer:02d}", model_name)
@@ -26,10 +28,18 @@ def generate_random_directions(d_model, n_layers, model_name):
     # 'EleutherAI/pythia-410m',
     # 'EleutherAI/pythia-1.4b',
     # 'EleutherAI/pythia-2.8b',
+    # "gemma-7b",
+    # "gemma-2b",
+    "qwen-7b",
+    "qwen-1.8b",
 ]
 for model in tqdm(MODELS):
-    model = HookedTransformer.from_pretrained(model)
+    print(f"Loading model {model}")
+    model = HookedTransformer.from_pretrained(model, dtype="bfloat16")
     d_model = model.cfg.d_model
     n_layers = model.cfg.n_layers
+    print(f"Generating random directions for {model.cfg.model_name}")
     generate_random_directions(d_model, n_layers, model.cfg.model_name)
+    del model
+    torch.cuda.empty_cache()
 # %%
+3 −2		.vscode/cspell.json
+3 −2		.vscode/extensions.json
+29 −11		.vscode/settings.json
+126 −146		demos/LLaMA.ipynb
+0 −588		demos/Mistral.ipynb
+376 −0		demos/Qwen.ipynb
+2 −0		docs/source/conf.py
+1,240 −1,370		poetry.lock
+91 −105		pyproject.toml
+79 −1		tests/acceptance/test_hooked_transformer.py
+5 −1		tests/unit/test_grouped_query_attention.py
+154 −89		transformer_lens/HookedTransformer.py
+3 −0		transformer_lens/HookedTransformerConfig.py
+109 −29		transformer_lens/components.py
+654 −30		transformer_lens/loading_from_pretrained.py
+5 −1		transformer_lens/utils.py