activation_store/nbs/example.ipynb

In [1]:

%reload_ext autoreload
%autoreload 2

In [2]:

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from activation_store.collect import activation_store

import torch

Load model

In [3]:

model_name = "Qwen/Qwen2.5-0.5B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    attn_implementation="eager",  # flex_attention  flash_attention_2 sdpa eager
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

Load data and tokenize

In [4]:

N = 20
max_length = 256

imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)


def proc(row):
    messages = [
        {"role":"user", "content": row['prompt'] },
        {"role":"assistant", "content": row['chosen'] }
    ]
    return tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False, return_dict=True, max_length=max_length)

ds2 = imdb.map(proc).with_format("torch")
new_cols = set(ds2.column_names) - set(imdb.column_names)
ds2 = ds2.select_columns(new_cols)
ds2

Out [4]:

Dataset({
    features: ['input_ids', 'attention_mask'],
    num_rows: 20
})

Data loader

In [5]:

from torch.utils.data import DataLoader
def collate_fn(examples):
    # Pad the batch to max length within this batch
    return tokenizer.pad(
        examples,
        padding=True,
        return_tensors="pt",
    )
ds = DataLoader(ds2, batch_size=4, num_workers=0, collate_fn=collate_fn)
print(ds)

<torch.utils.data.dataloader.DataLoader object at 0x76557465f770>

Collect activations

In [7]:

# choose layers to cache
layers = [k for k,v in model.named_modules() if 'mlp.down_proj' in k]
layers

Out [7]:

['model.layers.0.mlp.down_proj',
 'model.layers.1.mlp.down_proj',
 'model.layers.2.mlp.down_proj',
 'model.layers.3.mlp.down_proj',
 'model.layers.4.mlp.down_proj',
 'model.layers.5.mlp.down_proj',
 'model.layers.6.mlp.down_proj',
 'model.layers.7.mlp.down_proj',
 'model.layers.8.mlp.down_proj',
 'model.layers.9.mlp.down_proj',
 'model.layers.10.mlp.down_proj',
 'model.layers.11.mlp.down_proj',
 'model.layers.12.mlp.down_proj',
 'model.layers.13.mlp.down_proj',
 'model.layers.14.mlp.down_proj',
 'model.layers.15.mlp.down_proj',
 'model.layers.16.mlp.down_proj',
 'model.layers.17.mlp.down_proj',
 'model.layers.18.mlp.down_proj',
 'model.layers.19.mlp.down_proj',
 'model.layers.20.mlp.down_proj',
 'model.layers.21.mlp.down_proj',
 'model.layers.22.mlp.down_proj',
 'model.layers.23.mlp.down_proj']

In [ ]:

f = activation_store(ds, model, layers=layers, writer_batch_size=10)
f

2025-02-15 21:58:37.654 | INFO     | activation_store.collect:activation_store:70 - creating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__9b3f4b0da96e9ad5.parquet

collecting activations:   0%|          | 0/5 [00:00<?, ?it/s]

You're using a Qwen2TokenizerFast 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.

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[8], line 2
      1 f = activation_store(ds, model, layers=layers, writer_batch_size=10)
----> 2 ds_a

NameError: name 'ds_a' is not defined

In [ ]:

from datasets import Dataset
Dataset.from_parquet(str(f))

Generating train split: 0 examples [00:00, ? examples/s]

Dataset({
    features: ['act-model.layers.0.mlp.down_proj', 'act-model.layers.1.mlp.down_proj', 'act-model.layers.2.mlp.down_proj', 'act-model.layers.3.mlp.down_proj', 'act-model.layers.4.mlp.down_proj', 'act-model.layers.5.mlp.down_proj', 'act-model.layers.6.mlp.down_proj', 'act-model.layers.7.mlp.down_proj', 'act-model.layers.8.mlp.down_proj', 'act-model.layers.9.mlp.down_proj', 'act-model.layers.10.mlp.down_proj', 'act-model.layers.11.mlp.down_proj', 'act-model.layers.12.mlp.down_proj', 'act-model.layers.13.mlp.down_proj', 'act-model.layers.14.mlp.down_proj', 'act-model.layers.15.mlp.down_proj', 'act-model.layers.16.mlp.down_proj', 'act-model.layers.17.mlp.down_proj', 'act-model.layers.18.mlp.down_proj', 'act-model.layers.19.mlp.down_proj', 'act-model.layers.20.mlp.down_proj', 'act-model.layers.21.mlp.down_proj', 'act-model.layers.22.mlp.down_proj', 'act-model.layers.23.mlp.down_proj', 'logits', 'hidden_states'],
    num_rows: 20
})

In [ ]:

In [ ]:

ds_a[0:2]['logits'].shape

In [ ]:

ds_a[0:2]['model.layers.0.mlp.down_proj'].shape

In [ ]: