Complete type checking for OLMo support (builds on #816)

debugging/hf-tl-logit-comparator.ipynb

@@ -0,0 +1,266 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Logit Comparator for HuggingFace and TransformerLens Outputs\n",
+    "This notebook is a quick and dirty tool to compare the logit outputs of a HuggingFace model and a TransformerLens model via several different metrics. It is intended to help debug issues with the TransformerLens model, such as bugs in the model's implementation. If you identify any issues, please open an issue on the [GitHub repository](https://github.com/TransformerLensOrg/TransformerLens)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "from transformer_lens import HookedTransformer\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "if torch.backends.mps.is_available():\n",
+    "    device = \"mps\"\n",
+    "else:\n",
+    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "\n",
+    "# NBVAL_IGNORE_OUTPUT\n",
+    "_ = torch.set_grad_enabled(False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Comparator Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"EleutherAI/pythia-2.8b\"  # You can change this to any model name\n",
+    "sentence = \"The quick brown fox\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import login\n",
+    "login(token=\"\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get Transformers Logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "\n",
+    "def load_model(model_name=\"gpt2\"):\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "    model = AutoModelForCausalLM.from_pretrained(model_name)\n",
+    "    return model, tokenizer\n",
+    "\n",
+    "def get_logits(model, tokenizer, sentence, device):\n",
+    "    # Tokenize the input sentence\n",
+    "    inputs = tokenizer(sentence, return_tensors=\"pt\")\n",
+    "    \n",
+    "    # Move inputs to the device\n",
+    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
+    "    \n",
+    "    # Generate the logits\n",
+    "    with torch.no_grad():\n",
+    "        outputs = model(**inputs)\n",
+    "    \n",
+    "    # Get the logits for all tokens\n",
+    "    logits = outputs.logits\n",
+    "    \n",
+    "    return logits\n",
+    "\n",
+    "model, tokenizer = load_model(model_name)\n",
+    "model = model.to(device)\n",
+    "\n",
+    "hf_logits = get_logits(model, tokenizer, sentence, device)[:, -1, :]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get TransformerLens Logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = HookedTransformer.from_pretrained_no_processing(model_name, device=device)\n",
+    "tokens = model.to_tokens(sentence, prepend_bos=False)\n",
+    "tl_logits = model(tokens)[:, -1, :]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Compare Logit Distributions\n",
+    "Various metrics are used to compare the logit distributions of the two models. We don't yet have standard values for what constitutes a \"good\" logit comparison, so we are working on establishing benchmarks."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"HF Logits Shape: {hf_logits.shape}\")\n",
+    "print(f\"TL Logits Shape: {tl_logits.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Tensor Comparison"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "are_close = torch.allclose(tl_logits, hf_logits, rtol=1e-5, atol=1e-3)\n",
+    "print(f\"Are the logits close? {are_close}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Mean Squared Error"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Compare the logits with MSE\n",
+    "mse = torch.nn.functional.mse_loss(hf_logits, tl_logits)\n",
+    "print(f\"MSE: {mse}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Maximum Absolute Difference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "max_diff = torch.max(torch.abs(tl_logits - hf_logits))\n",
+    "print(f\"Max Diff: {max_diff}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Cosine Similarity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cosine_sim = F.cosine_similarity(tl_logits, hf_logits, dim=-1).mean()\n",
+    "print(f\"Cosine Sim: {cosine_sim}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### KL Divergence"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def kl_div(logits1: torch.Tensor, logits2: torch.Tensor) -> torch.Tensor:\n",
+    "    probs1 = F.softmax(logits1, dim=-1)\n",
+    "    probs2 = F.softmax(logits2, dim=-1)\n",
+    "    return F.kl_div(probs1.log(), probs2, reduction='batchmean')\n",
+    "\n",
+    "kl_tl_hf = kl_div(tl_logits, hf_logits)\n",
+    "kl_hf_tl = kl_div(hf_logits, tl_logits)\n",
+    "print(f\"KL(TL||HF): {kl_tl_hf}\")\n",
+    "print(f\"KL(HF||TL): {kl_hf_tl}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sae-l",
+   "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.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

demos/ARENA_Content.ipynb

@@ -40,8 +40,8 @@
     "\n",
     "    ipython = get_ipython()\n",
     "    # Code to automatically update the HookedTransformer code as its edited without restarting the kernel\n",
-    "    # ipython.run_line_magic(\"load_ext\", \"autoreload\")\n",
-    "    # ipython.run_line_magic(\"autoreload\", \"2\")\n",
+    "    ipython.run_line_magic(\"load_ext\", \"autoreload\")\n",
+    "    ipython.run_line_magic(\"autoreload\", \"2\")\n",
     "\n",
     "if IN_GITHUB or IN_COLAB:\n",
     "    %pip install torch\n",

demos/Activation_Patching_in_TL_Demo.ipynb

@@ -158,7 +158,8 @@
     }
    ],
    "source": [
-    "torch.set_grad_enabled(False)"
+    "# NBVAL_IGNORE_OUTPUT\n",
+    "_ = torch.set_grad_enabled(False)"
    ]
   },
   {

demos/BERT.ipynb

@@ -10,7 +10,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -36,17 +35,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Using renderer: colab\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh\n",
     "import plotly.io as pio\n",
@@ -108,22 +99,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<torch.autograd.grad_mode.set_grad_enabled at 0x2a285a790>"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "torch.set_grad_enabled(False)"
+    "# NBVAL_IGNORE_OUTPUT\n",
+    "_ = torch.set_grad_enabled(False)"
    ]
   },
   {