Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make the pipeline support training multiple components #163

Merged
merged 16 commits into from
Dec 22, 2023
Merged

Make the pipeline support training multiple components #163

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
a78bb40
Make the pipeline support multiple components
alan-cooney Dec 20, 2023
55cf8de
Change the weep to receive multiple cache names
alan-cooney Dec 20, 2023
54b70fe
Update sweep params
alan-cooney Dec 20, 2023
a5dace8
Fix optimizer related tests
alan-cooney Dec 20, 2023
b92c1b8
Fix activation resampler unit methods
alan-cooney Dec 21, 2023
5132810
Fix parameter type
alan-cooney Dec 21, 2023
5414888
Fix all activation resampler tests
alan-cooney Dec 21, 2023
bfb5d4d
Fix linting
alan-cooney Dec 21, 2023
8774c8b
Fix param
alan-cooney Dec 21, 2023
aaa76e0
Fix snaps
alan-cooney Dec 21, 2023
99b187b
Make all tests multi-component for activation resampler
alan-cooney Dec 21, 2023
7198830
Fix same forward pass test
alan-cooney Dec 21, 2023
29b95fc
Move to consistent num prefix
alan-cooney Dec 21, 2023
d0bd10a
Fix remaining tests
alan-cooney Dec 21, 2023
6b272bd
Fix warnings
alan-cooney Dec 21, 2023
9b7e543
Fix metric names
alan-cooney Dec 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .vscode/cspell.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@
"polysemanticity",
"precommit",
"pretokenized",
"pydantic",
"pyproject",
"pyright",
"pytest",
Expand Down
51 changes: 15 additions & 36 deletions docs/content/demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,18 +38,9 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"The autoreload extension is already loaded. To reload it, use:\n",
" %reload_ext autoreload\n"
]
}
],
"outputs": [],
"source": [
"# Check if we're in Colab\n",
"try:\n",
Expand All @@ -71,7 +62,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -89,7 +80,6 @@
" sweep,\n",
" SweepConfig,\n",
")\n",
"import wandb\n",
"\n",
"\n",
"os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
Expand All @@ -113,27 +103,19 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"SweepConfig(parameters=Hyperparameters(\n",
" source_data=SourceDataHyperparameters(dataset_path=Parameter(value=NeelNanda/c4-code-tokenized-2b), context_size=Parameter(value=128))\n",
" source_model=SourceModelHyperparameters(name=Parameter(value=gelu-2l), hook_site=Parameter(value=mlp_out), hook_layer=Parameter(value=0), hook_dimension=Parameter(value=512), dtype=Parameter(value=float32))\n",
" activation_resampler=ActivationResamplerHyperparameters(resample_interval=Parameter(value=197885952), max_n_resamples=Parameter(value=4), n_activations_activity_collate=Parameter(value=98942976), resample_dataset_size=Parameter(value=819200), threshold_is_dead_portion_fires=Parameter(value=1e-06))\n",
" autoencoder=AutoencoderHyperparameters(expansion_factor=Parameter(value=2))\n",
" loss=LossHyperparameters(l1_coefficient=Parameter(max=0.01, min=0.004))\n",
" optimizer=OptimizerHyperparameters(lr=Parameter(max=0.001, min=1e-05), adam_beta_1=Parameter(value=0.9), adam_beta_2=Parameter(value=0.99), adam_weight_decay=Parameter(value=0.0), amsgrad=Parameter(value=False), fused=Parameter(value=False))\n",
" pipeline=PipelineHyperparameters(log_frequency=Parameter(value=100), source_data_batch_size=Parameter(value=16), train_batch_size=Parameter(value=8192), max_store_size=Parameter(value=2998272), max_activations=Parameter(value=1999847424), checkpoint_frequency=Parameter(value=47972352), validation_frequency=Parameter(value=99999744), validation_number_activations=Parameter(value=8192))\n",
" random_seed=Parameter(value=49)\n",
"), method=<Method.RANDOM: 'random'>, metric=Metric(name=train/loss/total_loss, goal=minimize), command=None, controller=None, description=None, earlyterminate=None, entity=None, imageuri=None, job=None, kind=None, name=None, program=None, project=None)"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
"ename": "TypeError",
"evalue": "SourceModelHyperparameters.__init__() got an unexpected keyword argument 'hook_layer'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[3], line 12\u001b[0m\n\u001b[1;32m 1\u001b[0m sweep_config \u001b[38;5;241m=\u001b[39m SweepConfig(\n\u001b[1;32m 2\u001b[0m parameters\u001b[38;5;241m=\u001b[39mHyperparameters(\n\u001b[1;32m 3\u001b[0m activation_resampler\u001b[38;5;241m=\u001b[39mActivationResamplerHyperparameters(\n\u001b[1;32m 4\u001b[0m threshold_is_dead_portion_fires\u001b[38;5;241m=\u001b[39mParameter(\u001b[38;5;241m1e-6\u001b[39m),\n\u001b[1;32m 5\u001b[0m ),\n\u001b[1;32m 6\u001b[0m loss\u001b[38;5;241m=\u001b[39mLossHyperparameters(\n\u001b[1;32m 7\u001b[0m l1_coefficient\u001b[38;5;241m=\u001b[39mParameter(\u001b[38;5;28mmax\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-2\u001b[39m, \u001b[38;5;28mmin\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m4e-3\u001b[39m),\n\u001b[1;32m 8\u001b[0m ),\n\u001b[1;32m 9\u001b[0m optimizer\u001b[38;5;241m=\u001b[39mOptimizerHyperparameters(\n\u001b[1;32m 10\u001b[0m lr\u001b[38;5;241m=\u001b[39mParameter(\u001b[38;5;28mmax\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-3\u001b[39m, \u001b[38;5;28mmin\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-5\u001b[39m),\n\u001b[1;32m 11\u001b[0m ),\n\u001b[0;32m---> 12\u001b[0m source_model\u001b[38;5;241m=\u001b[39m\u001b[43mSourceModelHyperparameters\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 13\u001b[0m \u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mParameter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mgelu-2l\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 14\u001b[0m \u001b[43m \u001b[49m\u001b[43mcache_names\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mParameter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmlp_out\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 15\u001b[0m \u001b[43m \u001b[49m\u001b[43mhook_layer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mParameter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 16\u001b[0m \u001b[43m \u001b[49m\u001b[43mhook_dimension\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mParameter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m512\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 17\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 18\u001b[0m source_data\u001b[38;5;241m=\u001b[39mSourceDataHyperparameters(\n\u001b[1;32m 19\u001b[0m dataset_path\u001b[38;5;241m=\u001b[39mParameter(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNeelNanda/c4-code-tokenized-2b\u001b[39m\u001b[38;5;124m\"\u001b[39m),\n\u001b[1;32m 20\u001b[0m ),\n\u001b[1;32m 21\u001b[0m ),\n\u001b[1;32m 22\u001b[0m method\u001b[38;5;241m=\u001b[39mMethod\u001b[38;5;241m.\u001b[39mRANDOM,\n\u001b[1;32m 23\u001b[0m )\n\u001b[1;32m 24\u001b[0m sweep_config\n",
"\u001b[0;31mTypeError\u001b[0m: SourceModelHyperparameters.__init__() got an unexpected keyword argument 'hook_layer'"
]
}
],
"source": [
Expand All @@ -150,8 +132,7 @@
" ),\n",
" source_model=SourceModelHyperparameters(\n",
" name=Parameter(\"gelu-2l\"),\n",
" hook_site=Parameter(\"mlp_out\"),\n",
" hook_layer=Parameter(0),\n",
" cache_names=Parameter([\"blocks.0.hook_mlp_out\", \"blocks.1.hook_mlp_out\"]),\n",
" hook_dimension=Parameter(512),\n",
" ),\n",
" source_data=SourceDataHyperparameters(\n",
Expand All @@ -176,9 +157,7 @@
"metadata": {},
"outputs": [],
"source": [
"sweep(sweep_config=sweep_config)\n",
"\n",
"wandb.finish()"
"sweep(sweep_config=sweep_config)"
]
}
],
Expand Down
Loading