diff --git a/debugging/hf-tl-logit-comparator.ipynb b/debugging/hf-tl-logit-comparator.ipynb new file mode 100644 index 000000000..ee445c397 --- /dev/null +++ b/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 +} diff --git a/transformer_lens/HookedTransformerConfig.py b/transformer_lens/HookedTransformerConfig.py index e2fdc532e..4458705de 100644 --- a/transformer_lens/HookedTransformerConfig.py +++ b/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: diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 3146de0c2..a2a831e9f 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/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: diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py index 0b8489976..49dffbf04 100644 --- a/transformer_lens/loading_from_pretrained.py +++ b/transformer_lens/loading_from_pretrained.py @@ -875,6 +875,7 @@ def convert_hf_model_config(model_name: str, **kwargs): "rotary_dim": 128, "final_rms": True, "gated_mlp": True, + "rotary_base": 500000.0, } elif "Meta-Llama-3-70B" in official_model_name: cfg_dict = { @@ -894,6 +895,7 @@ def convert_hf_model_config(model_name: str, **kwargs): "rotary_dim": 128, "final_rms": True, "gated_mlp": True, + "rotary_base": 500000.0, } elif "Llama-3.2-1B" in official_model_name: cfg_dict = { @@ -913,6 +915,11 @@ def convert_hf_model_config(model_name: str, **kwargs): "rotary_dim": 64, "final_rms": True, "gated_mlp": True, + "rotary_base": 500000.0, + "use_NTK_by_parts_rope": True, + "NTK_by_parts_low_freq_factor": 1.0, + "NTK_by_parts_high_freq_factor": 4.0, + "NTK_by_parts_factor": 32.0, } elif "Llama-3.2-3B" in official_model_name: cfg_dict = { @@ -932,6 +939,11 @@ def convert_hf_model_config(model_name: str, **kwargs): "rotary_dim": 128, "final_rms": True, "gated_mlp": True, + "rotary_base": 500000.0, + "use_NTK_by_parts_rope": True, + "NTK_by_parts_low_freq_factor": 1.0, + "NTK_by_parts_high_freq_factor": 4.0, + "NTK_by_parts_factor": 32.0, } elif "Llama-3.1-8B" in official_model_name: cfg_dict = { @@ -951,6 +963,11 @@ def convert_hf_model_config(model_name: str, **kwargs): "rotary_dim": 128, "final_rms": True, "gated_mlp": True, + "rotary_base": 500000.0, + "use_NTK_by_parts_rope": True, + "NTK_by_parts_low_freq_factor": 1.0, + "NTK_by_parts_high_freq_factor": 4.0, + "NTK_by_parts_factor": 8.0, } elif "Llama-3.1-70B" in official_model_name: cfg_dict = { @@ -970,6 +987,11 @@ def convert_hf_model_config(model_name: str, **kwargs): "rotary_dim": 128, "final_rms": True, "gated_mlp": True, + "rotary_base": 500000.0, + "use_NTK_by_parts_rope": True, + "NTK_by_parts_low_freq_factor": 1.0, + "NTK_by_parts_high_freq_factor": 4.0, + "NTK_by_parts_factor": 8.0, } elif architecture == "GPTNeoForCausalLM": cfg_dict = {