activation_store/nbs/example.ipynb

In [1]:

%reload_ext autoreload
%autoreload 2

In [2]:

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
# os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

In [3]:

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from activation_store.collect import activation_store

import torch

Load model

In [4]:

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

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="eager",  # flex_attention  flash_attention_2 sdpa eager
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.paddding_side = "left"
tokenizer.truncation_side = "left"

Load data and tokenize

In [5]:

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 [5]:

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

Data loader

In [6]:

from torch.utils.data import DataLoader
from transformers.data import DataCollatorForLanguageModeling
collate_fn = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
ds = DataLoader(ds2, batch_size=4, num_workers=0, collate_fn=collate_fn)
print(ds)

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

Collect activations

In [7]:

# choose layers to cache
layer_groups = {
    'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')],
    'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')],
    'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')],
}
layer_groups

Out [7]:

{'mlp.down_proj': ['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'],
 'self_attn': ['model.layers.0.self_attn',
  'model.layers.1.self_attn',
  'model.layers.2.self_attn',
  'model.layers.3.self_attn',
  'model.layers.4.self_attn',
  'model.layers.5.self_attn',
  'model.layers.6.self_attn',
  'model.layers.7.self_attn',
  'model.layers.8.self_attn',
  'model.layers.9.self_attn',
  'model.layers.10.self_attn',
  'model.layers.11.self_attn',
  'model.layers.12.self_attn',
  'model.layers.13.self_attn',
  'model.layers.14.self_attn',
  'model.layers.15.self_attn',
  'model.layers.16.self_attn',
  'model.layers.17.self_attn',
  'model.layers.18.self_attn',
  'model.layers.19.self_attn',
  'model.layers.20.self_attn',
  'model.layers.21.self_attn',
  'model.layers.22.self_attn',
  'model.layers.23.self_attn'],
 'mlp.up_proj': ['model.layers.0.mlp.up_proj',
  'model.layers.1.mlp.up_proj',
  'model.layers.2.mlp.up_proj',
  'model.layers.3.mlp.up_proj',
  'model.layers.4.mlp.up_proj',
  'model.layers.5.mlp.up_proj',
  'model.layers.6.mlp.up_proj',
  'model.layers.7.mlp.up_proj',
  'model.layers.8.mlp.up_proj',
  'model.layers.9.mlp.up_proj',
  'model.layers.10.mlp.up_proj',
  'model.layers.11.mlp.up_proj',
  'model.layers.12.mlp.up_proj',
  'model.layers.13.mlp.up_proj',
  'model.layers.14.mlp.up_proj',
  'model.layers.15.mlp.up_proj',
  'model.layers.16.mlp.up_proj',
  'model.layers.17.mlp.up_proj',
  'model.layers.18.mlp.up_proj',
  'model.layers.19.mlp.up_proj',
  'model.layers.20.mlp.up_proj',
  'model.layers.21.mlp.up_proj',
  'model.layers.22.mlp.up_proj',
  'model.layers.23.mlp.up_proj']}

In [14]:

f = activation_store(ds, model, layers=layer_groups)
f

Out [14]:

2025-03-14 16:42:30.982 | INFO     | activation_store.collect:activation_store:134 - creating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__0e7d5dbf1c73cf7d.parquet

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

PosixPath('/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__0e7d5dbf1c73cf7d.parquet')

In [15]:

from datasets import Dataset
ds_a = Dataset.from_parquet(str(f)).with_format("torch")
ds_a

Out [15]:

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
})

In [16]:

ds_a[0:2]['hidden_states'].shape # [batch, layers, tokens, hidden_states]

Out [16]:

torch.Size([2, 25, 1, 896])

In [ ]: