\n",
"See this dropdown for some functions which you might find helpful, and how to use them.
\n",
"\n",
"First, we can run some code to inspect our current memory usage. Here's me running this code during the exercise set on SAE circuits, after having already loaded in the Gemma models from the previous section. This was on a Colab Pro notebook.\n",
"\n",
"```python\n",
"# Profile memory usage, and delete gemma models if we've loaded them in\n",
"namespace = globals().copy() | locals()\n",
"part32_utils.profile_pytorch_memory(namespace=namespace, filter_device=\"cuda:0\")\n",
"```\n",
"\n",
"Allocated = 35.88 GB\n",
"Total = 39.56 GB\n",
"Free = 3.68 GB\n",
"┌──────────────────────┬────────────────────────┬──────────┬─────────────┐\n",
"│ Name │ Object │ Device │ Size (GB) │\n",
"├──────────────────────┼────────────────────────┼──────────┼─────────────┤\n",
"│ gemma_2_2b │ HookedSAETransformer │ cuda:0 │ 11.94 │\n",
"│ gpt2 │ HookedSAETransformer │ cuda:0 │ 0.61 │\n",
"│ gemma_2_2b_sae │ SAE │ cuda:0 │ 0.28 │\n",
"│ sae_resid_dirs │ Tensor (4, 24576, 768) │ cuda:0 │ 0.28 │\n",
"│ gpt2_sae │ SAE │ cuda:0 │ 0.14 │\n",
"│ logits │ Tensor (4, 15, 50257) │ cuda:0 │ 0.01 │\n",
"│ logits_with_ablation │ Tensor (4, 15, 50257) │ cuda:0 │ 0.01 │\n",
"│ clean_logits │ Tensor (4, 15, 50257) │ cuda:0 │ 0.01 │\n",
"│ _ │ Tensor (16, 128, 768) │ cuda:0 │ 0.01 │\n",
"│ clean_sae_acts_post │ Tensor (4, 15, 24576) │ cuda:0 │ 0.01 │\n",
"└──────────────────────┴────────────────────────┴──────────┴─────────────┘\n",
"\n",
"From this, we see that we've allocated a lot of memory for the the Gemma model, so let's delete it. We'll also run some code to move any remaining objects on the GPU which are larger than 100MB to the CPU, and print the memory status again.\n",
"\n",
"```python\n",
"del gemma_2_2b\n",
"del gemma_2_2b_sae\n",
"\n",
"THRESHOLD = 0.1 # GB\n",
"for obj in gc.get_objects():\n",
" try:\n",
" if isinstance(obj, torch.nn.Module) and part32_utils.get_tensors_size(obj) / 1024**3 > THRESHOLD:\n",
" if hasattr(obj, \"cuda\"):\n",
" obj.cpu()\n",
" if hasattr(obj, \"reset\"):\n",
" obj.reset()\n",
" except:\n",
" pass\n",
"\n",
"# Move our gpt2 model & SAEs back to GPU (we'll need them for the exercises we're about to do)\n",
"gpt2.to(device)\n",
"gpt2_saes = {layer: sae.to(device) for layer, sae in gpt2_saes.items()}\n",
"\n",
"part32_utils.print_memory_status()\n",
"```\n",
"\n",
"Allocated = 14.90 GB\n",
"Reserved = 39.56 GB\n",
"Free = 24.66\n",
"\n",
"Mission success! We've managed to free up a lot of memory. Note that the code which moves all objects collected by the garbage collector to the CPU is often necessary to free up the memory. We can't just delete the objects directly because PyTorch can still sometimes keep references to them (i.e. their tensors) in memory. In fact, if you add code to the for loop above to print out `obj.shape` when `obj` is a tensor, you'll see that a lot of those tensors are actually Gemma model weights, even once you've deleted `gemma_2_2b`.\n",
"\n",
" "
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b9eee9f4",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-30T22:46:07.085477Z",
"iopub.status.busy": "2026-07-30T22:46:07.085263Z",
"iopub.status.idle": "2026-07-30T22:46:07.088789Z",
"shell.execute_reply": "2026-07-30T22:46:07.087889Z"
},
"id": "cell-0004",
"language": "python",
"papermill": {
"duration": 0.010694,
"end_time": "2026-07-30T22:46:07.089668+00:00",
"exception": false,
"start_time": "2026-07-30T22:46:07.078974+00:00",
"status": "completed"
},
"tags": [
"parameters"
]
},
"outputs": [],
"source": [
"# Parameters - These will be injected by papermill during parameterized test runs\n",
"core_log_dir = None # Directory to save analysis logs (if None, a temp directory will be created)"
]
},
{
"cell_type": "markdown",
"id": "b5d7d21f",
"metadata": {
"id": "cell-0005",
"language": "markdown",
"papermill": {
"duration": 0.004877,
"end_time": "2026-07-30T22:46:07.099723+00:00",
"exception": false,
"start_time": "2026-07-30T22:46:07.094846+00:00",
"status": "completed"
},
"tags": []
},
"source": [
"#### Imports"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "472b5841",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-30T22:46:07.111147Z",
"iopub.status.busy": "2026-07-30T22:46:07.110946Z",
"iopub.status.idle": "2026-07-30T22:46:10.339782Z",
"shell.execute_reply": "2026-07-30T22:46:10.338685Z"
},
"id": "cell-0006",
"language": "python",
"papermill": {
"duration": 3.236535,
"end_time": "2026-07-30T22:46:10.341398+00:00",
"exception": false,
"start_time": "2026-07-30T22:46:07.104863+00:00",
"status": "completed"
},
"tags": []
},
"outputs": [],
"source": [
"# Core imports\n",
"import interpretune as it # registered analysis ops will be available as it.![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
\n",
"