Skip to content

Commit 8f482fc

Browse files
bryce13950bryce13950
authored
Merge pull request #767 from TransformerLensOrg/dev
v2.8.1
2 parents b6e19d6 + c7837fb commit 8f482fc

File tree

4 files changed

+331
-2
lines changed

4 files changed

+331
-2
lines changed

‎debugging/hf-tl-logit-comparator.ipynb

Lines changed: 265 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,265 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Logit Comparator for HuggingFace and TransformerLens Outputs\n",
8+
"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)."
9+
]
10+
},
11+
{
12+
"cell_type": "code",
13+
"execution_count": null,
14+
"metadata": {},
15+
"outputs": [],
16+
"source": [
17+
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
18+
"from transformer_lens import HookedTransformer\n",
19+
"import torch\n",
20+
"import torch.nn.functional as F\n",
21+
"\n",
22+
"if torch.backends.mps.is_available():\n",
23+
" device = \"mps\"\n",
24+
"else:\n",
25+
" device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
26+
"\n",
27+
"torch.set_grad_enabled(False)"
28+
]
29+
},
30+
{
31+
"cell_type": "markdown",
32+
"metadata": {},
33+
"source": [
34+
"## Comparator Setup"
35+
]
36+
},
37+
{
38+
"cell_type": "code",
39+
"execution_count": 51,
40+
"metadata": {},
41+
"outputs": [],
42+
"source": [
43+
"model_name = \"EleutherAI/pythia-2.8b\" # You can change this to any model name\n",
44+
"sentence = \"The quick brown fox\""
45+
]
46+
},
47+
{
48+
"cell_type": "code",
49+
"execution_count": null,
50+
"metadata": {},
51+
"outputs": [],
52+
"source": [
53+
"from huggingface_hub import login\n",
54+
"login(token=\"\")"
55+
]
56+
},
57+
{
58+
"cell_type": "markdown",
59+
"metadata": {},
60+
"source": [
61+
"## Get Transformers Logits"
62+
]
63+
},
64+
{
65+
"cell_type": "code",
66+
"execution_count": null,
67+
"metadata": {},
68+
"outputs": [],
69+
"source": [
70+
"import torch\n",
71+
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
72+
"\n",
73+
"def load_model(model_name=\"gpt2\"):\n",
74+
" tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
75+
" model = AutoModelForCausalLM.from_pretrained(model_name)\n",
76+
" return model, tokenizer\n",
77+
"\n",
78+
"def get_logits(model, tokenizer, sentence, device):\n",
79+
" # Tokenize the input sentence\n",
80+
" inputs = tokenizer(sentence, return_tensors=\"pt\")\n",
81+
" \n",
82+
" # Move inputs to the device\n",
83+
" inputs = {k: v.to(device) for k, v in inputs.items()}\n",
84+
" \n",
85+
" # Generate the logits\n",
86+
" with torch.no_grad():\n",
87+
" outputs = model(**inputs)\n",
88+
" \n",
89+
" # Get the logits for all tokens\n",
90+
" logits = outputs.logits\n",
91+
" \n",
92+
" return logits\n",
93+
"\n",
94+
"model, tokenizer = load_model(model_name)\n",
95+
"model = model.to(device)\n",
96+
"\n",
97+
"hf_logits = get_logits(model, tokenizer, sentence, device)[:, -1, :]"
98+
]
99+
},
100+
{
101+
"cell_type": "markdown",
102+
"metadata": {},
103+
"source": [
104+
"## Get TransformerLens Logits"
105+
]
106+
},
107+
{
108+
"cell_type": "code",
109+
"execution_count": null,
110+
"metadata": {},
111+
"outputs": [],
112+
"source": [
113+
"model = HookedTransformer.from_pretrained_no_processing(model_name, device=device)\n",
114+
"tokens = model.to_tokens(sentence, prepend_bos=False)\n",
115+
"tl_logits = model(tokens)[:, -1, :]"
116+
]
117+
},
118+
{
119+
"cell_type": "markdown",
120+
"metadata": {},
121+
"source": [
122+
"## Compare Logit Distributions\n",
123+
"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."
124+
]
125+
},
126+
{
127+
"cell_type": "markdown",
128+
"metadata": {},
129+
"source": [
130+
"### Shape"
131+
]
132+
},
133+
{
134+
"cell_type": "code",
135+
"execution_count": null,
136+
"metadata": {},
137+
"outputs": [],
138+
"source": [
139+
"print(f\"HF Logits Shape: {hf_logits.shape}\")\n",
140+
"print(f\"TL Logits Shape: {tl_logits.shape}\")"
141+
]
142+
},
143+
{
144+
"cell_type": "markdown",
145+
"metadata": {},
146+
"source": [
147+
"### Tensor Comparison"
148+
]
149+
},
150+
{
151+
"cell_type": "code",
152+
"execution_count": null,
153+
"metadata": {},
154+
"outputs": [],
155+
"source": [
156+
"are_close = torch.allclose(tl_logits, hf_logits, rtol=1e-5, atol=1e-3)\n",
157+
"print(f\"Are the logits close? {are_close}\")"
158+
]
159+
},
160+
{
161+
"cell_type": "markdown",
162+
"metadata": {},
163+
"source": [
164+
"### Mean Squared Error"
165+
]
166+
},
167+
{
168+
"cell_type": "code",
169+
"execution_count": null,
170+
"metadata": {},
171+
"outputs": [],
172+
"source": [
173+
"# Compare the logits with MSE\n",
174+
"mse = torch.nn.functional.mse_loss(hf_logits, tl_logits)\n",
175+
"print(f\"MSE: {mse}\")"
176+
]
177+
},
178+
{
179+
"cell_type": "markdown",
180+
"metadata": {},
181+
"source": [
182+
"### Maximum Absolute Difference"
183+
]
184+
},
185+
{
186+
"cell_type": "code",
187+
"execution_count": null,
188+
"metadata": {},
189+
"outputs": [],
190+
"source": [
191+
"max_diff = torch.max(torch.abs(tl_logits - hf_logits))\n",
192+
"print(f\"Max Diff: {max_diff}\")"
193+
]
194+
},
195+
{
196+
"cell_type": "markdown",
197+
"metadata": {},
198+
"source": [
199+
"### Cosine Similarity"
200+
]
201+
},
202+
{
203+
"cell_type": "code",
204+
"execution_count": null,
205+
"metadata": {},
206+
"outputs": [],
207+
"source": [
208+
"cosine_sim = F.cosine_similarity(tl_logits, hf_logits, dim=-1).mean()\n",
209+
"print(f\"Cosine Sim: {cosine_sim}\")"
210+
]
211+
},
212+
{
213+
"cell_type": "markdown",
214+
"metadata": {},
215+
"source": [
216+
"### KL Divergence"
217+
]
218+
},
219+
{
220+
"cell_type": "code",
221+
"execution_count": null,
222+
"metadata": {},
223+
"outputs": [],
224+
"source": [
225+
"def kl_div(logits1: torch.Tensor, logits2: torch.Tensor) -> torch.Tensor:\n",
226+
" probs1 = F.softmax(logits1, dim=-1)\n",
227+
" probs2 = F.softmax(logits2, dim=-1)\n",
228+
" return F.kl_div(probs1.log(), probs2, reduction='batchmean')\n",
229+
"\n",
230+
"kl_tl_hf = kl_div(tl_logits, hf_logits)\n",
231+
"kl_hf_tl = kl_div(hf_logits, tl_logits)\n",
232+
"print(f\"KL(TL||HF): {kl_tl_hf}\")\n",
233+
"print(f\"KL(HF||TL): {kl_hf_tl}\")"
234+
]
235+
},
236+
{
237+
"cell_type": "code",
238+
"execution_count": null,
239+
"metadata": {},
240+
"outputs": [],
241+
"source": []
242+
}
243+
],
244+
"metadata": {
245+
"kernelspec": {
246+
"display_name": "sae-l",
247+
"language": "python",
248+
"name": "python3"
249+
},
250+
"language_info": {
251+
"codemirror_mode": {
252+
"name": "ipython",
253+
"version": 3
254+
},
255+
"file_extension": ".py",
256+
"mimetype": "text/x-python",
257+
"name": "python",
258+
"nbconvert_exporter": "python",
259+
"pygments_lexer": "ipython3",
260+
"version": "3.12.4"
261+
}
262+
},
263+
"nbformat": 4,
264+
"nbformat_minor": 2
265+
}

‎transformer_lens/HookedTransformerConfig.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,18 @@ class HookedTransformerConfig:
181181
output_logits_soft_cap (float): An optional softcap for output logits, currently only used
182182
in Gemma-2 (see attn_scores_soft_cap for details). Defaults to -1.0, which means not
183183
set.
184+
use_NTK_by_parts_rope (bool): Whether to apply the "NTK-by-parts" method when using Rotary
185+
Positional Embedding. This method adjusts the interpolation based on frequency factors
186+
for different parts of the hidden dimensions. See Section 3.2 in
187+
https://arxiv.org/pdf/2309.00071 for details. Defaults to False.
188+
NTK_by_parts_low_freq_factor (float): The threshold applied to low-frequency hidden
189+
dimensions during interpolation when using the "NTK-by-parts" method. Defaults to 1.0.
190+
NTK_by_parts_high_freq_factor (float): The threshold applied to high-frequency hidden
191+
dimensions during interpolation in the "NTK-by-parts" method. Defaults to 4.0.
192+
NTK_by_parts_factor (float): The overall factor used in the "NTK-by-parts" method that
193+
affects the rate of change between low and high-frequency interpolation strategies.
194+
Defaults to 8.0.
195+
184196
185197
"""
186198

@@ -246,6 +258,10 @@ class HookedTransformerConfig:
246258
use_normalization_before_and_after: bool = False
247259
attn_scores_soft_cap: float = -1.0
248260
output_logits_soft_cap: float = -1.0
261+
use_NTK_by_parts_rope: bool = False
262+
NTK_by_parts_low_freq_factor: float = 1.0
263+
NTK_by_parts_high_freq_factor: float = 4.0
264+
NTK_by_parts_factor: float = 8.0
249265

250266
def __post_init__(self):
251267
if self.n_heads == -1:

‎transformer_lens/components/abstract_attention.py

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import math
12
from abc import ABC
23
from typing import Dict, Optional, Tuple, Union
34

@@ -478,8 +479,33 @@ def calculate_sin_cos_rotary(
478479
pos = torch.arange(n_ctx, dtype=high_precision)
479480
dim = torch.arange(rotary_dim // 2, dtype=high_precision)
480481

481-
# A set of frequencies evenly spaced in log space
482-
freq = base ** (dim / (rotary_dim / 2))
482+
# Llama-3.1 uses NTK-by-Parts Rotary Embedding introduced in Section 3.2 in https://arxiv.org/pdf/2309.00071
483+
# Implementation copied from https://github.com/huggingface/transformers/blob/v4.46.0/src/transformers/modeling_rope_utils.py#L310
484+
if self.cfg.use_NTK_by_parts_rope:
485+
inv_freq = 1.0 / (
486+
base ** (torch.arange(0, rotary_dim, 2, dtype=torch.int64).float() / rotary_dim)
487+
)
488+
factor = self.cfg.NTK_by_parts_factor
489+
low_freq_factor = self.cfg.NTK_by_parts_low_freq_factor
490+
high_freq_factor = self.cfg.NTK_by_parts_high_freq_factor
491+
old_context_len = n_ctx
492+
493+
low_freq_wavelen = old_context_len / low_freq_factor
494+
high_freq_wavelen = old_context_len / high_freq_factor
495+
496+
wavelen = 2 * math.pi / inv_freq
497+
inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
498+
smooth_factor = (old_context_len / wavelen - low_freq_factor) / (
499+
high_freq_factor - low_freq_factor
500+
)
501+
smoothed_inv_freq = (
502+
1 - smooth_factor
503+
) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
504+
is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
505+
inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
506+
freq = 1 / inv_freq_llama
507+
else:
508+
freq = base ** (dim / (rotary_dim / 2))
483509
if self.cfg.rotary_adjacent_pairs:
484510
freq = einops.repeat(freq, "d -> (d 2)")
485511
else:

0 commit comments

Comments
 (0)