Merge pull request #767 from TransformerLensOrg/dev

v2.8.1

debugging/hf-tl-logit-comparator.ipynb

@@ -0,0 +1,265 @@
+{
+ "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",
+    "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
+}

transformer_lens/HookedTransformerConfig.py

@@ -181,6 +181,18 @@ class HookedTransformerConfig:
         output_logits_soft_cap (float): An optional softcap for output logits, currently only used
             in Gemma-2 (see attn_scores_soft_cap for details). Defaults to -1.0, which means not
             set.
+        use_NTK_by_parts_rope (bool): Whether to apply the "NTK-by-parts" method when using Rotary
+            Positional Embedding. This method adjusts the interpolation based on frequency factors
+            for different parts of the hidden dimensions. See Section 3.2 in
+            https://arxiv.org/pdf/2309.00071 for details. Defaults to False.
+        NTK_by_parts_low_freq_factor (float): The threshold applied to low-frequency hidden
+            dimensions during interpolation when using the "NTK-by-parts" method. Defaults to 1.0.
+        NTK_by_parts_high_freq_factor (float): The threshold applied to high-frequency hidden
+            dimensions during interpolation in the "NTK-by-parts" method. Defaults to 4.0.
+        NTK_by_parts_factor (float): The overall factor used in the "NTK-by-parts" method that
+            affects the rate of change between low and high-frequency interpolation strategies.
+            Defaults to 8.0.
+
 
     """
 
@@ -246,6 +258,10 @@ class HookedTransformerConfig:
     use_normalization_before_and_after: bool = False
     attn_scores_soft_cap: float = -1.0
     output_logits_soft_cap: float = -1.0
+    use_NTK_by_parts_rope: bool = False
+    NTK_by_parts_low_freq_factor: float = 1.0
+    NTK_by_parts_high_freq_factor: float = 4.0
+    NTK_by_parts_factor: float = 8.0
 
     def __post_init__(self):
         if self.n_heads == -1:

transformer_lens/components/abstract_attention.py

@@ -1,3 +1,4 @@
+import math
 from abc import ABC
 from typing import Dict, Optional, Tuple, Union
 
@@ -478,8 +479,33 @@ def calculate_sin_cos_rotary(
         pos = torch.arange(n_ctx, dtype=high_precision)
         dim = torch.arange(rotary_dim // 2, dtype=high_precision)
 
-        # A set of frequencies evenly spaced in log space
-        freq = base ** (dim / (rotary_dim / 2))
+        # Llama-3.1 uses NTK-by-Parts Rotary Embedding introduced in Section 3.2 in https://arxiv.org/pdf/2309.00071
+        # Implementation copied from https://github.com/huggingface/transformers/blob/v4.46.0/src/transformers/modeling_rope_utils.py#L310
+        if self.cfg.use_NTK_by_parts_rope:
+            inv_freq = 1.0 / (
+                base ** (torch.arange(0, rotary_dim, 2, dtype=torch.int64).float() / rotary_dim)
+            )
+            factor = self.cfg.NTK_by_parts_factor
+            low_freq_factor = self.cfg.NTK_by_parts_low_freq_factor
+            high_freq_factor = self.cfg.NTK_by_parts_high_freq_factor
+            old_context_len = n_ctx
+
+            low_freq_wavelen = old_context_len / low_freq_factor
+            high_freq_wavelen = old_context_len / high_freq_factor
+
+            wavelen = 2 * math.pi / inv_freq
+            inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
+            smooth_factor = (old_context_len / wavelen - low_freq_factor) / (
+                high_freq_factor - low_freq_factor
+            )
+            smoothed_inv_freq = (
+                1 - smooth_factor
+            ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
+            is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
+            inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
+            freq = 1 / inv_freq_llama
+        else:
+            freq = base ** (dim / (rotary_dim / 2))
         if self.cfg.rotary_adjacent_pairs:
             freq = einops.repeat(freq, "d -> (d 2)")
         else: