wassname/activation_store
1
0
Fork 0
mirror of https://github.com/wassname/activation_store.git synced 2026-08-11 11:14:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

arrayXd? no

Browse Source 
it looks like they are only for conversion to int, which is complex
This commit is contained in:
wassname
2025-02-16 09:26:00 +08:00
parent 8a61bfeba0
commit 86c85c9e76
3 changed files with 533 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
activation_store/collect.py
+119 -37
View File
@@ -1,47 +1,112 @@
from transformers import AutoModelForCausalLM
import torch
from datasets import Dataset
from tqdm.auto import tqdm
from torch.utils.data import DataLoader
from loguru import logger
import gc
from pathlib import Path
from baukit.nethook import TraceDict, recursive_copy
from einops import rearrange
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.fingerprint import Hasher
from transformers.modeling_outputs import ModelOutput
from activation_store.helpers.torch import clear_mem
from typing import Dict, Generator
import copy
import torch
from baukit.nethook import TraceDict, recursive_copy
from datasets import Dataset
from datasets.arrow_writer import ParquetWriter
from datasets.fingerprint import Hasher
from einops import rearrange
from loguru import logger
from torch import Tensor
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoModelForCausalLM
from transformers.modeling_outputs import ModelOutput
default_output_folder = (Path(__file__).parent.parent / "outputs").resolve()
def default_postprocess_result(input: dict, trace: TraceDict, output: ModelOutput) -> Dict[str, Tensor]:
def clear_mem():
gc.collect()
torch.cuda.empty_cache()
gc.collect()
def default_postprocess_result(
input: dict, trace: TraceDict, output: ModelOutput, model: AutoModelForCausalLM
) -> Dict[str, Tensor]:
"""add activations to output, and rearrange hidden states"""
# Baukit records the literal layer output, which varies by model. Here we assume that the output or the first part are activations we want
acts = {f'act-{k}':
v.output[0] if isinstance(v.output, tuple) else v.output
for k, v in trace.items()}
acts = {
f"act-{k}": v.output[0] if isinstance(v.output, tuple) else v.output
for k, v in trace.items()
}
output.hidden_states = rearrange(list(output.hidden_states), "l b t h -> b l t h")
o = dict(
attention_mask=input["attention_mask"],
**acts, **output
)
return o
def to_cpu(x):
"""
Trys to convert torch if possible a single item
"""
if isinstance(x, torch.Tensor):
x = x.cpu()
return x
else:
return x
def recursive_copy2(x, clone=None, detach=None, retain_grad=None):
"""
from baukit with addition of deep copy for non tensors
output.hidden_states = rearrange(list(output.hidden_states), 'l b t h -> b l t h')
return dict(**acts, **output)
Copies a reference to a tensor, or an object that contains tensors,
optionally detaching and cloning the tensor(s). If retain_grad is
true, the original tensors are marked to have grads retained.
"""
if not clone and not detach and not retain_grad:
return x
if isinstance(x, torch.Tensor):
if retain_grad:
if not x.requires_grad:
x.requires_grad = True
x.retain_grad()
elif detach:
x = x.detach()
if clone:
x = x.clone()
return x
# Only dicts, lists, and tuples (and subclasses) can be copied.
if isinstance(x, dict):
return type(x)({k: recursive_copy(v, clone=clone, detach=detach, retain_grad=retain_grad) for k, v in x.items()})
elif isinstance(x, (list, tuple)):
return type(x)([recursive_copy(v, clone=clone, detach=detach, retain_grad=retain_grad) for v in x])
else:
return copy.deepcopy(x)
@torch.no_grad
def generate_batches(loader: DataLoader, model: AutoModelForCausalLM, layers, postprocess_result=default_postprocess_result) -> Generator[Dict[str, Tensor], None, None]:
def generate_batches(
loader: DataLoader,
model: AutoModelForCausalLM,
layers,
postprocess_result=default_postprocess_result,
) -> Generator[Dict[str, Tensor], None, None]:
model.eval()
for batch in tqdm(loader, 'collecting activations'):
for batch in tqdm(loader, "collecting activations"):
device = next(model.parameters()).device
with torch.amp.autocast(device_type=device.type):
with TraceDict(model, layers) as trace:
out = model(**batch, use_cache=False, output_hidden_states=True, return_dict=True)
o = postprocess_result(batch, trace, out)
out = model(
**batch,
use_cache=False,
output_hidden_states=True,
return_dict=True,
)
o = postprocess_result(batch, trace, out, model)
# copy to avoid memory leaks
o = recursive_copy(o)
for k in o:
if not isinstance(o[k], torch.Tensor):
print('o', k, type(o[k]))
o = {k: to_cpu(v) for k, v in o.items()}
o = recursive_copy(o, clone=True, detach=True)
out = trace = batch = None
clear_mem()
yield o
@@ -52,7 +117,17 @@ def dataset_hash(**kwargs):
return suffix
def activation_store(loader: DataLoader, model: AutoModelForCausalLM, dataset_name='', layers=[], dataset_dir=default_output_folder, writer_batch_size=1, postprocess_result=default_postprocess_result) -> Dataset:
def activation_store(
loader: DataLoader,
model: AutoModelForCausalLM,
dataset_name="",
layers=[],
dataset_dir=default_output_folder,
writer_batch_size=1,
postprocess_result=default_postprocess_result,
features=None,
schema=None,
) -> Dataset:
"""
Collect activations from a model and store them in a dataset
@@ -76,20 +151,27 @@ def activation_store(loader: DataLoader, model: AutoModelForCausalLM, dataset_na
f.parent.mkdir(exist_ok=True, parents=True)
logger.info(f"creating dataset {f}")
iterator = generate_batches(loader, model, layers=layers, postprocess_result=postprocess_result)
iterator = generate_batches(
loader, model, layers=layers, postprocess_result=postprocess_result
)
with ParquetWriter(path=f, writer_batch_size=writer_batch_size,
embed_local_files=True
) as writer:
# batch_1 = next(iterator)
# Features.encode_batch(batch_1)
# features = Features({'x': Array2D(shape=(1, 3), dtype='int32')})
with ParquetWriter(
path=f, writer_batch_size=writer_batch_size, embed_local_files=True,
features=features, schema=schema,
) as writer:
# writer.write_batch(batch_1)
for bo in iterator:
bs = len(next(iter(bo.values())))
assert all(len(v) == bs for v in bo.values()), f"must return Dict[str,Tensor] and all tensors with same batch size a first dimension"
# or maybe better compression to `writer.write(example, key)` for each
assert all(len(v) == bs for v in bo.values()), (
"must return Dict[str,Tensor] and all tensors with same batch size a first dimension"
)
writer.write_batch(bo)
writer.finalize()
writer.finalize()
writer.close()
# ds = Dataset.from_file(str(f)).with_format("torch")
return f
activation_store/helpers/torch.py
-8
View File
@@ -1,8 +0,0 @@
import torch
import gc
def clear_mem():
gc.collect()
torch.cuda.empty_cache()
gc.collect()
nbs/example.ipynb
+414 -45
View File
@@ -18,10 +18,10 @@
"source": [
"from datasets import load_dataset\n",
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"from datasets import Dataset\n",
"import torch\n",
"\n",
"from activation_store.collect import activation_store\n",
"\n",
"import torch"
"from activation_store.collect import activation_store\n"
]
},
{
@@ -64,8 +64,8 @@
"data": {
"text/plain": [
"Dataset({\n",
" features: ['input_ids', 'attention_mask'],\n",
" num_rows: 20\n",
" features: ['attention_mask', 'input_ids'],\n",
" num_rows: 10\n",
"})"
]
},
@@ -75,8 +75,8 @@
}
],
"source": [
"N = 20\n",
"max_length = 256\n",
"N = 10\n",
"max_length = 128\n",
"\n",
"imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)\n",
"\n",
@@ -110,7 +110,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"<torch.utils.data.dataloader.DataLoader object at 0x76557465f770>\n"
"<torch.utils.data.dataloader.DataLoader object at 0x7089fb69f6e0>\n"
]
}
],
@@ -136,7 +136,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 6,
"metadata": {},
"outputs": [
{
@@ -168,7 +168,7 @@
" 'model.layers.23.mlp.down_proj']"
]
},
"execution_count": 7,
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
@@ -181,25 +181,25 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2025-02-15 21:58:37.654\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m70\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__9b3f4b0da96e9ad5.parquet\u001b[0m\n"
"\u001b[32m2025-02-16 09:16:55.292\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m122\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet\u001b[0m\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "fe95a697e5c0432e85d15707b07fd001",
"model_id": "90a9936ab9f94893a77fc79bf972a04f",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"collecting activations: 0%| | 0/5 [00:00<?, ?it/s]"
"collecting activations: 0%| | 0/3 [00:00<?, ?it/s]"
]
},
"metadata": {},
@@ -213,15 +213,14 @@
]
},
{
"ename": "NameError",
"evalue": "name 'ds_a' is not defined",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[8], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m f \u001b[38;5;241m=\u001b[39m activation_store(ds, model, layers\u001b[38;5;241m=\u001b[39mlayers, writer_batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mds_a\u001b[49m\n",
"\u001b[0;31mNameError\u001b[0m: name 'ds_a' is not defined"
]
"data": {
"text/plain": [
"PosixPath('/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet')"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
@@ -233,11 +232,25 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "1e1429e1d3224a2b8a5398f7a414911d",
"model_id": "06fafa5231674f4da16d4ddfab520bd7",
"version_major": 2,
"version_minor": 0
},
@@ -252,19 +265,352 @@
"data": {
"text/plain": [
"Dataset({\n",
" 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'],\n",
" num_rows: 20\n",
" features: ['attention_mask', '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'],\n",
" num_rows: 10\n",
"})"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# load\n",
"ds_a = Dataset.from_parquet(str(f)).with_format(\"torch\")\n",
"ds_a"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"DatasetInfo(description='', citation='', homepage='', license='', features={'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None), 'act-model.layers.0.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.1.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.2.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.3.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.4.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.5.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.6.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.7.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.8.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.9.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.10.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.11.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.12.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.13.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.14.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.15.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.16.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.17.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.18.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.19.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.20.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.21.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.22.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.23.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'logits': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'hidden_states': Sequence(feature=Sequence(feature=Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None), length=-1, id=None), length=-1, id=None)}, post_processed=None, supervised_keys=None, builder_name='parquet', dataset_name='parquet', config_name='default', version=0.0.0, splits={'train': SplitInfo(name='train', num_bytes=1391398926, num_examples=10, shard_lengths=[4, 6], dataset_name='parquet')}, download_checksums={'/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet': {'num_bytes': 1363203837, 'checksum': None}}, download_size=1363203837, post_processing_size=None, dataset_size=1391398926, size_in_bytes=2754602763)"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from datasets import Dataset\n",
"Dataset.from_parquet(str(f))"
"ds_a.info"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"torch.Size([2, 25, 453, 896])"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ds_a[0:2]['hidden_states'].shape"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"torch.Size([2, 453, 896])"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ds_a[0:2]['act-model.layers.0.mlp.down_proj'].shape"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"ename": "ZeroDivisionError",
"evalue": "division by zero",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mZeroDivisionError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[9], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\n",
"\u001b[0;31mZeroDivisionError\u001b[0m: division by zero"
]
}
],
"source": [
"1/0"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## With dtypes compression - wip"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"lost 2.02%\n"
]
}
],
"source": [
"def float_to_int8(x: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"Converts a floating point tensor to float16, then reinterprets as int16.\"\"\"\n",
" downcast = x.type(torch.float8_e4m3fn)\n",
" # if not downcast.isfinite().all():\n",
" # raise ValueError(\"Cannot convert to 16 bit: values are not finite\")\n",
"\n",
" return downcast.view(torch.int8)\n",
"\n",
"def int8_to_float32(x: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"Converts int16 to float16, then reinterprets as float32.\"\"\"\n",
" return x.view(torch.float8_e4m3fn).type(torch.float32)\n",
"\n",
"\n",
"x = torch.randn(2, 3, 4)\n",
"x2 = float_to_int8(x)\n",
"x3 = int8_to_float32(x2)\n",
"assert torch.isfinite(x3).all()\n",
"assert torch.allclose(x, x3, rtol=1e-1)\n",
"d = ((x-x3)/x).abs().mean()\n",
"print(f'lost {d:.2%}')"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {},
"outputs": [],
"source": [
"from activation_store.collect import default_postprocess_result\n",
"from datasets.features.features import cast_to_python_objects\n",
"# o = cast_to_python_objects(o, only_1d_for_numpy=True, optimize_list_casting=False)\n",
"\n",
"def float8_postprocess_result(\n",
" input, trace, output, model\n",
"):\n",
" o = default_postprocess_result(input, trace, output, model)\n",
" # o = cast_to_python_objects(o, only_1d_for_numpy=False, optimize_list_casting=False)\n",
"\n",
" for k, v in o.items():\n",
" if k=='attention_mask':\n",
" o[k] = v.to(torch.int8)\n",
" if isinstance(v, torch.Tensor) and torch.is_floating_point(v):\n",
" print(k, v.dtype, v.shape, 'to int8')\n",
" o[k] = float_to_int8(v.float())\n",
" else:\n",
" print('no conv', k, type(v))\n",
" # o = {k: float_to_int8(v) if isinstance(v, torch.Tensor) else v\n",
" # for k, v in o.items()}\n",
" return o"
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None),\n",
" 'act-model.layers.0.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.1.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.2.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.3.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.4.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.5.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.6.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.7.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.8.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.9.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.10.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.11.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.12.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.13.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.14.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.15.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.16.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.17.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.18.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.19.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.20.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.21.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.22.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'act-model.layers.23.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
" 'logits': Array3D(shape=(-1, 453, 151936), dtype='int8', id=None),\n",
" 'hidden_states': Array4D(shape=(-1, 25, 453, 896), dtype='int8', id=None)}"
]
},
"execution_count": 28,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from datasets.arrow_writer import OptimizedTypedSequence, _ArrayXDExtensionType\n",
"from datasets.features.features import Features, Array2D, Array3D, Array4D, Array5D\n",
"\n",
"# manually build features\n",
"optimized_int_type_by_col = {\n",
" \"attention_mask\": \"int8\", # binary tensor\n",
" \"special_tokens_mask\": \"int8\",\n",
" \"input_ids\": \"int32\", # typical vocab size: 0-50k (max ~500k, never > 1M)\n",
" \"token_type_ids\": \"int8\", # binary mask; some (XLNetModel) use an additional token represented by a 2\n",
"}\n",
"\n",
"def build_schema(d):\n",
" inferred_features = Features()\n",
" cols = d.keys()\n",
" for col in cols:\n",
" x = d[col]\n",
" if col in optimized_int_type_by_col:\n",
" dtype = optimized_int_type_by_col[col]\n",
" typed_sequence = OptimizedTypedSequence(x, col=col)\n",
" inferred_features[col] = typed_sequence.get_inferred_type()\n",
" else:\n",
" if x.ndim == 1:\n",
" inferred_features[col] = OptimizedTypedSequence(x, col=col)\n",
" inferred_features[col] = typed_sequence.get_inferred_type()\n",
" shape=(-1,)+x.shape[1:]\n",
" dtype = 'int8' if x.dtype == torch.float32 else x.dtype\n",
" if x.ndim == 2:\n",
" cls = Array2D\n",
" elif x.ndim == 3:\n",
" cls = Array3D\n",
" elif x.ndim == 4:\n",
" cls = Array4D\n",
" elif x.ndim == 5:\n",
" cls = Array5D\n",
" else:\n",
" raise ValueError(f\"Unsupported number of dimensions: {x.ndim}\")\n",
" inferred_features[col] = cls(dtype=dtype, shape=shape)\n",
" return inferred_features.arrow_schema\n",
" # Features.from_arrow_schema(schema)\n",
"\n",
"d = ds_a[0:2]\n",
"schema = build_schema(d)\n",
"schema\n",
"Features.from_arrow_schema(schema)"
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2025-02-16 09:25:08.798\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m152\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__c6184d05bf03be61.parquet\u001b[0m\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6fc1c814733a4a7ab65468f0d2ad0b2b",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"collecting activations: 0%| | 0/3 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"no conv attention_mask <class 'torch.Tensor'>\n",
"act-model.layers.0.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.1.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.2.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.3.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.4.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.5.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.6.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.7.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.8.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.9.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.10.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.11.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.12.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.13.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.14.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.15.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.16.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.17.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.18.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.19.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.20.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.21.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.22.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"act-model.layers.23.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
"logits torch.float16 torch.Size([4, 453, 151936]) to int8\n",
"hidden_states torch.float32 torch.Size([4, 25, 453, 896]) to int8\n"
]
},
{
"ename": "ArrowTypeError",
"evalue": "Could not convert tensor([[ 48, 25, -83, ..., 31, 45, 41],\n [ -76, -100, -94, ..., 26, -84, -117],\n [ -97, 26, 15, ..., -97, -107, -109],\n ...,\n [ 44, -94, -104, ..., -110, 18, 27],\n [ -77, 26, -77, ..., -100, 33, 43],\n [ -98, 22, -111, ..., -110, 14, -107]], dtype=torch.int8) with type Tensor: was not a sequence or recognized null for conversion to list type",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mArrowTypeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[29], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m f2 \u001b[38;5;241m=\u001b[39m \u001b[43mactivation_store\u001b[49m\u001b[43m(\u001b[49m\u001b[43mds\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlayers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlayers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[43m \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[43mpostprocess_result\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfloat8_postprocess_result\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 4\u001b[0m f2\n\u001b[1;32m 5\u001b[0m ds_a2 \u001b[38;5;241m=\u001b[39m Dataset\u001b[38;5;241m.\u001b[39mfrom_parquet(\u001b[38;5;28mstr\u001b[39m(f2))\u001b[38;5;241m.\u001b[39mwith_format(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtorch\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/activation_store/collect.py:172\u001b[0m, in \u001b[0;36mactivation_store\u001b[0;34m(loader, model, dataset_name, layers, dataset_dir, writer_batch_size, postprocess_result, features, schema)\u001b[0m\n\u001b[1;32m 168\u001b[0m bs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(bo\u001b[38;5;241m.\u001b[39mvalues())))\n\u001b[1;32m 169\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mall\u001b[39m(\u001b[38;5;28mlen\u001b[39m(v) \u001b[38;5;241m==\u001b[39m bs \u001b[38;5;28;01mfor\u001b[39;00m v \u001b[38;5;129;01min\u001b[39;00m bo\u001b[38;5;241m.\u001b[39mvalues()), (\n\u001b[1;32m 170\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmust return Dict[str,Tensor] and all tensors with same batch size a first dimension\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 171\u001b[0m )\n\u001b[0;32m--> 172\u001b[0m \u001b[43mwriter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbo\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 173\u001b[0m writer\u001b[38;5;241m.\u001b[39mfinalize()\n\u001b[1;32m 174\u001b[0m writer\u001b[38;5;241m.\u001b[39mclose()\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:605\u001b[0m, in \u001b[0;36mArrowWriter.write_batch\u001b[0;34m(self, batch_examples, writer_batch_size)\u001b[0m\n\u001b[1;32m 603\u001b[0m col_try_type \u001b[38;5;241m=\u001b[39m try_features[col] \u001b[38;5;28;01mif\u001b[39;00m try_features \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m col \u001b[38;5;129;01min\u001b[39;00m try_features \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 604\u001b[0m typed_sequence \u001b[38;5;241m=\u001b[39m OptimizedTypedSequence(col_values, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mcol_type, try_type\u001b[38;5;241m=\u001b[39mcol_try_type, col\u001b[38;5;241m=\u001b[39mcol)\n\u001b[0;32m--> 605\u001b[0m arrays\u001b[38;5;241m.\u001b[39mappend(\u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtyped_sequence\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m 606\u001b[0m inferred_features[col] \u001b[38;5;241m=\u001b[39m typed_sequence\u001b[38;5;241m.\u001b[39mget_inferred_type()\n\u001b[1;32m 607\u001b[0m schema \u001b[38;5;241m=\u001b[39m inferred_features\u001b[38;5;241m.\u001b[39marrow_schema \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpa_writer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mschema\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:252\u001b[0m, in \u001b[0;36mpyarrow.lib.array\u001b[0;34m()\u001b[0m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:114\u001b[0m, in \u001b[0;36mpyarrow.lib._handle_arrow_array_protocol\u001b[0;34m()\u001b[0m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:218\u001b[0m, in \u001b[0;36mTypedSequence.__arrow_array__\u001b[0;34m(self, type)\u001b[0m\n\u001b[1;32m 215\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 216\u001b[0m \u001b[38;5;66;03m# custom pyarrow types\u001b[39;00m\n\u001b[1;32m 217\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(pa_type, _ArrayXDExtensionType):\n\u001b[0;32m--> 218\u001b[0m storage \u001b[38;5;241m=\u001b[39m \u001b[43mto_pyarrow_listarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpa_type\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 219\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m pa\u001b[38;5;241m.\u001b[39mExtensionArray\u001b[38;5;241m.\u001b[39mfrom_storage(pa_type, storage)\n\u001b[1;32m 221\u001b[0m \u001b[38;5;66;03m# efficient np array to pyarrow array\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/features/features.py:1591\u001b[0m, in \u001b[0;36mto_pyarrow_listarray\u001b[0;34m(data, pa_type)\u001b[0m\n\u001b[1;32m 1589\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m any_np_array_to_pyarrow_listarray(data, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mpa_type\u001b[38;5;241m.\u001b[39mvalue_type)\n\u001b[1;32m 1590\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1591\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpa_type\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstorage_dtype\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:372\u001b[0m, in \u001b[0;36mpyarrow.lib.array\u001b[0;34m()\u001b[0m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:42\u001b[0m, in \u001b[0;36mpyarrow.lib._sequence_to_array\u001b[0;34m()\u001b[0m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/error.pxi:155\u001b[0m, in \u001b[0;36mpyarrow.lib.pyarrow_internal_check_status\u001b[0;34m()\u001b[0m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/error.pxi:92\u001b[0m, in \u001b[0;36mpyarrow.lib.check_status\u001b[0;34m()\u001b[0m\n",
"\u001b[0;31mArrowTypeError\u001b[0m: Could not convert tensor([[ 48, 25, -83, ..., 31, 45, 41],\n [ -76, -100, -94, ..., 26, -84, -117],\n [ -97, 26, 15, ..., -97, -107, -109],\n ...,\n [ 44, -94, -104, ..., -110, 18, 27],\n [ -77, 26, -77, ..., -100, 33, 43],\n [ -98, 22, -111, ..., -110, 14, -107]], dtype=torch.int8) with type Tensor: was not a sequence or recognized null for conversion to list type"
]
}
],
"source": [
"f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, \n",
" schema=schema,\n",
" postprocess_result=float8_postprocess_result)\n",
"f2\n",
"ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n",
"ds_a2.info"
]
},
{
@@ -272,7 +618,46 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
"source": [
"f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, postprocess_result=float8_postprocess_result)\n",
"f2"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from datasets import Dataset\n",
"# load\n",
"ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n",
"for c in ds_a2.column_names[1:]:\n",
" print(c)\n",
" ds_a2[c] = int8_to_float32(ds_a2[0:-1][c])\n",
"# ds_a2 = int8_to_float32(ds_a)\n",
"ds_a2.info"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"d = ds_a2[:][c]\n",
"print(c)\n",
"d.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ds_a2.info"
]
},
{
"cell_type": "code",
@@ -282,22 +667,6 @@
"source": [
"ds_a[0:2]['logits'].shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ds_a[0:2]['model.layers.0.mlp.down_proj'].shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {