mirror of
https://github.com/wassname/activation_store.git
synced 2026-09-09 11:16:26 +08:00
thiswillbeyourgithub
a1dd6d3d85
Signed-off-by: thiswillbeyourgithub <26625900+thiswillbeyourgithub@users.noreply.github.com>
61 lines
1.7 KiB
Markdown
61 lines
1.7 KiB
Markdown
# activation_store
|
|
|
|
Utility library to persistently store transformer activations on disk as huggingface dataset.
|
|
|
|
As these activations can be quite large (layers x batch x sequence x hidden_size), generating them to disk helps avoid out of memory errors.
|
|
|
|
Install using
|
|
```
|
|
pip install git+https://github.com/wassname/activation_store.git
|
|
```
|
|
|
|
|
|
## Examples
|
|
|
|
Full examples can be found in the [nbs folder](./nbs).
|
|
|
|
|
|
```py
|
|
layer_groups = {'mlp.down_proj': [
|
|
'model.layers.21.mlp.down_proj',
|
|
'model.layers.22.mlp.down_proj',
|
|
'model.layers.23.mlp.down_proj'],
|
|
'self_attn': [
|
|
'model.layers.21.self_attn',
|
|
'model.layers.22.self_attn',
|
|
'model.layers.23.self_attn'],
|
|
'mlp.up_proj': [
|
|
'model.layers.21.mlp.up_proj',
|
|
'model.layers.22.mlp.up_proj',
|
|
'model.layers.23.mlp.up_proj']}
|
|
|
|
# collect activations into a huggingface dataset
|
|
f = activation_store(loader=ds, model=model, layers=layer_groups)
|
|
f
|
|
# > Generating train split: 0 examples [00:00, ? examples/s]
|
|
# Dataset({
|
|
# features: ['mlp.down_proj', 'self_attn', 'mlp.up_proj', 'loss', 'logits', 'hidden_states'],
|
|
# num_rows: 20
|
|
# })
|
|
|
|
# it has this shape
|
|
ds_a = Dataset.from_parquet(str(f)).with_format("torch")
|
|
ds_a[0:2]['hidden_states'].shape # [batch, layers, tokens, hidden_states]
|
|
# torch.Size([2, 25, 1, 896])
|
|
|
|
```
|
|
|
|
|
|
## Development
|
|
```
|
|
git clone https//github.com/wassname/activation_store.git
|
|
uv sync
|
|
```
|
|
|
|
## TODO:
|
|
|
|
- [x] test compression: it's not worth the [complexity](https://github.com/EleutherAI/elk/blob/84e99a36a5050881d85f1510a2486ce46ac1f942/elk/extraction/extraction.py#L382)
|
|
- [x] add examples
|
|
- [ ] generate and collect activations
|
|
- A manual loop of forwards/generate, reusing kv_cache, and appending model outputs along the token dim. saving outputs too
|