dictionary_learning's People
Forkers
standardgalactic dronhazra g-w1 gavinratcliff techthiyanes jadenfiotto-kaufman rippletagging winnieyangwannan mikailkhona adamkarvonen rangell jannik-brinkmann empty-michael koayon migerovitch canrager kyoungrok0517 mbrukman alignment-lab-ai y12uc231 hannamw sun-wendy elanapearl jacksonkaunismaa ejmichauddictionary_learning's Issues
`EmptyStream` isn't used
Hi there, thanks for putting this library together! I noticed that you've defined EmptyStream
and even handle it in ActivationBuffer.__next__
, but it's never thrown. As an experiment I changed my dataset to have only a single point and this triggered an exception (which from my reading of the code it shouldn't have).
Traceback (most recent call last):
File "/home/joel/code/mamba-train/analyze_sae.py", line 180, in <module>
run(args)
File "/home/joel/code/mamba-train/analyze_sae.py", line 154, in run
result = evaluate(model, submodule, ae, buffer)
File "/home/joel/code/mamba-train/dictionary_learning/evaluation.py", line 82, in evaluate
acts = next(activations).to(device)
File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 103, in __next__
self.refresh()
File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 223, in refresh
self._refresh_std()
File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 156, in _refresh_std
tokens = self.tokenized_batch()["input_ids"]
File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 129, in tokenized_batch
texts = self.text_batch(batch_size=batch_size)
File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 123, in text_batch
return [next(self.data) for _ in range(batch_size)]
File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 123, in <listcomp>
return [next(self.data) for _ in range(batch_size)]
StopIteration
JumpReLU SAE training on the roadmap?
Just curious if anyone is thinking about implementing a training pipeline for JumpReLU SAEs! They have a couple of properties which are really desirable for something I'm working on.
`ActivationBuffer` used inconsistently re running out of data.
When you run trainSAE
, it expects to run through all of the data because the line for step, acts in enumerate(tqdm(activations, total=steps))
simply exits the loop when StopIteration
is raised. But evaluate
simply calls acts = next(activations).to(device)
, so it fails if the data runs out. I think ideally this case should be handled gracefully, but it should at least be documented.
Code:
from nnsight import LanguageModel
from dictionary_learning import ActivationBuffer
from dictionary_learning.training import trainSAE
from dictionary_learning.evaluation import evaluate
model = LanguageModel("EleutherAI/pythia-70m-deduped")
submodule = model.gpt_neox.layers[1].mlp
activation_dim = 512
dictionary_size = 16 * activation_dim
data = iter(
[
"This is some example data",
"In real life, for training a dictionary",
"you would need much more data than this",
]
)
buffer = ActivationBuffer(
data,
model,
submodule,
out_feats=activation_dim,
n_ctxs=3e4,
)
ae = trainSAE(
buffer,
activation_dim,
dictionary_size,
lr=3e-4,
sparsity_penalty=1e-3,
)
buffer = ActivationBuffer(
iter(["hello", "world"]),
model,
submodule,
out_feats=activation_dim,
n_ctxs=300,
)
print(evaluate(model, submodule, ae, buffer))
Result:
joel@simplex ~/c/mamba-sae (main)> python3.11 test.py
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
0it [00:00, ?it/s]
Traceback (most recent call last):
File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 94, in text_batch
return [
^
File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 95, in <listcomp>
next(self.data) for _ in range(batch_size)
^^^^^^^^^^^^^^^
StopIteration
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/Users/joel/code/mamba-sae/test.py", line 46, in <module>
print(evaluate(model, submodule, ae, buffer))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/joel/code/mamba-sae/dictionary_learning/evaluation.py", line 119, in evaluate
acts = next(activations).to(device)
^^^^^^^^^^^^^^^^^
File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 76, in __next__
self.refresh()
File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 164, in refresh
self._refresh_std()
File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 121, in _refresh_std
with self.model.invoke(self.text_batch(), truncation=True, max_length=self.ctx_len) as invoker:
^^^^^^^^^^^^^^^^^
File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 98, in text_batch
raise StopIteration("End of data stream reached")
StopIteration: End of data stream reached
`attention_mask` is required
b2bd5f0 added a requirement to have attention_mask
(in the dict returned from the tokenizer, if I understand correctly). My tokenizer / model doesn't have this. Should it be necessary?
Tensors must have same number of dimensions: got 2 and 3
Using basically the example code from README.md:
from nnsight import LanguageModel
from dictionary_learning import ActivationBuffer
from dictionary_learning.training import trainSAE
from dictionary_learning.evaluation import evaluate
model = LanguageModel("EleutherAI/pythia-70m-deduped")
submodule = model.gpt_neox.layers[1].mlp
activation_dim = 512
dictionary_size = 16 * activation_dim
data = iter(
[
"This is some example data",
"In real life, for training a dictionary",
"you would need much more data than this",
]
)
buffer = ActivationBuffer(
data,
model,
submodule,
out_feats=activation_dim,
n_ctxs=3e4,
)
ae = trainSAE(
buffer,
activation_dim,
dictionary_size,
lr=3e-4,
sparsity_penalty=1e-3,
)
Output:
(env) root@12702e9f3980:/workspace/mamba-sae# python3 test.py
Resolving data files: 100%|██████████████████████████████████████████| 30/30 [00:06<00:00, 4.89it/s]
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
0it [00:00, ?it/s]You're using a GPTNeoXTokenizerFast 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.
0it [00:17, ?it/s]
Traceback (most recent call last):
File "/workspace/mamba-sae/test.py", line 33, in <module>
ae = trainSAE(
File "/workspace/mamba-sae/dictionary_learning/training.py", line 180, in trainSAE
for step, acts in enumerate(tqdm(activations, total=steps)):
File "/workspace/mamba-sae/env/lib/python3.10/site-packages/tqdm/std.py", line 1181, in __iter__
for obj in iterable:
File "/workspace/mamba-sae/dictionary_learning/buffer.py", line 76, in __next__
self.refresh()
File "/workspace/mamba-sae/dictionary_learning/buffer.py", line 164, in refresh
self._refresh_std()
File "/workspace/mamba-sae/dictionary_learning/buffer.py", line 130, in _refresh_std
self.activations = t.cat([self.activations, hidden_states.value.to(self.device)], dim=0)
RuntimeError: Tensors must have same number of dimensions: got 2 and 3
The problem is that self.activations.shape
is torch.Size([0, 512])
while hidden_states.value.shape
is torch.Size([512, 128, 512])
. I'm not sure which is right.
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.