mirror of
https://github.com/wassname/activation_store.git
synced 2026-08-05 12:50:15 +08:00
700 lines
34 KiB
Plaintext
700 lines
34 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%reload_ext autoreload\n",
|
|
"%autoreload 2"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"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"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Load model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"model_name = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
|
|
"\n",
|
|
"model = AutoModelForCausalLM.from_pretrained(\n",
|
|
" model_name,\n",
|
|
" torch_dtype=\"auto\",\n",
|
|
" device_map=\"auto\",\n",
|
|
" attn_implementation=\"eager\", # flex_attention flash_attention_2 sdpa eager\n",
|
|
")\n",
|
|
"tokenizer = AutoTokenizer.from_pretrained(model_name)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Load data and tokenize"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"Dataset({\n",
|
|
" features: ['attention_mask', 'input_ids'],\n",
|
|
" num_rows: 10\n",
|
|
"})"
|
|
]
|
|
},
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"N = 10\n",
|
|
"max_length = 128\n",
|
|
"\n",
|
|
"imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)\n",
|
|
"\n",
|
|
"\n",
|
|
"def proc(row):\n",
|
|
" messages = [\n",
|
|
" {\"role\":\"user\", \"content\": row['prompt'] },\n",
|
|
" {\"role\":\"assistant\", \"content\": row['chosen'] }\n",
|
|
" ]\n",
|
|
" return tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False, return_dict=True, max_length=max_length)\n",
|
|
"\n",
|
|
"ds2 = imdb.map(proc).with_format(\"torch\")\n",
|
|
"new_cols = set(ds2.column_names) - set(imdb.column_names)\n",
|
|
"ds2 = ds2.select_columns(new_cols)\n",
|
|
"ds2"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Data loader"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"<torch.utils.data.dataloader.DataLoader object at 0x7089fb69f6e0>\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from torch.utils.data import DataLoader\n",
|
|
"def collate_fn(examples):\n",
|
|
" # Pad the batch to max length within this batch\n",
|
|
" return tokenizer.pad(\n",
|
|
" examples,\n",
|
|
" padding=True,\n",
|
|
" return_tensors=\"pt\",\n",
|
|
" )\n",
|
|
"ds = DataLoader(ds2, batch_size=4, num_workers=0, collate_fn=collate_fn)\n",
|
|
"print(ds)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Collect activations"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"['model.layers.0.mlp.down_proj',\n",
|
|
" 'model.layers.1.mlp.down_proj',\n",
|
|
" 'model.layers.2.mlp.down_proj',\n",
|
|
" 'model.layers.3.mlp.down_proj',\n",
|
|
" 'model.layers.4.mlp.down_proj',\n",
|
|
" 'model.layers.5.mlp.down_proj',\n",
|
|
" 'model.layers.6.mlp.down_proj',\n",
|
|
" 'model.layers.7.mlp.down_proj',\n",
|
|
" 'model.layers.8.mlp.down_proj',\n",
|
|
" 'model.layers.9.mlp.down_proj',\n",
|
|
" 'model.layers.10.mlp.down_proj',\n",
|
|
" 'model.layers.11.mlp.down_proj',\n",
|
|
" 'model.layers.12.mlp.down_proj',\n",
|
|
" 'model.layers.13.mlp.down_proj',\n",
|
|
" 'model.layers.14.mlp.down_proj',\n",
|
|
" 'model.layers.15.mlp.down_proj',\n",
|
|
" 'model.layers.16.mlp.down_proj',\n",
|
|
" 'model.layers.17.mlp.down_proj',\n",
|
|
" 'model.layers.18.mlp.down_proj',\n",
|
|
" 'model.layers.19.mlp.down_proj',\n",
|
|
" 'model.layers.20.mlp.down_proj',\n",
|
|
" 'model.layers.21.mlp.down_proj',\n",
|
|
" 'model.layers.22.mlp.down_proj',\n",
|
|
" 'model.layers.23.mlp.down_proj']"
|
|
]
|
|
},
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# choose layers to cache\n",
|
|
"layers = [k for k,v in model.named_modules() if 'mlp.down_proj' in k]\n",
|
|
"layers"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\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": "90a9936ab9f94893a77fc79bf972a04f",
|
|
"version_major": 2,
|
|
"version_minor": 0
|
|
},
|
|
"text/plain": [
|
|
"collecting activations: 0%| | 0/3 [00:00<?, ?it/s]"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"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.\n"
|
|
]
|
|
},
|
|
{
|
|
"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": [
|
|
"f = activation_store(ds, model, layers=layers, writer_batch_size=10)\n",
|
|
"f"
|
|
]
|
|
},
|
|
{
|
|
"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": "06fafa5231674f4da16d4ddfab520bd7",
|
|
"version_major": 2,
|
|
"version_minor": 0
|
|
},
|
|
"text/plain": [
|
|
"Generating train split: 0 examples [00:00, ? examples/s]"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"Dataset({\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": [
|
|
"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": 32,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"lost 0.01%\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"def float_to_int16(x: torch.Tensor) -> torch.Tensor:\n",
|
|
" \"\"\"Converts a floating point tensor to float16, then reinterprets as int16.\"\"\"\n",
|
|
" downcast = x.type(torch.float16)\n",
|
|
" # if not downcast.isfinite().all():\n",
|
|
" # raise ValueError(\"Cannot convert to 16 bit: values are not finite\")\n",
|
|
"\n",
|
|
" return downcast.view(torch.int16)\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.float16).type(torch.float32)\n",
|
|
"\n",
|
|
"\n",
|
|
"x = torch.randn(2, 3, 4)\n",
|
|
"x2 = float_to_int16(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": 41,
|
|
"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 float16_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 int16')\n",
|
|
" o[k] = float_to_int16(v)\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": 42,
|
|
"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='int16', id=None),\n",
|
|
" 'act-model.layers.1.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.2.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.3.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.4.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.5.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.6.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.7.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.8.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.9.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.10.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.11.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.12.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.13.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.14.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.15.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.16.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.17.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.18.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.19.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.20.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.21.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.22.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'act-model.layers.23.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n",
|
|
" 'logits': Array3D(shape=(-1, 453, 151936), dtype='int16', id=None),\n",
|
|
" 'hidden_states': Array4D(shape=(-1, 25, 453, 896), dtype='int16', id=None)}"
|
|
]
|
|
},
|
|
"execution_count": 42,
|
|
"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 = 'int16' 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": 43,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\u001b[32m2025-02-16 09:35:12.432\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m155\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": "1a2162645d9142c183936834d6815831",
|
|
"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 int16\n",
|
|
"act-model.layers.1.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.2.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.3.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.4.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.5.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.6.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.7.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.8.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.9.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.10.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.11.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.12.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.13.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.14.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.15.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.16.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.17.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.18.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.19.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.20.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.21.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.22.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"act-model.layers.23.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n",
|
|
"logits torch.float16 torch.Size([4, 453, 151936]) to int16\n",
|
|
"hidden_states torch.float32 torch.Size([4, 25, 453, 896]) to int16\n"
|
|
]
|
|
},
|
|
{
|
|
"ename": "TypeError",
|
|
"evalue": "Incompatible storage type list<item: list<item: int16>> for extension type extension<datasets.features.features.Array3DExtensionType<Array3DExtensionType>>",
|
|
"output_type": "error",
|
|
"traceback": [
|
|
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
|
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
|
|
"Cell \u001b[0;32mIn[43], 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[43mfeatures\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mFeatures\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_arrow_schema\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 4\u001b[0m \u001b[43m \u001b[49m\u001b[43mpostprocess_result\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfloat16_postprocess_result\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 5\u001b[0m f2\n\u001b[1;32m 6\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:175\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 171\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 172\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 173\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 174\u001b[0m )\n\u001b[0;32m--> 175\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 176\u001b[0m writer\u001b[38;5;241m.\u001b[39mfinalize()\n\u001b[1;32m 177\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:219\u001b[0m, in \u001b[0;36mTypedSequence.__arrow_array__\u001b[0;34m(self, type)\u001b[0m\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[1;32m 218\u001b[0m storage \u001b[38;5;241m=\u001b[39m to_pyarrow_listarray(data, pa_type)\n\u001b[0;32m--> 219\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mExtensionArray\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_storage\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstorage\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 221\u001b[0m \u001b[38;5;66;03m# efficient np array to pyarrow array\u001b[39;00m\n\u001b[1;32m 222\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, np\u001b[38;5;241m.\u001b[39mndarray):\n",
|
|
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:4354\u001b[0m, in \u001b[0;36mpyarrow.lib.ExtensionArray.from_storage\u001b[0;34m()\u001b[0m\n",
|
|
"\u001b[0;31mTypeError\u001b[0m: Incompatible storage type list<item: list<item: int16>> for extension type extension<datasets.features.features.Array3DExtensionType<Array3DExtensionType>>"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, \n",
|
|
" schema=schema,\n",
|
|
" features=Features.from_arrow_schema(schema),\n",
|
|
" postprocess_result=float16_postprocess_result)\n",
|
|
"f2\n",
|
|
"ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n",
|
|
"ds_a2.info"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 44,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# %debug"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"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",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"ds_a[0:2]['logits'].shape"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.12.3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|