Coverage for transformer_lens/HookedTransformer.py: 77%
832 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Hooked Transformer.
3The Hooked Transformer is the core part of TransformerLens.
5In common PyTorch model implementations (e.g. ones from HuggingFace) it's fairly easy to extract
6model weights, but much harder to extract activations. TransformerLens aims to simplify this task by
7attaching hooks to every notable activation within the model. This enables the inspection and/or
8alteration of activations in individual components like attention heads and MLP layers, facilitating
9a deeper understanding of the internal workings of transformers like GPT-2.
10"""
12from __future__ import annotations
14import logging
15import os
16import warnings
17from collections.abc import Generator
18from typing import (
19 Any,
20 Dict,
21 List,
22 NamedTuple,
23 Optional,
24 Tuple,
25 Type,
26 TypeVar,
27 Union,
28 cast,
29 overload,
30)
32import einops
33import numpy as np
34import torch
35import torch.nn as nn
36import tqdm.auto as tqdm
37from jaxtyping import Float, Int
38from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
39from transformers.models.auto.tokenization_auto import AutoTokenizer
40from transformers.tokenization_utils_base import PreTrainedTokenizerBase
41from typing_extensions import Literal
43import transformer_lens.loading_from_pretrained as loading
44import transformer_lens.utilities as utils
45from transformer_lens.ActivationCache import ActivationCache
47# Activation cache for run_with_cache; KV cache for generation
48from transformer_lens.cache.key_value_cache import TransformerLensKeyValueCache
49from transformer_lens.components import (
50 Embed,
51 LayerNorm,
52 LayerNormPre,
53 PosEmbed,
54 RMSNorm,
55 RMSNormPre,
56 TransformerBlock,
57 Unembed,
58)
59from transformer_lens.components.mlps.gated_mlp import GatedMLP
60from transformer_lens.components.mlps.mlp import MLP
61from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig
62from transformer_lens.FactoredMatrix import FactoredMatrix
63from transformer_lens.hook_points import HookPoint
64from transformer_lens.HookedRootModule import HookedRootModule
65from transformer_lens.loading_from_pretrained import NON_HF_HOSTED_MODEL_NAMES
66from transformer_lens.utilities import (
67 USE_DEFAULT_VALUE,
68 TypedModuleList,
69 apply_softcap,
70 get_best_available_device,
71 get_device_for_block_index,
72 init_kaiming_normal_,
73 init_kaiming_uniform_,
74 init_xavier_normal_,
75 init_xavier_uniform_,
76 softcap_enabled,
77)
78from transformer_lens.utilities.architectures import POST_NORM_ARCHITECTURES
79from transformer_lens.utilities.devices import move_to_and_update_config
80from transformer_lens.weight_processing import ProcessWeights
82SingleLoss = Float[torch.Tensor, ""] # Type alias for a single element tensor
83LossPerToken = Float[torch.Tensor, "batch pos-1"]
84Loss = Union[SingleLoss, LossPerToken]
86DTYPE_FROM_STRING = {
87 "float32": torch.float32,
88 "fp32": torch.float32,
89 "float16": torch.float16,
90 "fp16": torch.float16,
91 "bfloat16": torch.bfloat16,
92 "bf16": torch.bfloat16,
93}
95T = TypeVar("T", bound="HookedTransformer")
98class Output(NamedTuple):
99 """Output Named Tuple.
101 Named tuple object for if we want to output both logits and loss.
102 """
104 logits: Float[torch.Tensor, "batch pos d_vocab"]
105 loss: Loss
108class HookedTransformer(HookedRootModule):
109 """Hooked Transformer.
111 Implements a full Transformer using the components :doc:`here <transformer_lens.components>`,
112 with a :class:`transformer_lens.hook_points.HookPoint` on every interesting activation.
114 TransformerLens comes loaded with >50 GPT-style models. Typically you initialise it with one of
115 these via :meth:`from_pretrained`, although it can also be instantiated with randomly
116 initialized weights via :meth:`__init__`.
118 Once you've initialized the model, a common next step is to test it can do the task you're
119 investigating. This can be done with :func:`transformer_lens.utils.test_prompt`.
121 Tokenization notes
122 ------------------
124 :meth:`to_tokens`, :meth:`to_str_tokens`, :meth:`get_token_position`,
125 :meth:`forward` (string input), and :meth:`generate` accept ``prepend_bos``
126 to control BOS prepending. Resolution: explicit arg →
127 ``cfg.default_prepend_bos`` (defaults ``True``, even for non-BOS-trained
128 models — attention heads tend to use position 0 as a resting state).
129 **Pass ``prepend_bos=False`` when tokenizing a fragment of a larger
130 prompt** — off-by-one position errors usually trace back here.
132 Reconciliation with ``cfg.tokenizer_prepends_bos`` (set by
133 :meth:`set_tokenizer` for tokenizers that add BOS automatically) is
134 handled internally — pass the value you want; the framework adds or
135 strips manually as needed.
137 BPE/SentencePiece tokenizers treat ``"hello"``, ``" hello"``, and
138 ``"Hello"`` as distinct tokens. Concatenated prompts may not tokenize
139 as the sum of parts — inspect with :meth:`to_str_tokens` when in doubt.
140 """
142 ln_final: nn.Module
143 tokenizer: Optional[PreTrainedTokenizerBase]
144 blocks: TypedModuleList[TransformerBlock]
146 def __init__(
147 self,
148 cfg: Union[HookedTransformerConfig, Dict],
149 tokenizer: Optional[PreTrainedTokenizerBase] = None,
150 move_to_device: bool = True,
151 default_padding_side: Optional[Literal["left", "right"]] = None,
152 ):
153 """Model initialization.
155 Note that if you want to load the model from pretrained weights, you should use
156 :meth:`from_pretrained` instead.
158 Args:
159 cfg: The config to use for the model.
160 tokenizer: The tokenizer to use for the model. If not provided, it is inferred from
161 `cfg.tokenizer_name` or initialized to `None`. If `None`, then the model cannot be
162 passed strings, and d_vocab must be explicitly set.
163 move_to_device: Whether to move the model to the device specified in cfg.
164 device. Must be true if `n_devices` in the config is greater than 1, since the
165 model's layers will be split across multiple devices.
166 default_padding_side: Which side to pad on.
167 """
168 super().__init__()
169 warnings.warn(
170 "HookedTransformer is deprecated and will be removed in 4.0. Use "
171 "TransformerBridge.boot_transformers(...) instead, then call "
172 "enable_compatibility_mode() for HookedTransformer-equivalent numerics.",
173 DeprecationWarning,
174 stacklevel=2,
175 )
176 if isinstance(cfg, str): 176 ↛ 177line 176 didn't jump to line 177 because the condition on line 176 was never true
177 raise ValueError(
178 "Please pass in a config dictionary or HookedTransformerConfig object. If you want to load a "
179 "pretrained model, use HookedTransformer.from_pretrained() instead."
180 )
182 self.cfg = HookedTransformerConfig.unwrap(cfg)
183 if tokenizer is not None:
184 self.set_tokenizer(tokenizer, default_padding_side=default_padding_side)
185 elif self.cfg.tokenizer_name is not None:
186 # If we have a tokenizer name, we can load it from HuggingFace
187 if self.cfg.tokenizer_name in NON_HF_HOSTED_MODEL_NAMES: 187 ↛ 188line 187 didn't jump to line 188 because the condition on line 187 was never true
188 logging.warning(
189 "%s tokenizer not loaded. Please load manually.",
190 self.cfg.tokenizer_name,
191 )
192 else:
193 # Hugging Face defaults to use_fast to True
194 use_fast = True
195 # Phi model's fast tokenizer does not support adding a BOS token, use_fast
196 # should be False
197 if "phi" in self.cfg.tokenizer_name.lower(): 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 use_fast = False
199 huggingface_token = os.environ.get("HF_TOKEN", "")
200 add_bos_token = self.cfg.original_architecture not in [
201 "OlmoForCausalLM",
202 "OlmoeForCausalLM",
203 "Olmo2ForCausalLM",
204 "Qwen3ForCausalLM",
205 "PhiForCausalLM",
206 ]
207 self.set_tokenizer(
208 AutoTokenizer.from_pretrained(
209 self.cfg.tokenizer_name,
210 add_bos_token=add_bos_token,
211 trust_remote_code=self.cfg.trust_remote_code,
212 use_fast=use_fast,
213 token=huggingface_token if len(huggingface_token) > 0 else None,
214 ),
215 default_padding_side=default_padding_side,
216 )
217 else:
218 # If no tokenizer name is provided, we assume we're training on an algorithmic task and
219 # will pass in tokens directly. In this case, we don't need a tokenizer.
220 assert self.cfg.d_vocab != -1, "Must provide a tokenizer if d_vocab is not provided"
221 self.tokenizer = None
222 if default_padding_side != None: 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 logging.warning(
224 "default_padding_side is explicitly given but ignored because tokenizer is not set."
225 )
227 self.embed = Embed(self.cfg)
228 self.hook_embed = HookPoint() # [batch, pos, d_model]
230 if self.cfg.positional_embedding_type != "rotary":
231 self.pos_embed = PosEmbed(self.cfg)
232 self.hook_pos_embed = HookPoint() # [batch, pos, d__dictmodel]
234 if self.cfg.use_hook_tokens:
235 self.hook_tokens = HookPoint() # [batch, pos]
237 self.blocks = TypedModuleList(
238 [TransformerBlock(self.cfg, block_index) for block_index in range(self.cfg.n_layers)]
239 )
241 if self.cfg.normalization_type == "RMS":
242 self.ln_final = RMSNorm(self.cfg)
243 elif self.cfg.normalization_type == "RMSPre": 243 ↛ 244line 243 didn't jump to line 244 because the condition on line 243 was never true
244 self.ln_final = RMSNormPre(self.cfg)
245 elif self.cfg.normalization_type == "LN":
246 if self.cfg.final_rms: 246 ↛ 247line 246 didn't jump to line 247 because the condition on line 246 was never true
247 self.ln_final = RMSNorm(self.cfg)
248 else:
249 self.ln_final = LayerNorm(self.cfg)
250 elif self.cfg.normalization_type == "LNPre":
251 # We've folded in LayerNorm weights, so just need the center + scale parts
252 if self.cfg.final_rms:
253 self.ln_final = RMSNormPre(self.cfg)
254 else:
255 self.ln_final = LayerNormPre(self.cfg)
256 elif self.cfg.normalization_type is None: 256 ↛ 260line 256 didn't jump to line 260 because the condition on line 256 was always true
257 # If it's None, don't create either layer
258 pass
259 else:
260 logging.warning("Invalid normalization_type passed in %s", self.cfg.normalization_type)
261 self.unembed = Unembed(self.cfg)
263 if self.cfg.init_weights:
264 self.init_weights()
266 if move_to_device:
267 # We load the devices in a pipeline manner - the first device gets the embed and
268 # pos_embed layers and the first n_layers // n_devices blocks, the second gets the next
269 # n_layers // n_devices blocks ... the last gets the last n_layers // n_devices blocks,
270 # the final normalization layer (if it exists) and the unembed layer
271 self.move_model_modules_to_device()
273 # Helper variable to store a small (10K-20K) dataset of training data. Empty by default, can
274 # be loaded with load_sample_training_dataset
275 self.dataset = None
277 # Gives each module a parameter with its name (relative to this root module)
278 # Needed for HookPoints to work
279 self.setup()
281 def check_hooks_to_add(
282 self,
283 hook_point,
284 hook_point_name,
285 hook,
286 dir="fwd",
287 is_permanent=False,
288 prepend=False,
289 ) -> None:
290 if hook_point_name.endswith("attn.hook_result"):
291 assert (
292 self.cfg.use_attn_result
293 ), f"Cannot add hook {hook_point_name} if use_attn_result_hook is False"
294 if hook_point_name.endswith(("hook_q_input", "hook_k_input", "hook_v_input")):
295 assert (
296 self.cfg.use_split_qkv_input
297 ), f"Cannot add hook {hook_point_name} if use_split_qkv_input is False"
298 if hook_point_name.endswith("mlp_in"):
299 assert (
300 self.cfg.use_hook_mlp_in
301 ), f"Cannot add hook {hook_point_name} if use_hook_mlp_in is False"
302 if hook_point_name.endswith("attn_in"):
303 assert (
304 self.cfg.use_attn_in
305 ), f"Cannot add hook {hook_point_name} if use_attn_in is False"
307 def get_pos_offset(self, past_kv_cache, batch_size):
308 # If we're doing caching, then we reuse keys and values from previous runs, as that's the
309 # only way that past activations will affect the final logits. The cache contains those so
310 # we don't need to recompute them. This is useful for generating text. As we have absolute
311 # positional encodings, to implement this we have a `pos_offset` variable, defaulting to
312 # zero, which says to offset which positional encodings are used (cached keys and values
313 # were calculated with their own positional encodings).
314 if past_kv_cache is None:
315 pos_offset = 0
316 else:
317 (
318 cached_batch_size,
319 cache_ctx_length,
320 num_heads_in_cache,
321 d_head_in_cache,
322 ) = past_kv_cache[0].past_keys.shape
323 assert cached_batch_size == batch_size
324 if self.cfg.n_key_value_heads is None: 324 ↛ 327line 324 didn't jump to line 327 because the condition on line 324 was always true
325 assert num_heads_in_cache == self.cfg.n_heads
326 else:
327 assert num_heads_in_cache == self.cfg.n_key_value_heads
328 assert d_head_in_cache == self.cfg.d_head
329 pos_offset = cache_ctx_length
330 return pos_offset
332 def get_residual(
333 self,
334 embed,
335 pos_offset,
336 prepend_bos=USE_DEFAULT_VALUE,
337 attention_mask=None,
338 tokens=None,
339 return_shortformer_pos_embed=True,
340 device=None,
341 ):
342 if device is None:
343 device = get_device_for_block_index(0, self.cfg)
345 if tokens is None:
346 # Because tokens only need for defining batch size and sequence length, we can simply synthesize them
347 tokens = torch.ones((embed.size(0), embed.size(1))).int().to(device)
349 if self.cfg.positional_embedding_type == "standard":
350 pos_embed = self.hook_pos_embed(
351 self.pos_embed(tokens, pos_offset, attention_mask)
352 ) # [batch, pos, d_model]
353 residual = embed + pos_embed # [batch, pos, d_model]
354 shortformer_pos_embed = None
355 elif self.cfg.positional_embedding_type == "shortformer":
356 # If we're using shortformer style attention, we don't add the positional embedding to
357 # the residual stream. See HookedTransformerConfig for details
358 pos_embed = self.hook_pos_embed(
359 self.pos_embed(tokens, pos_offset, attention_mask)
360 ) # [batch, pos, d_model]
361 residual = embed
362 shortformer_pos_embed = pos_embed
363 elif self.cfg.positional_embedding_type == "rotary":
364 # Rotary doesn't use positional embeddings, instead they're applied when dot producting
365 # keys and queries. See HookedTransformerConfig for details
366 residual = embed
367 shortformer_pos_embed = None
368 elif self.cfg.positional_embedding_type == "alibi": 368 ↛ 373line 368 didn't jump to line 373 because the condition on line 368 was always true
369 # ALiBi does not add positional embeddings to word embeddings,instead it biases QK attention scores.
370 residual = embed
371 shortformer_pos_embed = None
372 else:
373 raise ValueError(
374 f"Invalid positional_embedding_type passed in {self.cfg.positional_embedding_type}"
375 )
377 if return_shortformer_pos_embed: 377 ↛ 380line 377 didn't jump to line 380 because the condition on line 377 was always true
378 return residual, shortformer_pos_embed
379 else:
380 return residual
382 def input_to_embed(
383 self,
384 input: Union[str, List[str], Int[torch.Tensor, "batch pos"]],
385 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
386 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
387 attention_mask: Optional[torch.Tensor] = None,
388 past_kv_cache: Optional[TransformerLensKeyValueCache] = None,
389 ) -> Tuple[
390 Float[torch.Tensor, "batch pos d_model"], # residual
391 Optional[Int[torch.Tensor, "batch pos"]], # tokens
392 Optional[Float[torch.Tensor, "batch pos d_model"]], # shortformer_pos_embed
393 Optional[torch.Tensor], # attention_mask [batch pos]
394 ]:
395 """Convert input to first residual stream.
397 Args:
398 input (Union[str, List[str], Int[torch.Tensor, "batch pos"]]): The input to the model.
399 prepend_bos (bool, optional): Overrides self.cfg.default_prepend_bos. Whether to prepend
400 the BOS token to the input (only applies when input is a string). Defaults to None,
401 implying usage of self.cfg.default_prepend_bos which is set to True unless specified
402 otherwise. Pass True or False to locally override the default.
403 padding_side ([Literal["left", "right"], optional): Overrides
404 self.tokenizer.padding_side. Specifies which side to pad when tokenizing
405 multiple strings of different lengths.
406 past_kv_cache (TransformerLensKeyValueCache, optional): If passed, we're doing caching
407 and attention_mask will be stored in the cache.
408 """
409 if isinstance(input, str) or isinstance(input, list):
410 # If text, convert to tokens (batch_size=1)
411 assert (
412 self.tokenizer is not None
413 ), "Must provide a tokenizer if passing a string to the model"
414 # This is only intended to support passing in a single string
415 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
416 else:
417 tokens = input
418 if len(tokens.shape) == 1: 418 ↛ 420line 418 didn't jump to line 420 because the condition on line 418 was never true
419 # If tokens are a rank 1 tensor, add a dummy batch dimension to avoid things breaking.
420 tokens = tokens[None]
421 if tokens.device.type != self.cfg.device: 421 ↛ 422line 421 didn't jump to line 422 because the condition on line 421 was never true
422 tokens = tokens.to(get_device_for_block_index(0, self.cfg))
424 if (
425 (self.tokenizer and self.tokenizer.padding_side == "left")
426 or attention_mask is not None
427 or past_kv_cache is not None
428 ):
429 # This means we need to have an explicit attention mask.
430 if attention_mask is None:
431 # If the padding side is left or we are using caching, we need to compute the attention
432 # mask for the adjustment of absolute positional embeddings and attention masking so
433 # that pad tokens are not attended.
434 if prepend_bos is USE_DEFAULT_VALUE:
435 prepend_bos = self.cfg.default_prepend_bos
436 if self.tokenizer is None: 436 ↛ 437line 436 didn't jump to line 437 because the condition on line 436 was never true
437 raise ValueError("Cannot compute attention mask without a tokenizer.")
438 attention_mask = utils.get_attention_mask(self.tokenizer, tokens, prepend_bos)
440 assert attention_mask.shape == tokens.shape, (
441 f"Attention mask shape {attention_mask.shape} does not match tokens shape "
442 f"{tokens.shape}"
443 )
444 attention_mask = attention_mask.to(get_device_for_block_index(0, self.cfg))
445 if past_kv_cache is not None:
446 # past_kv_cache is not None, so we're doing caching.
447 # We need to extend the previous attention_mask.
448 # Update the past_kv_cache with the new attention_mask (unless it's frozen)
449 attention_mask = past_kv_cache.append_attention_mask(attention_mask)
450 else:
451 # We separate this case from for computational efficiency.
452 attention_mask = None
454 batch_size = tokens.shape[0]
455 pos_offset = self.get_pos_offset(past_kv_cache, batch_size)
457 if self.cfg.use_hook_tokens:
458 tokens = self.hook_tokens(tokens)
460 embed = self.hook_embed(self.embed(tokens)) # [batch, pos, d_model]
461 residual, shortformer_pos_embed = self.get_residual(
462 embed,
463 pos_offset,
464 prepend_bos,
465 attention_mask,
466 tokens,
467 return_shortformer_pos_embed=True,
468 )
469 return residual, tokens, shortformer_pos_embed, attention_mask
471 @overload
472 def forward(
473 self,
474 input,
475 return_type: Literal["logits"],
476 loss_per_token: bool = False,
477 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
478 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
479 start_at_layer: Optional[int] = None,
480 tokens: Optional[Int[torch.Tensor, "batch pos"]] = None,
481 shortformer_pos_embed: Optional[Float[torch.Tensor, "batch pos d_model"]] = None,
482 attention_mask: Optional[torch.Tensor] = None, # [batch pos]
483 stop_at_layer: Optional[int] = None,
484 past_kv_cache: Optional[TransformerLensKeyValueCache] = None,
485 ) -> Loss:
486 ...
488 @overload
489 def forward(
490 self,
491 input,
492 return_type: Literal["loss"],
493 loss_per_token: bool = False,
494 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
495 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
496 start_at_layer: Optional[int] = None,
497 tokens: Optional[Int[torch.Tensor, "batch pos"]] = None,
498 shortformer_pos_embed: Optional[Float[torch.Tensor, "batch pos d_model"]] = None,
499 attention_mask: Optional[torch.Tensor] = None, # [batch pos]
500 stop_at_layer: Optional[int] = None,
501 past_kv_cache: Optional[TransformerLensKeyValueCache] = None,
502 ) -> Loss:
503 ...
505 @overload
506 def forward(
507 self,
508 input,
509 return_type: Literal["both"],
510 loss_per_token: bool = False,
511 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
512 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
513 start_at_layer: Optional[int] = None,
514 tokens: Optional[Int[torch.Tensor, "batch pos"]] = None,
515 shortformer_pos_embed: Optional[Float[torch.Tensor, "batch pos d_model"]] = None,
516 attention_mask: Optional[torch.Tensor] = None, # [batch pos]
517 stop_at_layer: Optional[int] = None,
518 past_kv_cache: Optional[TransformerLensKeyValueCache] = None,
519 ) -> Tuple[Float[torch.Tensor, "batch pos d_vocab"], Loss]:
520 ...
522 @overload
523 def forward(
524 self,
525 input,
526 return_type: Literal[None],
527 loss_per_token: bool = False,
528 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
529 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
530 start_at_layer: Optional[int] = None,
531 tokens: Optional[Int[torch.Tensor, "batch pos"]] = None,
532 shortformer_pos_embed: Optional[Float[torch.Tensor, "batch pos d_model"]] = None,
533 attention_mask: Optional[torch.Tensor] = None, # [batch pos]
534 stop_at_layer: Optional[int] = None,
535 past_kv_cache: Optional[TransformerLensKeyValueCache] = None,
536 ) -> None:
537 ...
539 def forward(
540 self,
541 input: Union[
542 str,
543 List[str],
544 Int[torch.Tensor, "batch pos"],
545 Float[torch.Tensor, "batch pos d_model"],
546 ],
547 return_type: Optional[str] = "logits",
548 loss_per_token: bool = False,
549 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
550 padding_side: Optional[Literal["left", "right"]] = USE_DEFAULT_VALUE,
551 start_at_layer: Optional[int] = None,
552 tokens: Optional[Int[torch.Tensor, "batch pos"]] = None,
553 shortformer_pos_embed: Optional[Float[torch.Tensor, "batch pos d_model"]] = None,
554 attention_mask: Optional[torch.Tensor] = None, # [batch pos]
555 stop_at_layer: Optional[int] = None,
556 past_kv_cache: Optional[TransformerLensKeyValueCache] = None,
557 ) -> Union[
558 None,
559 Float[torch.Tensor, "batch pos d_vocab"],
560 Loss,
561 Tuple[Float[torch.Tensor, "batch pos d_vocab"], Loss],
562 ]:
563 """Forward Pass.
565 Input is either a batch of tokens ([batch, pos]) or a text string, a string is automatically
566 tokenized to a batch of a single element. The prepend_bos flag only applies when inputting a
567 text string.
569 Note that loss is the standard "predict the next token" cross-entropy loss for GPT-2 style
570 language models - if you want a custom loss function, the recommended behaviour is returning
571 the logits and then applying your custom loss function.
573 Args:
574 return_type Optional[str]: The type of output to return. Can be one of: None (return
575 nothing, don't calculate logits), 'logits' (return logits), 'loss' (return
576 cross-entropy loss), 'both' (return logits and loss).
577 loss_per_token bool: Whether to return the (next token prediction) loss per token (True)
578 or average (False). Average loss is a scalar (averaged over position *and* batch),
579 per-token loss is a tensor ([batch, position-1]) - position-1 because we're
580 predicting the next token, and there's no specified next token for the final token.
581 Defaults to False.
582 prepend_bos Optional[bool]: Overrides self.cfg.default_prepend_bos. Whether to prepend
583 the BOS token to the input (only applies when input is a string). Defaults to None,
584 implying usage of self.cfg.default_prepend_bos which is set to True unless specified
585 otherwise. (Even for models not explicitly trained with a prepended BOS token, heads
586 often use the first position as a resting position and accordingly lose information
587 from the first token, so this empirically seems to give better results.) Pass True
588 or False to locally override the default.
589 padding_side Optional[Literal["left", "right"]]: Overrides self.tokenizer.padding_side.
590 Specifies which side to pad on when tokenizing multiple strings of different
591 lengths.
592 start_at_layer Optional[int]: If not None, start the forward pass at the specified
593 layer. Requires input to be the residual stream before the specified layer with
594 shape [batch, pos, d_model]. Inclusive - ie, start_at_layer = 0 skips the embedding
595 then runs the rest of the model. Supports negative indexing. start_at_layer = -1
596 only runs the final block and the unembedding. Defaults to None (run the full
597 model).
598 tokens: Optional[Int[torch.Tensor, "batch pos"]]: Tokenized input. Only use if
599 start_at_layer is not None and return type is "loss" or "both".
600 shortformer_pos_embed: Optional[Float[torch.Tensor, "batch pos d_model"]]: Positional
601 embedding for shortformer models. Only use if start_at_layer is not None and
602 self.cfg.positional_embedding_type == "shortformer".
603 attention_mask: Optional[torch.Tensor]: Override the attention mask used to ignore
604 padded tokens. If start_at_layer is not None and (self.tokenizer.padding_side ==
605 "left" or past_kv_cache is not None), this should be passed as the attention mask
606 is not computed automatically. Defaults to None.
607 stop_at_layer Optional[int]: If not None, stop the forward pass at the specified layer.
608 Exclusive - ie, stop_at_layer = 0 will only run the embedding layer, stop_at_layer =
609 1 will run the embedding layer and the first transformer block, etc. Supports
610 negative indexing. Useful for analysis of intermediate layers, eg finding neuron
611 activations in layer 3 of a 24 layer model. Defaults to None (run the full model).
612 If not None, we return the last residual stream computed.
613 past_kv_cache Optional[TransformerLensKeyValueCache]: If not None, keys and values
614 will be stored for every attention head (unless the cache is frozen). If there are
615 keys and values already in the cache, these will be prepended to the keys and values
616 for the new input, so that the new tokens can pay attention to previous tokens. This
617 is useful for generating text, because we don't need to repeat computation for
618 tokens that have already been through the model. Also caches attention_mask so
619 previous tokens are masked correctly (unless frozen). Padding should be ignored in
620 all cases, so it's okay to eg. pass in left padded tokens twice in a row.
621 Warning: Don't accidentally prepend_bos to the second half of a prompt.
622 Defaults to None (don't use caching).
623 """
625 with utils.LocallyOverridenDefaults(
626 self, prepend_bos=prepend_bos, padding_side=padding_side
627 ):
628 if start_at_layer is None:
629 (
630 residual,
631 tokens,
632 shortformer_pos_embed,
633 attention_mask,
634 ) = self.input_to_embed(
635 input,
636 prepend_bos=prepend_bos,
637 padding_side=padding_side,
638 attention_mask=attention_mask,
639 past_kv_cache=past_kv_cache,
640 )
641 else:
642 assert type(input) == torch.Tensor
643 residual = input
645 if start_at_layer is None:
646 start_at_layer = 0
647 # If we explicitly want to start or stop at a layer, we only iterate through the blocks
648 # between those indices. Note that start_at_layer is inclusive and stop_at_layer is
649 # exclusive.
650 # Eg: start_at_layer==None + stop_at_layer==0 means to only run the embed.
651 # Eg: start_at_layer==3 + stop_at_layer==-1 means to run from layer 3 until the end of the PENULTIMATE layer
652 blocks_and_idxs = list(zip(range(self.cfg.n_layers), self.blocks))
653 for i, block in blocks_and_idxs[start_at_layer:stop_at_layer]:
654 # Note that each block includes skip connections, so we don't need
655 # residual + block(residual)
656 # If we're using multiple GPUs, we need to send the residual and shortformer_pos_embed to the correct GPU
657 residual = residual.to(get_device_for_block_index(i, self.cfg))
658 if shortformer_pos_embed is not None:
659 shortformer_pos_embed = shortformer_pos_embed.to(
660 get_device_for_block_index(i, self.cfg)
661 )
663 residual = block(
664 residual,
665 # Cache contains a list of TransformerLensKeyValueCache objects, one for each
666 # block
667 past_kv_cache_entry=past_kv_cache[i] if past_kv_cache is not None else None,
668 shortformer_pos_embed=shortformer_pos_embed,
669 attention_mask=attention_mask,
670 ) # [batch, pos, d_model]
672 if stop_at_layer is not None:
673 # When we stop at an early layer, we end here rather than doing further computation
674 return residual
676 if self.cfg.normalization_type is not None:
677 residual = self.ln_final(residual) # [batch, pos, d_model]
678 if return_type is None:
679 return None
680 else:
681 logits = self.unembed(residual) # [batch, pos, d_vocab]
682 logits = apply_softcap(logits, self.cfg.output_logits_soft_cap)
683 if return_type == "logits":
684 return logits
685 else:
686 assert (
687 tokens is not None
688 ), "tokens must be passed in if return_type is 'loss' or 'both'"
689 loss = self.loss_fn(logits, tokens, attention_mask, per_token=loss_per_token)
690 if return_type == "loss": 690 ↛ 692line 690 didn't jump to line 692 because the condition on line 690 was always true
691 return loss
692 elif return_type == "both":
693 return Output(logits, loss)
694 else:
695 logging.warning(f"Invalid return_type passed in: {return_type}")
696 return None
698 def loss_fn(
699 self,
700 logits: Float[torch.Tensor, "batch pos d_vocab"],
701 tokens: Int[torch.Tensor, "batch pos"],
702 attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None,
703 per_token: bool = False,
704 ):
705 """Wrapper around `utils.lm_cross_entropy_loss`.
707 Used in forward() with return_type=="loss" or "both".
708 """
709 if tokens.device != logits.device: 709 ↛ 710line 709 didn't jump to line 710 because the condition on line 709 was never true
710 tokens = tokens.to(logits.device)
711 return utils.lm_cross_entropy_loss(logits, tokens, attention_mask, per_token)
713 @overload
714 def run_with_cache(
715 self, *model_args, return_cache_object: Literal[True] = True, **kwargs
716 ) -> Tuple[Output, ActivationCache]:
717 ...
719 @overload
720 def run_with_cache(
721 self, *model_args, return_cache_object: Literal[False], **kwargs
722 ) -> Tuple[Output, Dict[str, torch.Tensor]]:
723 ...
725 def run_with_cache(
726 self, *model_args, return_cache_object=True, remove_batch_dim=False, **kwargs
727 ) -> Tuple[
728 Union[
729 None,
730 Float[torch.Tensor, "batch pos d_vocab"],
731 Loss,
732 Tuple[Float[torch.Tensor, "batch pos d_vocab"], Loss],
733 ],
734 Union[ActivationCache, Dict[str, torch.Tensor]],
735 ]:
736 """Wrapper around `run_with_cache` in HookedRootModule.
738 If return_cache_object is True, this will return an ActivationCache object, with a bunch of
739 useful HookedTransformer specific methods, otherwise it will return a dictionary of
740 activations as in HookedRootModule.
741 """
742 out, cache_dict = super().run_with_cache(
743 *model_args, remove_batch_dim=remove_batch_dim, **kwargs
744 )
745 if return_cache_object: 745 ↛ 749line 745 didn't jump to line 749 because the condition on line 745 was always true
746 cache = ActivationCache(cache_dict, self, has_batch_dim=not remove_batch_dim)
747 return out, cache
748 else:
749 return out, cache_dict
751 def set_tokenizer(
752 self,
753 tokenizer,
754 default_padding_side=None,
755 ):
756 """Set the tokenizer to use for this model.
758 Args:
759 tokenizer (PreTrainedTokenizer): a pretrained HuggingFace tokenizer.
760 default_padding_side (str): "right" or "left", which side to pad on.
762 """
763 assert isinstance(
764 tokenizer, PreTrainedTokenizerBase
765 ), f"{type(tokenizer)} is not a supported tokenizer, please use PreTrainedTokenizer or PreTrainedTokenizerFast"
767 assert default_padding_side in [
768 "right",
769 "left",
770 None,
771 ], f"padding_side must be 'right', 'left' or 'None', got {default_padding_side}"
773 # Use a tokenizer that is initialized with add_bos_token=True as the default tokenizer.
774 # Such a tokenizer should be set as the default tokenizer because the tokenization of some
775 # tokenizers like LlamaTokenizer are different when bos token is automatically/manually
776 # prepended, and add_bos_token cannot be dynamically controlled after initialization
777 # (https://github.com/huggingface/transformers/issues/25886).
778 tokenizer_with_bos = tokenizer
779 if self.cfg.original_architecture not in [ 779 ↛ 786line 779 didn't jump to line 786 because the condition on line 779 was always true
780 "OlmoForCausalLM",
781 "OlmoeForCausalLM",
782 "Olmo2ForCausalLM",
783 ]:
784 tokenizer_with_bos = utils.get_tokenizer_with_bos(tokenizer)
786 self.tokenizer = tokenizer_with_bos
787 assert self.tokenizer is not None # keep mypy happy
789 # Use explicit value, else tokenizer default, else "right"
790 if default_padding_side is not None:
791 self.tokenizer.padding_side = default_padding_side
792 if self.tokenizer.padding_side is None: 792 ↛ 793line 792 didn't jump to line 793 because the condition on line 792 was never true
793 self.tokenizer.padding_side = "right"
795 # Detect whether tokenizer actually prepends BOS to control prepend_bos dynamically
796 self.cfg.tokenizer_prepends_bos = len(self.tokenizer.encode("")) > 0
798 if self.tokenizer.eos_token is None: 798 ↛ 799line 798 didn't jump to line 799 because the condition on line 798 was never true
799 self.tokenizer.eos_token = "<|endoftext|>"
800 if self.tokenizer.pad_token is None:
801 self.tokenizer.pad_token = self.tokenizer.eos_token
802 if self.tokenizer.bos_token is None: 802 ↛ 803line 802 didn't jump to line 803 because the condition on line 802 was never true
803 self.tokenizer.bos_token = self.tokenizer.eos_token
805 # Infer vocab size from tokenizer
806 if self.cfg.d_vocab == -1:
807 self.cfg.d_vocab = max(self.tokenizer.vocab.values()) + 1
808 if self.cfg.d_vocab_out == -1:
809 self.cfg.d_vocab_out = self.cfg.d_vocab
811 def to_tokens(
812 self,
813 input: Union[str, List[str]],
814 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
815 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
816 move_to_device: bool = True,
817 truncate: bool = True,
818 ) -> Int[torch.Tensor, "batch pos"]:
819 """Converts a string to a tensor of tokens.
821 See the class-level "Tokenization notes" for full ``prepend_bos``
822 semantics, the ``default_prepend_bos`` /
823 ``tokenizer_prepends_bos`` interaction, and the whitespace-
824 sensitivity gotcha. **Pass ``prepend_bos=False`` whenever you're
825 tokenizing only part of a prompt.**
827 Args:
828 input (Union[str, List[str]]): The input to tokenize.
829 prepend_bos (bool, optional): Overrides ``self.cfg.default_prepend_bos``.
830 Defaults to ``USE_DEFAULT_VALUE`` (use the cfg setting). Pass ``True``
831 or ``False`` to override locally.
832 padding_side (Union[Literal["left", "right"], None], optional): Overrides
833 self.tokenizer.padding_side. Specifies which side to pad when tokenizing
834 multiple strings of different lengths.
835 move_to_device (bool): Whether to move the output tensor of tokens to the device the
836 model lives on. Defaults to True
837 truncate (bool): If the output tokens are too long,
838 whether to truncate the output tokens to the model's max context window. Does nothing
839 for shorter inputs. Defaults to True.
840 """
841 with utils.LocallyOverridenDefaults(
842 self, prepend_bos=prepend_bos, padding_side=padding_side
843 ):
844 assert self.tokenizer is not None, "Cannot use to_tokens without a tokenizer"
845 assert (
846 self.cfg.tokenizer_prepends_bos is not None
847 ), "Set the tokenizer for the model by calling set_tokenizer"
849 if self.cfg.default_prepend_bos and not self.cfg.tokenizer_prepends_bos: 849 ↛ 851line 849 didn't jump to line 851 because the condition on line 849 was never true
850 # We want to prepend bos but the tokenizer doesn't automatically do it, so we add it manually
851 input = utils.get_input_with_manually_prepended_bos(self.tokenizer.bos_token, input)
853 tokens = self.tokenizer(
854 input,
855 return_tensors="pt",
856 padding=True,
857 truncation=truncate,
858 max_length=self.cfg.n_ctx if truncate else None,
859 )["input_ids"]
861 if not self.cfg.default_prepend_bos and self.cfg.tokenizer_prepends_bos:
862 # We don't want to prepend bos but the tokenizer does it automatically, so we remove it manually
863 tokens = utils.get_tokens_with_bos_removed(self.tokenizer, tokens)
865 if move_to_device:
866 tokens = tokens.to(self.cfg.device)
867 return tokens
869 def to_string(
870 self,
871 tokens: Union[
872 List[int],
873 Int[torch.Tensor, ""],
874 Int[torch.Tensor, "batch pos"],
875 Int[torch.Tensor, "pos"],
876 np.ndarray,
877 List[Int[torch.Tensor, "pos"]],
878 ],
879 ) -> Union[str, List[str]]:
880 """Tokens to String(s).
882 Converts a tensor of tokens to a string (if rank 1) or a list of strings (if rank 2).
884 Accepts lists of tokens and numpy arrays as inputs too (and converts to tensors internally)
885 """
886 assert self.tokenizer is not None, "Cannot use to_string without a tokenizer"
888 if not isinstance(tokens, torch.Tensor):
889 # We allow lists to be input
890 tokens = torch.tensor(tokens)
892 # I'm not sure what exactly clean_up_tokenization_spaces does, but if
893 # it's set, then tokenization is no longer invertible, and some tokens
894 # with a bunch of whitespace get collapsed together
895 if len(tokens.shape) == 2:
896 return self.tokenizer.batch_decode(tokens, clean_up_tokenization_spaces=False)
897 elif len(tokens.shape) <= 1: 897 ↛ 900line 897 didn't jump to line 900 because the condition on line 897 was always true
898 return self.tokenizer.decode(tokens, clean_up_tokenization_spaces=False)
899 else:
900 raise ValueError(f"Invalid shape passed in: {tokens.shape}")
902 def to_str_tokens(
903 self,
904 input: Union[
905 str,
906 Int[torch.Tensor, "pos"],
907 Int[torch.Tensor, "1 pos"],
908 Int[np.ndarray, "pos"],
909 Int[np.ndarray, "1 pos"],
910 list,
911 ],
912 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
913 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
914 ) -> Union[List[str], List[List[str]]]:
915 """Map text, a list of text or tokens to a list of tokens as strings.
917 See the class-level "Tokenization notes" for full ``prepend_bos``
918 semantics. **Pass ``prepend_bos=False`` whenever you're tokenizing
919 only part of a prompt.**
921 String inputs that exceed ``model.cfg.n_ctx`` are truncated.
923 Args:
924 input (Union[str, list, torch.Tensor]): The input - either a string or a tensor of
925 tokens. If tokens, should be a tensor of shape [pos] or [1, pos].
926 prepend_bos (bool, optional): Overrides ``self.cfg.default_prepend_bos``. Only
927 applies when ``input`` is a string. Defaults to ``USE_DEFAULT_VALUE``
928 (use the cfg setting). Pass ``True`` or ``False`` to override locally.
929 padding_side (Union[Literal["left", "right"], None], optional): Overrides
930 self.tokenizer.padding_side. Specifies which side to pad when tokenizing multiple
931 strings of different lengths.
933 Returns:
934 str_tokens: List of individual tokens as strings
935 """
936 with utils.LocallyOverridenDefaults(
937 self, prepend_bos=prepend_bos, padding_side=padding_side
938 ):
939 assert self.tokenizer is not None # keep mypy happy
940 tokens: Union[np.ndarray, torch.Tensor]
941 if isinstance(input, list):
942 return list(
943 map(
944 lambda tokens: self.to_str_tokens(tokens, prepend_bos, padding_side),
945 input,
946 )
947 ) # type: ignore
948 elif isinstance(input, str):
949 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)[
950 0
951 ]
952 # Gemma tokenizer expects a batch dimension
953 if "gemma" in self.tokenizer.name_or_path and tokens.ndim == 1: 953 ↛ 954line 953 didn't jump to line 954 because the condition on line 953 was never true
954 tokens = tokens.unsqueeze(1)
955 elif isinstance(input, torch.Tensor):
956 tokens = input
957 tokens = tokens.squeeze() # Get rid of a trivial batch dimension
958 if tokens.dim() == 0:
959 # Don't pass dimensionless tensor
960 tokens = tokens.unsqueeze(0)
961 assert (
962 tokens.dim() == 1
963 ), f"Invalid tokens input to to_str_tokens, has shape: {tokens.shape}"
964 elif isinstance(input, np.ndarray): 964 ↛ 974line 964 didn't jump to line 974 because the condition on line 964 was always true
965 tokens = input
966 tokens = tokens.squeeze() # Get rid of a trivial batch dimension
967 if tokens.ndim == 0:
968 # Don't pass dimensionless tensor
969 tokens = np.expand_dims(tokens, axis=0)
970 assert (
971 tokens.ndim == 1
972 ), f"Invalid tokens input to to_str_tokens, has shape: {tokens.shape}"
973 else:
974 raise ValueError(f"Invalid input type to to_str_tokens: {type(input)}")
975 # v5 compat: wrap each token so batch_decode decodes them individually
976 if isinstance(tokens, np.ndarray):
977 tokens_list = [[int(t)] for t in tokens]
978 else:
979 tokens_list = [[int(t)] for t in tokens.tolist()]
980 str_tokens = self.tokenizer.batch_decode(
981 tokens_list, clean_up_tokenization_spaces=False
982 )
983 return str_tokens
985 def to_single_token(self, string):
986 """Map a string that makes up a single token to the id for that token.
988 Raises an error for strings that are not a single token! If uncertain use to_tokens.
989 """
991 # We use the to_tokens method, do not append a BOS token
992 token = self.to_tokens(string, prepend_bos=False).squeeze()
993 # If token shape is non-empty, raise error
994 assert not token.shape, f"Input string: {string} is not a single token!"
995 return token.item()
997 def to_single_str_token(self, int_token: int) -> str:
998 # Gives the single token corresponding to an int in string form
999 assert isinstance(int_token, int)
1000 token = self.to_str_tokens(torch.tensor([int_token]))
1001 assert len(token) == 1
1002 return cast(str, token[0])
1004 def get_token_position(
1005 self,
1006 single_token: Union[str, int],
1007 input: Union[str, Union[Float[torch.Tensor, "pos"], Float[torch.Tensor, "1 pos"]]],
1008 mode="first",
1009 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
1010 padding_side: Optional[Union[Literal["left", "right"], None]] = USE_DEFAULT_VALUE,
1011 ):
1012 """Get the position of a single_token in a string or sequence of tokens.
1014 Raises an error if the token is not present.
1016 When ``input`` is a string it's tokenized internally — see the
1017 class-level "Tokenization notes" for ``prepend_bos`` semantics.
1018 Off-by-one position errors usually mean ``prepend_bos`` is on
1019 when it shouldn't be (or vice versa); pass ``prepend_bos=False``
1020 when ``input`` is a fragment of a larger prompt.
1022 Args:
1023 single_token (Union[str, int]): The token to search for. Can
1024 be a token index, or a string (but the string must correspond to a single token).
1025 input (Union[str, torch.Tensor]): The sequence to
1026 search in. Can be a string or a rank 1 tensor of tokens or a rank 2 tensor of tokens
1027 with a dummy batch dimension.
1028 mode (str, optional): If there are multiple matches, which match to return. Supports
1029 "first" or "last". Defaults to "first".
1030 prepend_bos (bool, optional): Overrides ``self.cfg.default_prepend_bos``. Only
1031 applies when ``input`` is a string. Defaults to ``USE_DEFAULT_VALUE``
1032 (use the cfg setting). Pass ``True`` or ``False`` to override locally.
1033 padding_side (Union[Literal["left", "right"], None], optional): Overrides
1034 self.tokenizer.padding_side. Specifies which side to pad when tokenizing multiple
1035 strings of different lengths.
1036 """
1037 if isinstance(input, str):
1038 # If the input is a string, convert to tensor
1039 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
1040 else:
1041 tokens = input
1043 if len(tokens.shape) == 2:
1044 # If the tokens have shape [1, seq_len], flatten to [seq_len]
1045 assert (
1046 tokens.shape[0] == 1
1047 ), f"If tokens are rank two, they must have shape [1, seq_len], not {tokens.shape}"
1048 tokens = tokens[0]
1050 if isinstance(single_token, str):
1051 # If the single token is a string, convert to an integer
1052 single_token = self.to_single_token(single_token)
1053 elif isinstance(single_token, torch.Tensor): 1053 ↛ 1054line 1053 didn't jump to line 1054 because the condition on line 1053 was never true
1054 single_token = single_token.item()
1056 indices = torch.arange(len(tokens), device=tokens.device)[tokens == single_token]
1057 assert len(indices) > 0, "The token does not occur in the prompt"
1058 if mode == "first":
1059 return indices[0].item()
1060 elif mode == "last": 1060 ↛ 1063line 1060 didn't jump to line 1063 because the condition on line 1060 was always true
1061 return indices[-1].item()
1062 else:
1063 raise ValueError(f"mode must be 'first' or 'last', not {mode}")
1065 def tokens_to_residual_directions(
1066 self,
1067 tokens: Union[
1068 str,
1069 int,
1070 Int[torch.Tensor, ""],
1071 Int[torch.Tensor, "pos"],
1072 Int[torch.Tensor, "batch pos"],
1073 ],
1074 ) -> Union[
1075 Float[torch.Tensor, "d_model"],
1076 Float[torch.Tensor, "pos d_model"],
1077 Float[torch.Tensor, "batch pos d_model"],
1078 ]:
1079 """Map tokens to a tensor with the unembedding vector for those tokens.
1081 I.e. the vector in the residual stream that we dot with to the get the logit for that token.
1083 WARNING: If you use this without folding in LayerNorm, the results will be misleading and
1084 may be incorrect, as the LN weights change the unembed map. This is done automatically with
1085 the fold_ln flag on from_pretrained
1087 WARNING 2: LayerNorm scaling will scale up or down the effective direction in the residual
1088 stream for each output token on any given input token position.
1089 ActivationCache.apply_ln_to_stack will apply the appropriate scaling to these directions.
1091 Args:
1092 tokens (Union[str, int, torch.Tensor]): The token(s). If a single token, can be a single
1093 element tensor, an integer, or string. If string, will be mapped to a single token
1094 using to_single_token, and an error raised if it's multiple tokens. The method also
1095 works for a batch of input tokens.
1097 Returns:
1098 residual_direction torch.Tensor: The unembedding vector for the token(s), a stack of
1099 [d_model] tensor.
1100 """
1101 if isinstance(tokens, torch.Tensor) and tokens.numel() > 1:
1102 # If the tokens are a tensor, and have more than one element, assume they are a batch of
1103 # tokens.
1104 residual_directions = self.W_U[:, tokens]
1105 residual_directions = einops.rearrange(
1106 residual_directions, "d_model ... -> ... d_model"
1107 )
1108 return residual_directions
1109 else:
1110 # Otherwise there is a single token
1111 if isinstance(tokens, str):
1112 token = self.to_single_token(tokens)
1113 elif isinstance(tokens, int):
1114 token = tokens
1115 elif isinstance(tokens, torch.Tensor) and tokens.numel() == 1: 1115 ↛ 1118line 1115 didn't jump to line 1118 because the condition on line 1115 was always true
1116 token = tokens.item()
1117 else:
1118 raise ValueError(f"Invalid token type: {type(tokens)}")
1119 residual_direction = self.W_U[:, token]
1120 return residual_direction
1122 def to( # type: ignore
1123 self,
1124 device_or_dtype: Union[torch.device, str, torch.dtype],
1125 print_details: bool = True,
1126 ):
1127 return move_to_and_update_config(self, device_or_dtype, print_details)
1129 def cuda(self: T, device: Optional[Union[int, torch.device]] = None) -> T:
1130 # TODO: Add support for kwargs
1131 if isinstance(device, int):
1132 return self.to(f"cuda:{device}")
1133 elif device is None:
1134 return self.to("cuda")
1135 else:
1136 return self.to(device)
1138 def cpu(self: T) -> T:
1139 return self.to(torch.device("cpu"))
1141 def mps(self: T) -> T:
1142 """Warning: MPS may produce silently incorrect results. See #1178."""
1143 return self.to(torch.device("mps"))
1145 def move_model_modules_to_device(self):
1146 self.embed.to(get_best_available_device(self.cfg))
1147 self.hook_embed.to(get_best_available_device(self.cfg))
1148 if self.cfg.positional_embedding_type != "rotary":
1149 self.pos_embed.to(get_best_available_device(self.cfg))
1150 self.hook_pos_embed.to(get_best_available_device(self.cfg))
1152 if hasattr(self, "ln_final"):
1153 self.ln_final.to(get_best_available_device(self.cfg))
1154 self.unembed.to(get_best_available_device(self.cfg))
1155 for i, block in enumerate(self.blocks):
1156 block.to(get_best_available_device(self.cfg))
1158 @classmethod
1159 def from_pretrained(
1160 cls: Type[T],
1161 model_name: str,
1162 fold_ln: bool = True,
1163 center_writing_weights: bool = True,
1164 center_unembed: bool = True,
1165 refactor_factored_attn_matrices: bool = False,
1166 checkpoint_index: Optional[int] = None,
1167 checkpoint_value: Optional[int] = None,
1168 checkpoint_label: Optional[int] = None,
1169 hf_model: Optional[PreTrainedModel] = None,
1170 device: Optional[Union[str, torch.device]] = None,
1171 n_devices: int = 1,
1172 tokenizer: Optional[PreTrainedTokenizerBase] = None,
1173 move_to_device: bool = True,
1174 fold_value_biases: bool = True,
1175 default_prepend_bos: Optional[bool] = None,
1176 default_padding_side: Optional[Literal["left", "right"]] = None,
1177 dtype="float32",
1178 first_n_layers: Optional[int] = None,
1179 n_ctx: Optional[int] = None,
1180 **from_pretrained_kwargs,
1181 ) -> T:
1182 """Load in a Pretrained Model.
1184 Load in pretrained model weights to the HookedTransformer format and optionally to do some
1185 processing to make the model easier to interpret. Currently supports loading from most
1186 autoregressive HuggingFace models (``gpt2``, ``neo``, ``gptj``, ``opt``...) and from a range
1187 of toy models and SoLU models trained by Neel Nanda. The full list is available in the docs
1188 under :doc:`model properties</generated/model_properties_table>`. Also supports loading from
1189 a checkpoint for checkpointed models (currently, models trained by NeelNanda and the
1190 stanford-crfm models (using parameters ``checkpoint_index`` and ``checkpoint_value``).
1192 See :meth:`load_and_process_state_dict` for details on the processing (folding layer norm,
1193 centering the unembedding and centering the writing weights).
1195 Example:
1197 >>> from transformer_lens import HookedTransformer
1198 >>> model = HookedTransformer.from_pretrained("tiny-stories-1M")
1199 Loaded pretrained model tiny-stories-1M into HookedTransformer
1201 Args:
1202 model_name: The model name - must be an element of
1203 :const:`transformer_lens.loading_from_pretrained.OFFICIAL_MODEL_NAMES` or an alias
1204 of one. The full list of available models can be found in the docs under :doc:`model
1205 properties</generated/model_properties_table>`.
1206 fold_ln: Whether to fold in the LayerNorm weights to the
1207 subsequent linear layer. This does not change the computation.
1209 `LayerNorm
1210 <https://wandb.ai/wandb_fc/LayerNorm/reports/Layer-Normalization-in-Pytorch-With-Examples---VmlldzoxMjk5MTk1>`_
1211 is a common regularization technique used in transformers. Unlike BatchNorm, it
1212 cannot be turned off at inference time, as it significantly alters the mathematical
1213 function implemented by the transformer.
1215 When `fold_ln` is set to True, LayerNorm (with weights :math:`w_{ln}` and
1216 :math:`b_{ln}`) followed by a linear layer (:math:`W + b`) is optimized to
1217 LayerNormPre (just centering & normalizing) followed by a new linear layer with
1218 :math:`W_{eff} = w[:, \text{None}] * W` (element-wise multiplication) and
1219 :math:`b_{eff} = b + b_{ln} @ W`. This transformation is computationally equivalent
1220 and simplifies the model's interpretability. It essentially merges LayerNorm weights
1221 into the subsequent linear layer's weights, which is handled by HookedTransformer
1222 when loading pre-trained weights. Set `fold_ln` to False when loading a state dict
1223 if you wish to turn this off.
1225 Mathematically, LayerNorm is defined as follows:
1227 .. math::
1228 x_1 &= x_0 - \\text{mean}(x_0)
1230 x_2 &= \\frac{x_1}{\\sqrt{\\text{mean}(x_1^2)}}
1232 x_3 &= x_2 \\cdot w
1234 x_4 &= x_3 + b
1236 For further details, refer to `this document
1237 <https://transformer-circuits.pub/2021/framework/index.html#:~:text=Handling%20Layer%20Normalization>`_.
1238 center_writing_weights: Whether to center weights
1239 writing to the residual stream (ie set mean to be zero). Due to LayerNorm this
1240 doesn't change the computation.
1242 A related idea to folding layernorm (``fold_ln``) - *every* component reading an
1243 input from the residual stream is preceded by a LayerNorm, which means that the mean
1244 of a residual stream vector (ie the component in the direction of all ones) never
1245 matters. This means we can remove the all ones component of weights and biases whose
1246 output *writes* to the residual stream. Mathematically, ``W_writing -=
1247 W_writing.mean(dim=1, keepdim=True)``.
1248 center_unembed: Whether to center W_U (ie set mean
1249 to be zero). Softmax is translation invariant so this doesn't affect log probs or
1250 loss, but does change logits.
1252 The logits are fed into a softmax. Softmax is translation invariant (eg, adding 1 to
1253 every logit doesn't change the output), so we can simplify things by setting the
1254 mean of the logits to be zero. This is equivalent to setting the mean of every
1255 output vector of ``W_U`` to zero. In code, ``W_U -= W_U.mean(dim=-1,
1256 keepdim=True)``.
1257 refactor_factored_attn_matrices: Whether to convert the factored
1258 matrices (W_Q & W_K, and W_O & W_V) to be "even". Defaults to False
1259 checkpoint_index: If loading from a checkpoint, the index of
1260 the checkpoint to load.
1261 checkpoint_value: If loading from a checkpoint, the value of
1262 the checkpoint to load, ie the step or token number (each model has checkpoints
1263 labelled with exactly one of these). E.g. ``1000`` for a checkpoint taken at step
1264 1000 or after 1000 tokens. If `checkpoint_index` is also specified, this will be
1265 ignored.
1266 checkpoint_label: Alias for ``checkpoint_value`` kept for backwards compatibility with
1267 older docs and downstream code. Cannot be combined with ``checkpoint_value``.
1268 hf_model: If you have already loaded in the
1269 HuggingFace model, you can pass it in here rather than needing to recreate the
1270 object. Defaults to None.
1271 device: The device to load the model onto. By
1272 default will load to CUDA if available, else CPU.
1273 n_devices: The number of devices to split the model
1274 across. Defaults to 1. If greater than 1, `device` must be cuda.
1275 tokenizer: The tokenizer to use for the model. If not
1276 provided, it is inferred from cfg.tokenizer_name or initialized to None. If None,
1277 then the model cannot be passed strings, and d_vocab must be explicitly set.
1278 move_to_device: Whether to move the model to the device specified in
1279 cfg. device. Must be true if `n_devices` in the config is greater than 1, since the
1280 model's layers will be split across multiple devices.
1281 fold_value_biases: Each attention head has a value bias. Values are averaged to create
1282 mixed values (``z``), weighted by the attention pattern, but as the bias is
1283 constant, its contribution to ``z`` is exactly the same. The output of a head is ``z
1284 @ W_O``, and so the value bias just linearly adds to the output of the head. This
1285 means that the value bias of a head has nothing to do with the head, and is just a
1286 constant added to the attention layer outputs. We can take the sum across these and
1287 b_O to get an "effective bias" for the layer. In code, we set ``b_V=0``. and ``b_O =
1288 (b_V @ W_O).sum(dim=0) + b_O``.
1290 The technical derivation of this is as follows. ``v = residual @ W_V[h] +
1291 broadcast_b_V[h]`` for each head ``h`` (where ``b_V`` is broadcast up from shape
1292 ``d_head`` to shape ``[position, d_head]``). And ``z = pattern[h] @ v = pattern[h] @
1293 residual @ W_V[h] + pattern[h] @ broadcast_b_V[h]``. Because ``pattern[h]`` is
1294 ``[destination_position, source_position]`` and ``broadcast_b_V`` is constant along
1295 the ``(source_)position`` dimension, we're basically just multiplying it by the sum
1296 of the pattern across the ``source_position`` dimension, which is just ``1``. So it
1297 remains exactly the same, and so is just broadcast across the destination positions.
1298 default_prepend_bos: Default behavior of whether to prepend the BOS
1299 token when the methods of HookedTransformer process input text to tokenize (only
1300 when input is a string).
1301 Resolution order for default_prepend_bos:
1302 1. If user passes value explicitly, use that value
1303 2. Model-specific default from cfg_dict if it exists (e.g. for bloom models it's False)
1304 3. Global default (True)
1306 Even for models not explicitly trained with the BOS token, heads often use the first position as a resting position
1307 and accordingly lose information from the first token, so this empirically seems to give better
1308 results. Note that you can also locally override the default behavior by passing in
1309 prepend_bos=True/False when you call a method that processes the input string.
1310 from_pretrained_kwargs: Any other optional argument passed to
1311 HuggingFace's from_pretrained (e.g. "cache_dir" or "torch_dtype"). Also passed to
1312 other HuggingFace functions when compatible. For some models or arguments it doesn't
1313 work, especially for models that are not internally loaded with HuggingFace's
1314 from_pretrained (e.g. SoLU models).
1315 dtype: What data type to load the model in (also sets the dtype of
1316 the HuggingFace model). Set to bfloat16 or float16 if you get out of memory errors when loading
1317 the model.
1318 default_padding_side: Which side to pad on when tokenizing.
1319 Resolution order for default_padding_side:
1320 1. If user passes value explicitly, use that value
1321 2. If tokenizer has a default padding side, use that value
1322 3. Global default ("right")
1323 first_n_layers: If specified, only load the first n layers of the model.
1324 """
1325 import warnings
1327 warnings.warn(
1328 "HookedTransformer.from_pretrained is deprecated and will be removed in a "
1329 "future major release. Use TransformerBridge.boot_transformers(...) instead, "
1330 "then call enable_compatibility_mode() for HookedTransformer-equivalent "
1331 "numerics. See docs/source/content/migrating_to_v3.md.",
1332 DeprecationWarning,
1333 stacklevel=2,
1334 )
1336 if checkpoint_value is not None and checkpoint_label is not None:
1337 raise ValueError(
1338 "Specify checkpoint_value or checkpoint_label, not both — they are aliases."
1339 )
1340 elif checkpoint_label is not None:
1341 checkpoint_value = checkpoint_label
1343 if model_name.lower().startswith("t5"): 1343 ↛ 1344line 1343 didn't jump to line 1344 because the condition on line 1343 was never true
1344 raise RuntimeError(
1345 "Execution stopped: Please use HookedEncoderDecoder to load T5 models instead of HookedTransformer."
1346 )
1348 if model_name.lower().startswith("bert"): 1348 ↛ 1349line 1348 didn't jump to line 1349 because the condition on line 1348 was never true
1349 raise RuntimeError(
1350 "Execution stopped: Please use HookedEncoder to load BERT-style models instead of HookedTransformer."
1351 )
1353 assert not (
1354 from_pretrained_kwargs.get("load_in_8bit", False)
1355 or from_pretrained_kwargs.get("load_in_4bit", False)
1356 ), "Quantization not supported"
1358 if hf_model is not None: 1358 ↛ 1359line 1358 didn't jump to line 1359 because the condition on line 1358 was never true
1359 assert hasattr(hf_model, "config"), "PreTrainedModel must have a config attribute"
1360 hf_cfg = hf_model.config.to_dict()
1361 qc = hf_cfg.get("quantization_config", {})
1362 load_in_4bit = qc.get("load_in_4bit", False)
1363 load_in_8bit = qc.get("load_in_8bit", False)
1364 quant_method = qc.get("quant_method", "")
1365 assert not load_in_8bit, "8-bit quantization is not supported"
1366 assert not (
1367 load_in_4bit and ("llama" not in model_name.lower())
1368 ), "Quantization is only supported for Llama models"
1369 if load_in_4bit:
1370 assert (
1371 qc.get("quant_method", "") == "bitsandbytes"
1372 ), "Only bitsandbytes quantization is supported"
1373 elif quant_method:
1374 # Anything other than the supported bitsandbytes 4-bit Llama
1375 # flow reaches the converters, which slice `.weight` directly:
1376 # packed or scale-separated storage yields wrong numbers rather
1377 # than an error. Refuse instead.
1378 raise NotImplementedError(
1379 f"HookedTransformer cannot convert a {quant_method!r}-quantized "
1380 "checkpoint: its weight converters read weights directly, and "
1381 "packed or scale-separated storage would silently produce wrong "
1382 "values. Load the model dequantized, or use TransformerBridge "
1383 "for a quantized forward pass."
1384 )
1385 else:
1386 hf_cfg = {}
1388 if isinstance(dtype, str):
1389 # Convert from string to a torch dtype
1390 dtype = DTYPE_FROM_STRING[dtype]
1391 if "torch_dtype" in from_pretrained_kwargs: 1391 ↛ 1393line 1391 didn't jump to line 1393 because the condition on line 1391 was never true
1392 # Backwards compat: torch_dtype overrides dtype
1393 dtype = from_pretrained_kwargs["torch_dtype"]
1395 if ( 1395 ↛ 1399line 1395 didn't jump to line 1399 because the condition on line 1395 was never true
1396 (from_pretrained_kwargs.get("torch_dtype", None) == torch.float16)
1397 or dtype == torch.float16
1398 ) and device in ["cpu", None]:
1399 logging.warning("float16 models may not work on CPU. Consider using a GPU or bfloat16.")
1401 # Get the model name used in HuggingFace, rather than the alias.
1402 official_model_name = loading.get_official_model_name(model_name)
1404 # Load config (includes checkpoint info if applicable)
1405 cfg = loading.get_pretrained_model_config(
1406 official_model_name,
1407 hf_cfg=hf_cfg,
1408 checkpoint_index=checkpoint_index,
1409 checkpoint_value=checkpoint_value,
1410 fold_ln=fold_ln,
1411 device=device,
1412 n_devices=n_devices,
1413 default_prepend_bos=default_prepend_bos,
1414 dtype=dtype,
1415 first_n_layers=first_n_layers,
1416 n_ctx=n_ctx,
1417 **from_pretrained_kwargs,
1418 )
1420 if cfg.positional_embedding_type == "shortformer":
1421 if fold_ln: 1421 ↛ 1427line 1421 didn't jump to line 1427 because the condition on line 1421 was always true
1422 logging.warning(
1423 "You tried to specify fold_ln=True for a shortformer model, but this can't be done! Setting fold_"
1424 "ln=False instead."
1425 )
1426 fold_ln = False
1427 if center_unembed: 1427 ↛ 1433line 1427 didn't jump to line 1433 because the condition on line 1427 was always true
1428 logging.warning(
1429 "You tried to specify center_unembed=True for a shortformer model, but this can't be done! "
1430 "Setting center_unembed=False instead."
1431 )
1432 center_unembed = False
1433 if center_writing_weights: 1433 ↛ 1441line 1433 didn't jump to line 1441 because the condition on line 1433 was always true
1434 logging.warning(
1435 "You tried to specify center_writing_weights=True for a shortformer model, but this can't be done! "
1436 "Setting center_writing_weights=False instead."
1437 )
1438 center_writing_weights = False
1439 # Post-norm architectures are incompatible with fold_ln/center_writing_weights,
1440 # both of which assume the norm gain sits on a sublayer's input.
1441 if cfg.original_architecture in POST_NORM_ARCHITECTURES: 1441 ↛ 1442line 1441 didn't jump to line 1442 because the condition on line 1441 was never true
1442 if fold_ln:
1443 logging.warning(
1444 f"fold_ln=True is incompatible with {cfg.original_architecture}'s "
1445 "post-norm architecture. Setting fold_ln=False."
1446 )
1447 fold_ln = False
1448 if center_writing_weights:
1449 logging.warning(
1450 f"center_writing_weights=True is incompatible with "
1451 f"{cfg.original_architecture}'s post-norm architecture. "
1452 "Setting center_writing_weights=False."
1453 )
1454 center_writing_weights = False
1455 if center_unembed and softcap_enabled(cfg.output_logits_soft_cap): 1455 ↛ 1456line 1455 didn't jump to line 1456 because the condition on line 1455 was never true
1456 logging.warning(
1457 "You tried to specify center_unembed=True for a model using logit softcap, but this can't be done! Softcapping is not invariant upon adding a constant "
1458 "Setting center_unembed=False instead."
1459 )
1460 center_unembed = False
1462 # Get the state dict of the model (ie a mapping of parameter names to tensors), processed to
1463 # match the HookedTransformer parameter names.
1464 state_dict = loading.get_pretrained_state_dict(
1465 official_model_name, cfg, hf_model, dtype=dtype, **from_pretrained_kwargs
1466 )
1468 # Create the HookedTransformer object
1469 model = cls(
1470 cfg,
1471 tokenizer,
1472 move_to_device=False,
1473 default_padding_side=default_padding_side,
1474 )
1476 model.load_and_process_state_dict(
1477 state_dict,
1478 fold_ln=fold_ln,
1479 center_writing_weights=center_writing_weights,
1480 center_unembed=center_unembed,
1481 fold_value_biases=fold_value_biases,
1482 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
1483 )
1485 if move_to_device: 1485 ↛ 1488line 1485 didn't jump to line 1488 because the condition on line 1485 was always true
1486 model.move_model_modules_to_device()
1488 print(f"Loaded pretrained model {model_name} into HookedTransformer")
1489 return model
1491 @classmethod
1492 def from_pretrained_no_processing(
1493 cls,
1494 model_name: str,
1495 fold_ln=False,
1496 center_writing_weights=False,
1497 center_unembed=False,
1498 refactor_factored_attn_matrices=False,
1499 fold_value_biases=False,
1500 dtype=torch.float32,
1501 default_prepend_bos=None,
1502 default_padding_side=None,
1503 **from_pretrained_kwargs,
1504 ):
1505 """Wrapper for from_pretrained.
1507 Wrapper for from_pretrained with all boolean flags related to simplifying the model set to
1508 False. Refer to from_pretrained for details.
1509 """
1510 return cls.from_pretrained(
1511 model_name,
1512 fold_ln=fold_ln,
1513 center_writing_weights=center_writing_weights,
1514 center_unembed=center_unembed,
1515 fold_value_biases=fold_value_biases,
1516 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
1517 dtype=dtype,
1518 default_prepend_bos=default_prepend_bos,
1519 default_padding_side=default_padding_side,
1520 **from_pretrained_kwargs,
1521 )
1523 def init_weights(self):
1524 """Initialize weights.
1526 LayerNorm weights are already initialized to 1.0, and all biases are initialized to 0.0
1527 (including LayerNorm), so this just initializes weight matrices.
1529 Weight matrices are set to empty by default (to save space + compute, since they're the bulk
1530 of the parameters), so it is important to call this if you are not loading in pretrained
1531 weights! Note that this function assumes that weight names being with `W_`.
1533 Set seed here to ensure determinism.
1535 This does NOT follow the default PyTorch scheme, which is the following: all linear layers
1536 use uniform(-1/sqrt(fan_in), 1/sqrt(fan_in)) for weights, and uniform(-1/sqrt(fan_in),
1537 1/sqrt(fan_in)) for biases. For biases, fan_in is computed using the fan_in for the weight
1538 matrix of the linear layer. Note that it *does not actually* use Kaiming initialization,
1539 despite the fact that it calls the function.
1541 However, for Transformer blocks, it instead initializes biases to zero and weights using Xavier uniform, that
1542 is: uniform(-sqrt(6 / (fan_in + fan_out)), sqrt(6 / (fan_in + fan_out))) for weights.
1544 We split off the initialization into separate functions because muP initialization handles
1545 different parts of the model differently.
1546 """
1548 if self.cfg.seed is not None: 1548 ↛ 1549line 1548 didn't jump to line 1549 because the condition on line 1548 was never true
1549 torch.manual_seed(self.cfg.seed)
1551 if self.cfg.init_mode == "gpt2": 1551 ↛ 1553line 1551 didn't jump to line 1553 because the condition on line 1551 was always true
1552 self._init_weights_gpt2()
1553 elif self.cfg.init_mode == "xavier_uniform":
1554 self._init_weights_xavier(dist_type="uniform")
1555 elif self.cfg.init_mode == "xavier_normal":
1556 self._init_weights_xavier(dist_type="normal")
1557 elif self.cfg.init_mode == "kaiming_uniform":
1558 self._init_weights_kaiming(dist_type="uniform")
1559 elif self.cfg.init_mode == "kaiming_normal":
1560 self._init_weights_kaiming(dist_type="normal")
1561 elif self.cfg.init_mode == "muP":
1562 self._init_weights_muP(dist_type="normal") # muP uses normal initialization
1564 def _init_weights_gpt2(self):
1565 """Initialize weights with GPT-2 initialization. Biases are initialized to 0.0 and weights
1566 are initialized to N(0, 0.64/d_model) if initializer_range is not set, otherwise std is initializer_range.
1567 """
1568 for name, param in self.named_parameters():
1569 if "W_" in name:
1570 nn.init.normal_(param, std=self.cfg.initializer_range)
1572 def _init_weights_xavier(self, dist_type="normal"):
1573 """
1574 Initialize weights with Xavier initialization -- that is, scale the weights by sqrt(6 /
1575 (fan_in + fan_out)) for a [-1, 1] uniform distribution, or sqrt(2 / (fan_in + fan_out)) for a
1576 standard normal.
1578 Note that since TransformerLens implements the matrices in the opposite orientation to what
1579 torch does (e.g. it's d_in x d_out, not d_out x d_in as in torch), we need to calculate it
1580 ourselves.
1581 """
1582 gain = self.cfg.initializer_range
1583 for name, param in self.named_parameters():
1584 if "W_" in name:
1585 if dist_type == "uniform":
1586 init_xavier_uniform_(param, gain=gain)
1587 elif dist_type == "normal":
1588 init_xavier_normal_(param, gain=gain)
1590 def _init_weights_kaiming(self, dist_type="uniform"):
1591 """
1592 Initialize weights with Kaiming initialization -- that is, scale the weights by
1593 c / sqrt(fan_in), where c = sqrt(2) if the params were immediately preceded by a relu and 1 for
1594 everything else.
1596 Note that the numbers are actually incorrect here when you're using a nonlinearity other
1597 than relu, e.g. the correct c for SiLu is ~1.74, for tanh it's 5/3 ~= 1.67, and for GeLU it's ~1.57.
1598 But this is unlikely to matter in practice.
1600 I'm just using fan_mode = "fan_in" for now, but it should be trivial to add fan_out.
1602 Again, we have to implement it ourselves because of the orientation of the matrices.
1603 """
1604 gain = self.cfg.initializer_range
1605 for name, param in self.named_parameters():
1606 if "W_" in name:
1607 if dist_type == "uniform":
1608 init_kaiming_uniform_(param, gain=gain, nonlinearity="relu", mode="fan_in")
1609 elif dist_type == "normal":
1610 init_kaiming_normal_(param, gain=gain, nonlinearity="relu", mode="fan_in")
1612 def _init_weights_muP(self, dist_type="uniform"):
1613 """
1614 Initialize weights with muParameterization. This involves scaling output weights by a factor
1615 of 1/fan_in, input weights and biases by 1, everything else by a factor of 1/sqrt(fan_in).
1617 Also, you need to use muAdamW, which rescales the learning rate for output weights and
1618 hidden weights by a factor of 1/fan_in.
1620 All biases are still assumed to be initialized to 0.0, so we only need to change the
1621 weights.
1622 """
1623 for name, param in self.named_parameters():
1624 if "W_" in name:
1625 fan_in, _ = utils.calc_fan_in_and_fan_out(param)
1626 if "embed" in name:
1627 scale = float(1)
1628 elif "unembed" in name:
1629 scale = 1 / fan_in
1630 else:
1631 scale = 1 / fan_in**0.5
1633 if dist_type == "uniform":
1634 scale *= 3**0.5
1635 nn.init.uniform_(param, -scale, scale)
1636 elif dist_type == "normal":
1637 nn.init.normal_(param, std=scale)
1639 def load_and_process_state_dict(
1640 self,
1641 state_dict: Dict[str, torch.Tensor],
1642 fold_ln: bool = True,
1643 center_writing_weights: bool = True,
1644 center_unembed: bool = True,
1645 fold_value_biases: bool = True,
1646 refactor_factored_attn_matrices: bool = False,
1647 ):
1648 """Load & Process State Dict.
1650 Load a state dict into the model, and to apply processing to simplify it. The state dict is
1651 assumed to be in the HookedTransformer format.
1653 See the relevant method (same name as the flag) for more details on the folding, centering
1654 and processing flags.
1656 Args:
1657 state_dict (dict): The state dict of the model, in HookedTransformer format. fold_ln
1658 fold_ln (bool, optional): Whether to fold in the LayerNorm weights to the
1659 subsequent linear layer. This does not change the computation. Defaults to True.
1660 center_writing_weights (bool, optional): Whether to center weights writing to the
1661 residual stream (ie set mean to be zero). Due to LayerNorm this doesn't change the
1662 computation. Defaults to True.
1663 center_unembed (bool, optional): Whether to center W_U (ie set mean to be zero).
1664 Softmax is translation invariant so this doesn't affect log probs or loss, but does
1665 change logits. Defaults to True.
1666 fold_value_biases (bool, optional): Whether to fold the value biases into the output
1667 bias. Because attention patterns add up to 1, the value biases always have a
1668 constant effect on a layer's output, and it doesn't matter which head a bias is
1669 associated with. We can factor this all into a single output bias to the layer, and
1670 make it easier to interpret the head's output.
1671 refactor_factored_attn_matrices (bool, optional): Whether to convert the factored
1672 matrices (W_Q & W_K, and W_O & W_V) to be "even". Defaults to False.
1673 model_name (str, optional): checks the model name for special cases of state dict
1674 loading. Only used for Redwood 2L model currently.
1675 """
1676 if self.cfg.dtype not in [torch.float32, torch.float64] and fold_ln: 1676 ↛ 1677line 1676 didn't jump to line 1677 because the condition on line 1676 was never true
1677 logging.warning(
1678 "With reduced precision, it is advised to use `from_pretrained_no_processing` instead of `from_pretrained`."
1679 )
1681 if ( 1681 ↛ 1686line 1681 didn't jump to line 1686 because the condition on line 1681 was never true
1682 self.cfg.dtype not in [torch.float32, torch.float64]
1683 and self.cfg.num_experts
1684 and self.cfg.num_experts > 1
1685 ):
1686 logging.warning(
1687 "When running MoE models, it is advised to use a higher precision data type. See docs for more info."
1688 )
1690 state_dict = self.fill_missing_keys(state_dict)
1691 if fold_ln:
1692 if self.cfg.normalization_type not in ["LN", "LNPre", "RMS", "RMSPre"]: 1692 ↛ 1693line 1692 didn't jump to line 1693 because the condition on line 1692 was never true
1693 logging.warning(
1694 "You are not using LayerNorm or RMSNorm, so the layer norm weights can't be folded! Skipping"
1695 )
1696 fold_ln = False
1697 else:
1698 ln_keys_present = any(
1699 k.endswith((".ln1.w", ".ln2.w", "ln_final.w")) for k in state_dict
1700 )
1701 if not ln_keys_present: 1701 ↛ 1702line 1701 didn't jump to line 1702 because the condition on line 1701 was never true
1702 logging.warning(
1703 "fold_ln=True but no LayerNorm weights found in state_dict. "
1704 "The model may have been saved with already-folded LayerNorms. "
1705 "Skipping fold."
1706 )
1707 fold_ln = False
1708 else:
1709 if self.cfg.normalization_type == "LN":
1710 self.cfg.normalization_type = "LNPre"
1711 self.ln_final = LayerNormPre(self.cfg)
1712 for layer in self.blocks:
1713 layer.ln1 = LayerNormPre(self.cfg)
1714 layer.ln2 = LayerNormPre(self.cfg)
1715 if self.cfg.is_layer_norm_activation(): 1715 ↛ 1716line 1715 didn't jump to line 1716 because the condition on line 1715 was never true
1716 layer.mlp.ln = LayerNormPre(self.cfg)
1717 elif self.cfg.normalization_type == "RMS":
1718 self.cfg.normalization_type = "RMSPre"
1719 self.ln_final = RMSNormPre(self.cfg)
1720 for layer in self.blocks:
1721 layer.ln1 = RMSNormPre(self.cfg)
1722 layer.ln2 = RMSNormPre(self.cfg)
1723 if self.cfg.is_layer_norm_activation(): 1723 ↛ 1724line 1723 didn't jump to line 1724 because the condition on line 1723 was never true
1724 layer.mlp.ln = RMSNormPre(self.cfg)
1726 # Use the centralized ProcessWeights class for all weight processing
1727 # (fold_ln is passed through — if we skipped above, it's now False)
1728 state_dict = ProcessWeights.process_weights(
1729 state_dict,
1730 self.cfg,
1731 fold_ln=fold_ln,
1732 center_writing_weights=center_writing_weights,
1733 center_unembed=center_unembed,
1734 fold_value_biases=fold_value_biases,
1735 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
1736 )
1738 if self.cfg.load_in_4bit: 1738 ↛ 1741line 1738 didn't jump to line 1741 because the condition on line 1738 was never true
1739 # with quantization, parameters should be assigned
1740 # so that quantization settings are not lost
1741 self.load_state_dict(state_dict, assign=True, strict=False)
1742 else:
1743 state_dict_keys = list(state_dict.keys())
1744 for key in state_dict_keys:
1745 self.load_state_dict({key: state_dict[key]}, strict=False)
1746 del state_dict[key]
1748 if fold_ln:
1749 self.setup()
1751 def fill_missing_keys(self, state_dict):
1752 return loading.fill_missing_keys(self, state_dict)
1754 def fold_layer_norm(
1755 self, state_dict: Dict[str, torch.Tensor], fold_biases=True, center_weights=True
1756 ):
1757 """Fold Layer Norm. Can also be used to fold RMS Norm, when fold_biases and center_weights are set to False.
1759 Takes in a state dict from a pretrained model, formatted to be consistent with
1760 HookedTransformer but with LayerNorm weights and biases. Folds these into the neighbouring
1761 weights. See further_comments.md for more details.
1763 Args:
1764 state_dict (Dict[str, torch.Tensor]): State dict of pretrained model.
1765 fold_biases (bool): Enables folding of LN biases. Should be disabled when RMS Norm is used.
1766 center_weights (bool): Enables the centering of weights after folding in LN. Should be disabled when RMS Norm is used.
1767 """
1768 return ProcessWeights.fold_layer_norm(state_dict, self.cfg, fold_biases, center_weights)
1770 def center_writing_weights(self, state_dict: Dict[str, torch.Tensor]):
1771 """Center Writing Weights.
1773 Centers the weights of the model that write to the residual stream - W_out, W_E, W_pos and
1774 W_out. This is done by subtracting the mean of the weights from the weights themselves. This
1775 is done in-place. See fold_layer_norm for more details.
1776 """
1777 return ProcessWeights.center_writing_weights(state_dict, self.cfg)
1779 def center_unembed(self, state_dict: Dict[str, torch.Tensor]):
1780 """Center the unembedding weights W_U.
1782 This is done by subtracting the mean of the weights from the weights themselves. This is
1783 done in-place. As softmax is translation invariant, this changes the logits but not the log
1784 probs, and makes the model logits (slightly) more interpretable - when trying to understand
1785 how components contribute to the logits, we'll be less misled by components that just add
1786 something to every logit.
1787 """
1788 return ProcessWeights.center_unembed(state_dict)
1790 def fold_value_biases(self, state_dict: Dict[str, torch.Tensor]):
1791 """Fold the value biases into the output bias.
1793 Because attention patterns add up to 1, the value biases always have a constant effect on a
1794 head's output. Further, as the outputs of each head in a layer add together, each head's
1795 value bias has a constant effect on the *layer's* output, which can make it harder to
1796 interpret the effect of any given head, and it doesn't matter which head a bias is
1797 associated with. We can factor this all into a single output bias to the layer, and make it
1798 easier to interpret the head's output. Formally, we take b_O_new = b_O_original +
1799 sum_head(b_V_head @ W_O_head).
1800 """
1801 return ProcessWeights.fold_value_biases(state_dict, self.cfg)
1803 def refactor_factored_attn_matrices(self, state_dict: Dict[str, torch.Tensor]):
1804 """Experimental method for managing queries, keys and values.
1806 As argued in [A Mathematical Framework for Transformer
1807 Circuits](https://transformer-circuits.pub/2021/framework/index.html), queries, keys and
1808 values are somewhat arbitrary intermediate terms when computing with the low rank factored
1809 matrices W_QK = W_Q @ W_K.T and W_OV = W_V @ W_O, and these matrices are the only thing
1810 determining head behaviour. But there are many ways to find a low rank factorization to a
1811 given matrix, and hopefully some of these are more interpretable than others! This method is
1812 one attempt, which makes all of the matrices have orthogonal rows or columns, W_O into a
1813 rotation and W_Q and W_K having the nth column in each having the same norm. The formula is
1814 $W_V = U @ S,W_O=Vh.T,W_Q=U@S.sqrt(),W_K=Vh@S.sqrt()$.
1816 More details:
1818 If W_OV = U @ S @ Vh.T in its singular value decomposition, (where S is in R^d_head not
1819 R^d_model, as W_OV is low rank), W_OV = (U @ S) @ (Vh.T) is an equivalent low rank
1820 factorisation, where rows/columns of each matrix are orthogonal! So setting $W_V=US$ and
1821 $W_O=Vh.T$ works just as well. I *think* this is a more interpretable setup, because now
1822 $W_O$ is just a rotation, and doesn't change the norm, so $z$ has the same norm as the
1823 result of the head.
1825 For $W_QK = W_Q @ W_K.T$ we use the refactor $W_Q = U @ S.sqrt()$ and $W_K = Vh @ S.sqrt()$,
1826 which is also equivalent ($S==S.sqrt() @ S.sqrt()$ as $S$ is diagonal). Here we keep the
1827 matrices as having the same norm, since there's not an obvious asymmetry between the keys
1828 and queries.
1830 Biases are more fiddly to deal with. For OV it's pretty easy - we just need (x @ W_V + b_V)
1831 @ W_O + b_O to be preserved, so we can set b_V' = 0. and b_O' = b_V @ W_O + b_O (note that
1832 b_V in R^{head_index x d_head} while b_O in R^{d_model}, so we need to sum b_V @ W_O along
1833 the head_index dimension too).
1835 For QK it's messy - we need to preserve the bilinear form of (x @ W_Q + b_Q) * (y @ W_K +
1836 b_K), which is fairly messy. To deal with the biases, we concatenate them to W_Q and W_K to
1837 simulate a d_model+1 dimensional input (whose final coordinate is always 1), do the SVD
1838 factorization on this effective matrix, then separate out into final weights and biases.
1839 """
1840 return ProcessWeights.refactor_factored_attn_matrices(state_dict, self.cfg)
1842 def set_use_attn_result(self, use_attn_result: bool):
1843 """Toggle whether to explicitly calculate and expose the result for each attention head.
1845 Useful for interpretability but can easily burn through GPU memory.
1846 """
1847 self.cfg.use_attn_result = use_attn_result
1849 def set_use_split_qkv_input(self, use_split_qkv_input: bool):
1850 """
1851 Toggles whether to allow editing of the separate Q, K, and V inputs to each attention head.
1852 """
1853 self.cfg.use_split_qkv_input = use_split_qkv_input
1855 def set_use_hook_mlp_in(self, use_hook_mlp_in: bool):
1856 """Toggles whether to allow storing and editing inputs to each MLP layer."""
1858 assert not self.cfg.attn_only, "Can't use hook_mlp_in with attn_only model"
1859 self.cfg.use_hook_mlp_in = use_hook_mlp_in
1861 def set_use_attn_in(self, use_attn_in: bool):
1862 """
1863 Toggles whether to allow editing of inputs to each attention head.
1864 """
1865 assert (
1866 self.cfg.n_key_value_heads is None
1867 ), "Can't use attn_in with GroupedQueryAttention, please use split_qkv_input instead"
1868 self.cfg.use_attn_in = use_attn_in
1870 def set_ungroup_grouped_query_attention(self, ungroup_grouped_query_attention: bool):
1871 """
1872 Toggles whether to ungroup the grouped key and value heads in models with grouped query attention (GQA).
1873 """
1874 self.cfg.ungroup_grouped_query_attention = ungroup_grouped_query_attention
1876 def process_weights_(
1877 self,
1878 fold_ln: bool = True,
1879 center_writing_weights: bool = True,
1880 center_unembed: bool = True,
1881 refactor_factored_attn_matrices: bool = False,
1882 ):
1883 """Wrapper around `load_and_process_state_dict`.
1885 Wrapper around load_and_process_state_dict to allow for in-place processing of the weights.
1886 This is useful if using HookedTransformer for training, if we then want to analyse a cleaner
1887 version of the same model.
1888 """
1889 state_dict = self.state_dict()
1890 self.load_and_process_state_dict(
1891 state_dict,
1892 fold_ln=fold_ln,
1893 center_writing_weights=center_writing_weights,
1894 center_unembed=center_unembed,
1895 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
1896 )
1898 @torch.inference_mode()
1899 def generate(
1900 self,
1901 input: Union[
1902 str,
1903 List[str],
1904 Int[torch.Tensor, "batch pos"],
1905 Float[torch.Tensor, "batch pos hidden_size"],
1906 ] = "",
1907 max_new_tokens: int = 10,
1908 stop_at_eos: bool = True,
1909 eos_token_id: Optional[int] = None,
1910 do_sample: bool = True,
1911 top_k: Optional[int] = None,
1912 top_p: Optional[float] = None,
1913 temperature: float = 1.0,
1914 freq_penalty: float = 0.0,
1915 use_past_kv_cache: bool = True,
1916 prepend_bos: Optional[bool] = USE_DEFAULT_VALUE,
1917 padding_side: Optional[Literal["left", "right"]] = USE_DEFAULT_VALUE,
1918 return_type: Optional[str] = "input",
1919 verbose: bool = True,
1920 **generation_kwargs,
1921 ) -> Union[
1922 str,
1923 List[str],
1924 Int[torch.Tensor, "batch pos_plus_new_tokens"],
1925 Float[torch.Tensor, "batch pos_plus_new_tokens hidden_size"],
1926 Any, # transformers.utils.ModelOutput to accommodate output_logits=True.
1927 # Using Any due to beartype's forward reference resolution limitations.
1928 # See: https://github.com/beartype/beartype/issues/546
1929 ]:
1930 """Sample Tokens from the Model.
1932 Sample tokens from the model until the model outputs eos_token or max_new_tokens is reached.
1934 To avoid fiddling with ragged tensors, if we input a batch of text and some sequences finish
1935 (by producing an EOT token), we keep running the model on the entire batch, but throw away
1936 the output for a finished sequence and just keep adding EOTs to pad.
1938 Args:
1939 input (Union[str, List[str], Int[torch.Tensor, "batch pos"], Float[torch.Tensor, "batch pos hidden_size"]]):
1940 A text string (this will be converted to a batch of tokens with batch
1941 size 1), a list of strings, batch of tokens or a tensor of precomputed embeddings of shape
1942 [batch, pos, hidden_size].
1943 max_new_tokens (int): Maximum number of tokens to generate.
1944 stop_at_eos (bool): If True, stop generating tokens when the model outputs eos_token.
1945 eos_token_id (Optional[Union[int, Sequence]]): The token ID to use for end
1946 of sentence. If None, use the tokenizer's eos_token_id - required if using
1947 stop_at_eos. It's also possible to provide a list of token IDs (not just the
1948 eos_token_id), in which case the generation will stop when any of them are output
1949 (useful e.g. for stable_lm).
1950 do_sample (bool): If True, sample from the model's output distribution. Otherwise, use
1951 greedy search (take the max logit each time).
1952 top_k (int): Number of tokens to sample from. If None, sample from all tokens.
1953 top_p (float): Probability mass to sample from. If 1.0, sample from all tokens. If <1.0,
1954 we take the top tokens with cumulative probability >= top_p.
1955 temperature (float): Temperature for sampling. Higher values will make the model more
1956 random (limit of temp -> 0 is just taking the top token, limit of temp -> inf is
1957 sampling from a uniform distribution).
1958 freq_penalty (float): Frequency penalty for sampling - how much to penalise previous
1959 tokens. Higher values will make the model more random. Works only with str and tokens input.
1960 use_past_kv_cache (bool): If True, create and use cache to speed up generation.
1961 prepend_bos (bool, optional): Overrides self.cfg.default_prepend_bos. Whether to prepend
1962 the BOS token to the input (applicable when input is a string). Defaults to None,
1963 implying usage of self.cfg.default_prepend_bos (default is True unless specified
1964 otherwise). Pass True or False to override the default.
1965 padding_side (Union[Literal["left", "right"], None], optional): Overrides
1966 self.tokenizer.padding_side. Specifies which side to pad when tokenizing
1967 multiple strings of different lengths. For batched list inputs, left-padding
1968 is forced internally for correct generation behavior.
1969 return_type (Optional[str]): The type of the output to return - a string or a list of strings ('str'),
1970 a tensor of tokens ('tokens'), a tensor of output embeddings ('embeds') or whatever the format of the
1971 input was ('input').
1972 verbose (bool): If True, show tqdm progress bars for generation.
1974 Returns:
1975 outputs (str, List[str], Int[torch.Tensor, "batch pos_plus_new_tokens"], Float[torch.Tensor,
1976 "batch pos_plus_new_tokens hidden_size"]): generated sequence. Str, tokens or embeddings.
1977 If input is embeddings and return type is tokens or string, returns only new generated sequence.
1978 In other cases returns sequence including input sequence.
1979 """
1981 with utils.LocallyOverridenDefaults(
1982 self, prepend_bos=prepend_bos, padding_side=padding_side
1983 ):
1984 assert isinstance(input, (str, torch.Tensor, list)) and (
1985 isinstance(input, list)
1986 and all(isinstance(i, str) for i in input)
1987 or not isinstance(input, list)
1988 ), "Input must be either string, torch.Tensor, or List[str]"
1990 assert return_type in [
1991 "input",
1992 "str",
1993 "tokens",
1994 "embeds",
1995 ], "return_type must be one of ['input', 'str', 'tokens', 'embeds']"
1997 if return_type == "input":
1998 if isinstance(input, (str, list)):
1999 return_type = "str"
2000 elif input.ndim == 2:
2001 return_type = "tokens"
2002 else:
2003 return_type = "embeds"
2005 # initial_attention_mask is always computed so that single-prompt and
2006 # batched generation go through the same masked code path, producing
2007 # consistent results for the same prompt regardless of batching.
2008 initial_attention_mask: Optional[torch.Tensor] = None
2009 _is_batched_list = isinstance(input, list) and len(input) > 1
2011 if isinstance(input, (str, list)):
2012 input_type = "str"
2013 assert (
2014 self.tokenizer is not None
2015 ), "Must provide a tokenizer if passing a string to the model"
2016 if _is_batched_list:
2017 # Force left-padding for batched generation so real tokens
2018 # are flush-right and logits[:, -1, :] is always correct.
2019 input = self.to_tokens(input, prepend_bos=prepend_bos, padding_side="left")
2020 else:
2021 input = self.to_tokens(
2022 input, prepend_bos=prepend_bos, padding_side=padding_side
2023 )
2024 elif input.ndim == 2:
2025 input_type = "tokens"
2026 else:
2027 input_type = "embeds"
2029 input_tokens = input if input_type in ["str", "tokens"] else None
2030 batch_size, ctx_length = input.shape[0], input.shape[1]
2032 # Compute initial attention mask. For batched inputs with padding,
2033 # this correctly masks pad tokens. For single/unpadded inputs, this
2034 # is all-ones which matches the no-mask code path but ensures both
2035 # go through the same PosEmbed/attention logic for consistency.
2036 if input_tokens is not None and self.tokenizer is not None:
2037 _prepend_bos = (
2038 self.cfg.default_prepend_bos
2039 if prepend_bos is USE_DEFAULT_VALUE
2040 else (False if prepend_bos is None else prepend_bos)
2041 )
2042 # Temporarily set padding_side="left" so get_attention_mask
2043 # scans for leading pads (matching the left-padded tokens).
2044 _orig_padding_side = self.tokenizer.padding_side
2045 if _is_batched_list:
2046 self.tokenizer.padding_side = "left"
2047 initial_attention_mask = utils.get_attention_mask(
2048 self.tokenizer, input_tokens, _prepend_bos
2049 )
2050 if _is_batched_list:
2051 self.tokenizer.padding_side = _orig_padding_side
2052 device = get_device_for_block_index(0, self.cfg)
2053 input = input.to(device)
2054 if input_tokens is not None:
2055 # Re-alias to the moved tensor: input_tokens must live on the model's
2056 # device so later concatenations with sampled tokens (freq_penalty
2057 # sampling and the final output) don't mix devices.
2058 input_tokens = input
2059 if use_past_kv_cache:
2060 past_kv_cache = TransformerLensKeyValueCache.init_cache(
2061 self.cfg, self.cfg.device, batch_size
2062 )
2063 else:
2064 past_kv_cache = None
2066 # Only `output_logits` is supported from HF generation kwargs
2067 output_logits_flag = False
2068 if generation_kwargs:
2069 if "output_logits" in generation_kwargs:
2070 output_logits_flag = bool(generation_kwargs.pop("output_logits"))
2071 # Warn about unsupported keys
2072 accepted_keys = {"output_logits", "return_dict_in_generate"}
2073 unsupported_keys = [k for k in generation_kwargs.keys() if k not in accepted_keys]
2074 # Ignore `return_dict_in_generate`
2075 if "return_dict_in_generate" in generation_kwargs:
2076 generation_kwargs.pop("return_dict_in_generate")
2077 # Warn and drop unsupported keys
2078 if unsupported_keys:
2079 import warnings
2081 warnings.warn(
2082 f"HookedTransformer.generate received unsupported generation kwargs; ignoring: {unsupported_keys}",
2083 UserWarning,
2084 )
2085 # Remove unsupported keys
2086 for k in unsupported_keys:
2087 generation_kwargs.pop(k, None)
2089 # Collect per-step logits if requested
2090 logits_seq_list: Optional[List[torch.Tensor]] = [] if output_logits_flag else None
2092 shortformer_pos_embed = None
2093 embeds = input if input_type == "embeds" else self.embed(input)
2095 assert isinstance(embeds, torch.Tensor) and embeds.ndim == 3
2097 stop_tokens: List[int] = []
2098 eos_token_for_padding = 0
2099 if stop_at_eos:
2100 tokenizer_has_eos_token = (
2101 self.tokenizer is not None and self.tokenizer.eos_token_id is not None
2102 )
2103 if eos_token_id is None:
2104 assert (
2105 tokenizer_has_eos_token
2106 ), "Must pass a eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
2107 assert self.tokenizer is not None
2108 eos_token_id = self.tokenizer.eos_token_id
2110 if isinstance(eos_token_id, int): 2110 ↛ 2115line 2110 didn't jump to line 2115 because the condition on line 2110 was always true
2111 stop_tokens = [eos_token_id]
2112 eos_token_for_padding = eos_token_id
2113 else:
2114 # eos_token_id is a Sequence (e.g. list or tuple)
2115 stop_tokens = eos_token_id
2116 if tokenizer_has_eos_token:
2117 assert self.tokenizer is not None
2118 eos_token_for_padding = self.tokenizer.eos_token_id
2119 else:
2120 eos_token_for_padding = eos_token_id[0]
2122 # An array to track which sequences in the batch have finished.
2123 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
2125 # Currently nothing in HookedTransformer changes with eval, but this is here in case
2126 # that changes in the future.
2127 self.eval()
2128 sampled_tokens_list: List[torch.Tensor] = []
2129 for index in tqdm.tqdm(range(max_new_tokens), disable=not verbose):
2130 pos_offset = self.get_pos_offset(past_kv_cache, batch_size)
2132 # Extend the initial attention mask with 1s for generated tokens.
2133 attention_mask: Optional[torch.Tensor] = None
2134 if initial_attention_mask is not None:
2135 n_new = len(sampled_tokens_list)
2136 if n_new > 0:
2137 ones = torch.ones(
2138 batch_size,
2139 n_new,
2140 dtype=initial_attention_mask.dtype,
2141 device=device,
2142 )
2143 attention_mask = torch.cat([initial_attention_mask.to(device), ones], dim=1)
2144 else:
2145 attention_mask = initial_attention_mask.to(device)
2146 residual, shortformer_pos_embed = self.get_residual(
2147 embeds,
2148 pos_offset,
2149 return_shortformer_pos_embed=True,
2150 device=device,
2151 attention_mask=attention_mask,
2152 )
2154 # While generating, we keep generating logits, throw away all but the final logits,
2155 # and then use those logits to sample from the distribution We keep adding the
2156 # sampled tokens to the end of tokens.
2157 start_at_layer = 0 # Make forward returns embeddings
2158 if use_past_kv_cache:
2159 # We just take the final tokens, as a [batch, 1] tensor
2160 if index > 0:
2161 logits = self.forward(
2162 residual[:, -1:],
2163 return_type="logits",
2164 prepend_bos=prepend_bos,
2165 padding_side=padding_side,
2166 past_kv_cache=past_kv_cache,
2167 start_at_layer=start_at_layer,
2168 shortformer_pos_embed=shortformer_pos_embed,
2169 attention_mask=attention_mask,
2170 )
2171 else:
2172 logits = self.forward(
2173 residual,
2174 return_type="logits",
2175 prepend_bos=prepend_bos,
2176 padding_side=padding_side,
2177 past_kv_cache=past_kv_cache,
2178 start_at_layer=start_at_layer,
2179 shortformer_pos_embed=shortformer_pos_embed,
2180 attention_mask=attention_mask,
2181 )
2182 else:
2183 # We input the entire sequence, as a [batch, pos] tensor, since we aren't using
2184 # the cache.
2185 logits = self.forward(
2186 residual,
2187 return_type="logits",
2188 prepend_bos=prepend_bos,
2189 padding_side=padding_side,
2190 start_at_layer=start_at_layer,
2191 shortformer_pos_embed=shortformer_pos_embed,
2192 attention_mask=attention_mask,
2193 )
2194 final_logits = logits[:, -1, :]
2196 if output_logits_flag:
2197 assert logits_seq_list is not None
2198 logits_seq_list.append(final_logits.clone())
2200 if do_sample:
2201 if input_type in [
2202 "str",
2203 "tokens",
2204 ]: # Those types of inputs support frequency penalty
2205 assert input_tokens is not None
2206 sampled_tokens = utils.sample_logits(
2207 final_logits,
2208 top_k=top_k,
2209 top_p=top_p,
2210 temperature=temperature,
2211 freq_penalty=freq_penalty,
2212 tokens=torch.cat(
2213 (input_tokens, torch.cat(sampled_tokens_list, dim=1)), dim=1
2214 )
2215 if "sampled_tokens" in locals()
2216 else input_tokens,
2217 ).to(get_device_for_block_index(0, self.cfg))
2218 else:
2219 sampled_tokens = utils.sample_logits(
2220 final_logits, top_k=top_k, top_p=top_p, temperature=temperature
2221 ).to(get_device_for_block_index(0, self.cfg))
2222 else:
2223 sampled_tokens = final_logits.argmax(-1).to(
2224 get_device_for_block_index(0, self.cfg)
2225 )
2226 sampled_tokens_list.append(sampled_tokens.unsqueeze(1))
2227 if stop_at_eos:
2228 # For all unfinished sequences, add on the next token. If a sequence was
2229 # finished, throw away the generated token and add eos_token_for_padding
2230 # instead.
2231 sampled_tokens[finished_sequences] = eos_token_for_padding
2232 finished_sequences.logical_or_(
2233 torch.isin(
2234 sampled_tokens.to(self.cfg.device),
2235 torch.tensor(stop_tokens).to(self.cfg.device),
2236 )
2237 )
2239 embeds = torch.hstack([embeds, self.embed(sampled_tokens.unsqueeze(-1))])
2241 if stop_at_eos and finished_sequences.all(): 2241 ↛ 2242line 2241 didn't jump to line 2242 because the condition on line 2241 was never true
2242 break
2244 sampled_tokens = torch.cat(sampled_tokens_list, dim=1)
2245 if input_type in ["str", "tokens"]:
2246 assert input_tokens is not None
2247 output_tokens = torch.cat((input_tokens, sampled_tokens), dim=1)
2248 else:
2249 output_tokens = sampled_tokens
2251 if return_type == "str":
2252 assert self.tokenizer is not None
2253 decoded_texts: List[str] = [
2254 cast(str, self.tokenizer.decode(tokens, skip_special_tokens=True))
2255 for tokens in output_tokens
2256 ]
2257 result: Any = decoded_texts[0] if len(decoded_texts) == 1 else decoded_texts
2258 elif return_type == "tokens":
2259 result = cast(Any, output_tokens)
2260 else:
2261 result = cast(Any, embeds)
2263 if output_logits_flag:
2264 # Return HF ModelOutput format
2265 from transformers.utils import ModelOutput # type: ignore
2267 def _logits_to_tuple(logits_list: list[torch.Tensor]) -> tuple[torch.Tensor, ...]:
2268 assert logits_list is not None
2269 return tuple(logits_list)
2271 try:
2272 from transformers.generation.utils import GenerateDecoderOnlyOutput
2274 return GenerateDecoderOnlyOutput(
2275 sequences=cast(torch.LongTensor, output_tokens),
2276 # HF's type hint tuple[FloatTensor] is really tuple[FloatTensor, ...]
2277 logits=_logits_to_tuple(logits_seq_list), # type: ignore[arg-type]
2278 )
2279 except (ImportError, AttributeError):
2280 # Fallback for older transformers versions
2281 # `sequences` expects a tensor of token ids
2282 return ModelOutput(sequences=output_tokens, logits=_logits_to_tuple(logits_seq_list)) # type: ignore[arg-type]
2283 else:
2284 return result
2286 @torch.inference_mode()
2287 def generate_stream(
2288 self,
2289 input: Union[str, Float[torch.Tensor, "batch pos"]] = "",
2290 max_new_tokens: int = 10,
2291 max_tokens_per_yield: int = 25,
2292 stop_at_eos: bool = True,
2293 eos_token_id: Optional[int] = None,
2294 do_sample: bool = True,
2295 top_k: Optional[int] = None,
2296 top_p: Optional[float] = None,
2297 temperature: float = 1.0,
2298 freq_penalty: float = 0.0,
2299 use_past_kv_cache: bool = True,
2300 prepend_bos: Optional[bool] = USE_DEFAULT_VALUE,
2301 padding_side: Optional[Literal["left", "right"]] = USE_DEFAULT_VALUE,
2302 return_type: Optional[str] = "input",
2303 verbose: bool = True,
2304 ) -> Generator[Union[Int[torch.Tensor, "batch"], str], None, None]:
2305 """Stream tokens from the Model as they are generated.
2307 Sample tokens from the model until the model outputs eos_token or max_new_tokens is reached,
2308 yielding batches of tokens progressively during generation rather than waiting for the entire
2309 sequence to be generated.
2311 To avoid fiddling with ragged tensors, if we input a batch of text and some sequences finish
2312 (by producing an EOT token), we keep running the model on the entire batch, but throw away
2313 the output for a finished sequence and just keep adding EOTs to pad.
2315 This supports entering a single string, but not a list of strings - if the strings don't
2316 tokenize to exactly the same length, this gets messy. If that functionality is needed,
2317 convert them to a batch of tokens and input that instead.
2319 Args:
2320 input (Union[str, Int[torch.Tensor, "batch pos"])]): Either a batch of tokens ([batch,
2321 pos]) or a text string (this will be converted to a batch of tokens with batch size
2322 1).
2323 max_new_tokens (int): Maximum number of tokens to generate.
2324 max_tokens_per_yield (int): Maximum number of tokens to accumulate before yielding.
2325 Controls how frequently the function yields tokens during generation.
2326 stop_at_eos (bool): If True, stop generating tokens when the model outputs eos_token.
2327 eos_token_id (Optional[Union[int, Sequence]]): The token ID to use for end
2328 of sentence. If None, use the tokenizer's eos_token_id - required if using
2329 stop_at_eos. It's also possible to provide a list of token IDs (not just the
2330 eos_token_id), in which case the generation will stop when any of them are output
2331 (useful e.g. for stable_lm).
2332 do_sample (bool): If True, sample from the model's output distribution. Otherwise, use
2333 greedy search (take the max logit each time).
2334 top_k (int): Number of tokens to sample from. If None, sample from all tokens.
2335 top_p (float): Probability mass to sample from. If 1.0, sample from all tokens. If <1.0,
2336 we take the top tokens with cumulative probability >= top_p.
2337 temperature (float): Temperature for sampling. Higher values will make the model more
2338 random (limit of temp -> 0 is just taking the top token, limit of temp -> inf is
2339 sampling from a uniform distribution).
2340 freq_penalty (float): Frequency penalty for sampling - how much to penalise previous
2341 tokens. Higher values will make the model more random.
2342 use_past_kv_cache (bool): If True, create and use cache to speed up generation.
2343 prepend_bos (bool, optional): Overrides self.cfg.default_prepend_bos. Whether to prepend
2344 the BOS token to the input (applicable when input is a string). Defaults to None,
2345 implying usage of self.cfg.default_prepend_bos (default is True unless specified
2346 otherwise). Pass True or False to override the default.
2347 padding_side (Union[Literal["left", "right"], None], optional): Overrides
2348 self.tokenizer.padding_side. Specifies which side to pad when tokenizing multiple
2349 strings of different lengths.
2350 return_type (Optional[str]): The type of the output to return - either a string (str),
2351 a tensor of tokens (tensor) or whatever the format of the input was (input).
2352 verbose (bool): If True, show tqdm progress bars for generation.
2354 Yields:
2355 outputs (Union[Int[torch.Tensor, "batch"], str]): Batches of generated tokens, yielded
2356 progressively during generation. Each yield contains accumulated tokens since the last
2357 yield, up to max_tokens_per_yield.
2358 """
2360 with utils.LocallyOverridenDefaults(
2361 self, prepend_bos=prepend_bos, padding_side=padding_side
2362 ):
2363 if type(input) == str: 2363 ↛ 2370line 2363 didn't jump to line 2370 because the condition on line 2363 was always true
2364 # If text, convert to tokens (batch_size=1)
2365 assert (
2366 self.tokenizer is not None
2367 ), "Must provide a tokenizer if passing a string to the model"
2368 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
2369 else:
2370 assert isinstance(input, torch.Tensor), "Input must be a tensor when not a string"
2371 tokens = input
2373 if return_type == "input": 2373 ↛ 2379line 2373 didn't jump to line 2379 because the condition on line 2373 was always true
2374 if type(input) == str: 2374 ↛ 2377line 2374 didn't jump to line 2377 because the condition on line 2374 was always true
2375 return_type = "str"
2376 else:
2377 return_type = "tensor"
2379 assert isinstance(tokens, torch.Tensor)
2380 batch_size, ctx_length = tokens.shape
2381 device = get_device_for_block_index(0, self.cfg)
2382 tokens = tokens.to(device)
2383 if use_past_kv_cache: 2383 ↛ 2388line 2383 didn't jump to line 2388 because the condition on line 2383 was always true
2384 past_kv_cache = TransformerLensKeyValueCache.init_cache(
2385 self.cfg, self.cfg.device, batch_size
2386 )
2387 else:
2388 past_kv_cache = None
2390 stop_tokens: List[int] = []
2391 eos_token_for_padding = 0
2392 if stop_at_eos: 2392 ↛ 2416line 2392 didn't jump to line 2416 because the condition on line 2392 was always true
2393 tokenizer_has_eos_token = (
2394 self.tokenizer is not None and self.tokenizer.eos_token_id is not None
2395 )
2396 if eos_token_id is None: 2396 ↛ 2403line 2396 didn't jump to line 2403 because the condition on line 2396 was always true
2397 assert (
2398 tokenizer_has_eos_token
2399 ), "Must pass a eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
2400 assert self.tokenizer is not None
2401 eos_token_id = self.tokenizer.eos_token_id
2403 if isinstance(eos_token_id, int): 2403 ↛ 2408line 2403 didn't jump to line 2408 because the condition on line 2403 was always true
2404 stop_tokens = [eos_token_id]
2405 eos_token_for_padding = eos_token_id
2406 else:
2407 # eos_token_id is a Sequence (e.g. list or tuple)
2408 stop_tokens = eos_token_id
2409 if tokenizer_has_eos_token:
2410 assert self.tokenizer is not None
2411 eos_token_for_padding = self.tokenizer.eos_token_id
2412 else:
2413 eos_token_for_padding = eos_token_id[0]
2415 # An array to track which sequences in the batch have finished.
2416 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
2418 accumulated_tokens: Optional[torch.Tensor] = None
2419 tokens_since_last_yield = 0
2421 # Currently nothing in HookedTransformer changes with eval, but this is here in case
2422 # that changes in the future.
2423 self.eval()
2424 for index in tqdm.tqdm(range(max_new_tokens), disable=not verbose):
2425 # While generating, we keep generating logits, throw away all but the final logits,
2426 # and then use those logits to sample from the distribution We keep adding the
2427 # sampled tokens to the end of tokens.
2428 if use_past_kv_cache: 2428 ↛ 2449line 2428 didn't jump to line 2449 because the condition on line 2428 was always true
2429 # We just take the final tokens, as a [batch, 1] tensor
2430 if index > 0:
2431 logits = self.forward(
2432 tokens[:, -1:],
2433 return_type="logits",
2434 prepend_bos=prepend_bos,
2435 padding_side=padding_side,
2436 past_kv_cache=past_kv_cache,
2437 )
2438 else:
2439 logits = self.forward(
2440 tokens,
2441 return_type="logits",
2442 prepend_bos=prepend_bos,
2443 padding_side=padding_side,
2444 past_kv_cache=past_kv_cache,
2445 )
2446 else:
2447 # We input the entire sequence, as a [batch, pos] tensor, since we aren't using
2448 # the cache.
2449 logits = self.forward(
2450 tokens,
2451 return_type="logits",
2452 prepend_bos=prepend_bos,
2453 padding_side=padding_side,
2454 )
2455 final_logits = logits[:, -1, :]
2457 if do_sample: 2457 ↛ 2458line 2457 didn't jump to line 2458 because the condition on line 2457 was never true
2458 sampled_tokens = utils.sample_logits(
2459 final_logits,
2460 top_k=top_k,
2461 top_p=top_p,
2462 temperature=temperature,
2463 freq_penalty=freq_penalty,
2464 tokens=tokens,
2465 ).to(get_device_for_block_index(0, self.cfg))
2466 else:
2467 sampled_tokens = final_logits.argmax(-1).to(
2468 get_device_for_block_index(0, self.cfg)
2469 )
2471 if stop_at_eos: 2471 ↛ 2483line 2471 didn't jump to line 2483 because the condition on line 2471 was always true
2472 # For all unfinished sequences, add on the next token. If a sequence was
2473 # finished, throw away the generated token and add eos_token_for_padding
2474 # instead.
2475 sampled_tokens[finished_sequences] = eos_token_for_padding
2476 finished_sequences.logical_or_(
2477 torch.isin(
2478 sampled_tokens.to(self.cfg.device),
2479 torch.tensor(stop_tokens).to(self.cfg.device),
2480 )
2481 )
2483 new_tokens = sampled_tokens.unsqueeze(-1)
2485 # Accumulate tokens until we hit max_tokens_per_yield
2486 if index == 0:
2487 accumulated_tokens = torch.cat([tokens, new_tokens], dim=-1)
2488 tokens_since_last_yield = accumulated_tokens.shape[1]
2489 else:
2490 if accumulated_tokens is None:
2491 accumulated_tokens = new_tokens
2492 else:
2493 accumulated_tokens = torch.cat([accumulated_tokens, new_tokens], dim=-1)
2494 tokens_since_last_yield += 1
2496 if tokens_since_last_yield >= max_tokens_per_yield:
2497 yield accumulated_tokens
2498 tokens_since_last_yield = 0
2499 accumulated_tokens = None
2501 tokens = torch.cat([tokens, new_tokens], dim=-1)
2503 if stop_at_eos and finished_sequences.all(): 2503 ↛ 2505line 2503 didn't jump to line 2505 because the condition on line 2503 was never true
2504 # Yield any remaining accumulated tokens before breaking
2505 if accumulated_tokens is not None:
2506 yield accumulated_tokens
2507 break
2509 # Only yield remaining tokens if we didn't already yield them in the break case
2510 if accumulated_tokens is not None and not (stop_at_eos and finished_sequences.all()): 2510 ↛ exitline 2510 didn't jump to the function exit
2511 yield accumulated_tokens
2513 @property
2514 def n_params_total(self) -> int:
2515 """Total number of parameters in the model, including embeddings, biases,
2516 and layer norm weights.
2518 This complements ``self.cfg.n_params``, which counts only the "hidden
2519 weight" parameters (attention projections + MLP weights, excluding
2520 embeddings/biases/layer norms) following the
2521 `scaling laws paper <https://arxiv.org/pdf/2001.08361.pdf>`_ convention.
2523 Use this when you want the actual parameter count for memory budgeting,
2524 comparison with HuggingFace's ``model.num_parameters()``, or alignment
2525 with reported model sizes in papers (e.g. the Pythia suite).
2527 Returns:
2528 int: ``sum(p.numel() for p in self.parameters())``
2529 """
2530 return sum(p.numel() for p in self.parameters())
2532 # Give access to all weights as properties.
2533 @property
2534 def W_U(self) -> Float[torch.Tensor, "d_model d_vocab"]:
2535 """Convenience to get the unembedding matrix.
2537 I.e. the linear map from the final residual stream to the output logits).
2538 """
2539 return self.unembed.W_U
2541 @property
2542 def b_U(self) -> Float[torch.Tensor, "d_vocab"]:
2543 return self.unembed.b_U
2545 @property
2546 def W_E(self) -> Float[torch.Tensor, "d_vocab d_model"]:
2547 """Convenience to get the embedding matrix."""
2548 return self.embed.W_E
2550 @property
2551 def W_pos(self) -> Float[torch.Tensor, "n_ctx d_model"]:
2552 """Convenience function to get the positional embedding.
2554 Only works on models with absolute positional embeddings!
2555 """
2556 return self.pos_embed.W_pos
2558 @property
2559 def W_E_pos(self) -> Float[torch.Tensor, "d_vocab+n_ctx d_model"]:
2560 """Concatenated W_E and W_pos.
2562 Used as a full (overcomplete) basis of the input space, useful for full QK and full OV
2563 circuits.
2564 """
2565 return torch.cat([self.W_E, self.W_pos], dim=0)
2567 # Layer-specific weights are stacked into one massive tensor and given as properties for
2568 # convenience and a cache is used to avoid repeated computation. Often a useful convenience when
2569 # we want to do analysis on weights across all layers. If GPU memory is a bottleneck, don't use
2570 # these properties!
2572 @property
2573 def W_K(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]:
2574 """Stack the key weights across all layers."""
2575 return torch.stack([block.attn.W_K for block in self.blocks], dim=0)
2577 @property
2578 def W_Q(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]:
2579 """Stack the query weights across all layers."""
2580 return torch.stack([block.attn.W_Q for block in self.blocks], dim=0)
2582 @property
2583 def W_V(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]:
2584 """Stack the value weights across all layers."""
2585 return torch.stack([block.attn.W_V for block in self.blocks], dim=0)
2587 @property
2588 def W_O(self) -> Float[torch.Tensor, "n_layers n_heads d_head d_model"]:
2589 """Stack the attn output weights across all layers."""
2590 return torch.stack([block.attn.W_O for block in self.blocks], dim=0)
2592 @property
2593 def W_in(self) -> Float[torch.Tensor, "n_layers d_model d_mlp"]:
2594 """Stack the MLP input weights across all layers."""
2595 return torch.stack(
2596 [cast(Union[MLP, GatedMLP], block.mlp).W_in for block in self.blocks], dim=0
2597 )
2599 @property
2600 def W_gate(self) -> Union[Float[torch.Tensor, "n_layers d_model d_mlp"], None]:
2601 """Stack the MLP gate weights across all layers.
2603 Only works for models with gated MLPs.
2604 """
2605 if self.cfg.gated_mlp:
2606 return torch.stack([cast(GatedMLP, block.mlp).W_gate for block in self.blocks], dim=0)
2607 else:
2608 return None
2610 @property
2611 def W_out(self) -> Float[torch.Tensor, "n_layers d_mlp d_model"]:
2612 """Stack the MLP output weights across all layers."""
2613 return torch.stack(
2614 [cast(Union[MLP, GatedMLP], block.mlp).W_out for block in self.blocks], dim=0
2615 )
2617 @property
2618 def b_K(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]:
2619 """Stack the key biases across all layers."""
2620 return torch.stack([block.attn.b_K for block in self.blocks], dim=0)
2622 @property
2623 def b_Q(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]:
2624 """Stack the query biases across all layers."""
2625 return torch.stack([block.attn.b_Q for block in self.blocks], dim=0)
2627 @property
2628 def b_V(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]:
2629 """Stack the value biases across all layers."""
2630 return torch.stack([block.attn.b_V for block in self.blocks], dim=0)
2632 @property
2633 def b_O(self) -> Float[torch.Tensor, "n_layers d_model"]:
2634 """Stack the attn output biases across all layers."""
2635 return torch.stack([block.attn.b_O for block in self.blocks], dim=0)
2637 @property
2638 def b_in(self) -> Float[torch.Tensor, "n_layers d_mlp"]:
2639 """Stack the MLP input biases across all layers."""
2640 return torch.stack(
2641 [cast(Union[MLP, GatedMLP], block.mlp).b_in for block in self.blocks], dim=0
2642 )
2644 @property
2645 def b_out(self) -> Float[torch.Tensor, "n_layers d_model"]:
2646 """Stack the MLP output biases across all layers."""
2647 return torch.stack(
2648 [cast(Union[MLP, GatedMLP], block.mlp).b_out for block in self.blocks], dim=0
2649 )
2651 @property
2652 def QK(self):
2653 return FactoredMatrix(self.W_Q, self.W_K.transpose(-2, -1))
2655 @property
2656 def OV(self):
2657 return FactoredMatrix(self.W_V, self.W_O)
2659 # Various utility functions
2660 def accumulated_bias(
2661 self, layer: int, mlp_input: bool = False, include_mlp_biases=True
2662 ) -> Float[torch.Tensor, "d_model"]:
2663 """Accumulated Bias.
2665 Returns the accumulated bias from all layer outputs (ie the b_Os and b_outs), up to the
2666 input of layer L.
2668 Args:
2669 layer (int): Layer number, in [0, n_layers]. layer==0 means no layers, layer==n_layers
2670 means all layers.
2671 mlp_input (bool): If True, we take the bias up to the input of the MLP
2672 of layer L (ie we include the bias from the attention output of the current layer,
2673 otherwise just biases from previous layers)
2674 include_mlp_biases (bool): Whether to include the biases of MLP layers. Often useful to
2675 have as False if we're expanding attn_out into individual heads, but keeping mlp_out
2676 as is.
2678 Returns:
2679 bias (torch.Tensor): [d_model], accumulated bias
2680 """
2681 accumulated_bias = torch.zeros(self.cfg.d_model, device=self.cfg.device)
2683 for i in range(layer):
2684 block = self.blocks[i]
2685 accumulated_bias += cast(torch.Tensor, block.attn.b_O)
2686 if include_mlp_biases:
2687 accumulated_bias += cast(torch.Tensor, block.mlp.b_out)
2688 if mlp_input:
2689 assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer"
2690 block = self.blocks[layer]
2691 accumulated_bias += cast(torch.Tensor, block.attn.b_O)
2692 return accumulated_bias
2694 def all_composition_scores(
2695 self, mode
2696 ) -> Float[torch.Tensor, "n_layers n_heads n_layers n_heads"]:
2697 """All Composition Scores.
2699 Returns the Composition scores for all pairs of heads, as a L1, H1, L2, H2 tensor (which is
2700 upper triangular on the first and third axes).
2702 See
2703 https://transformer-circuits.pub/2021/framework/index.html#:~:text=The%20above%20diagram%20shows%20Q%2D%2C%20K%2D%2C%20and%20V%2DComposition
2704 for three metrics used.
2706 Args:
2707 mode (str): One of ["Q", "K", "V"], the mode to use for the composition score.
2708 """
2709 left = self.OV
2710 if mode == "Q":
2711 right = self.QK
2712 elif mode == "K":
2713 right = self.QK.T
2714 elif mode == "V":
2715 right = self.OV
2716 else:
2717 raise ValueError(f"mode must be one of ['Q', 'K', 'V'] not {mode}")
2719 scores = utils.composition_scores(left, right, broadcast_dims=True)
2720 # Mask scores to be zero for all pairs with the right head in the same layer or earlier
2721 # layer than the left head.
2722 mask = (
2723 torch.arange(self.cfg.n_layers, device=self.cfg.device)[:, None, None, None]
2724 < torch.arange(self.cfg.n_layers, device=self.cfg.device)[None, None, :, None]
2725 )
2726 scores = torch.where(mask, scores, torch.zeros_like(scores))
2727 return scores
2729 def all_head_labels(self):
2730 """Returns a list of all head names in the model."""
2731 return [f"L{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)]
2733 def load_sample_training_dataset(self, **kwargs):
2734 """Load Sample Training Dataset.
2736 Helper function to load in a 10K-20K dataset of elements from the model's training data
2737 distribution.
2739 Wrapper around utils.get_dataset, which identifies the appropriate dataset the pretrained
2740 models. Each dataset has a 'text' field, which contains the relevant info, some have several
2741 meta data fields.
2743 Kwargs will be passed to utils.get_dataset (e.g. cache_dir to set download location)
2745 Notes:
2747 - PT-2's training data is not open source. OpenWebText is a replication (links with
2748 >3 karma on Reddit)
2749 - OPT's training data is not open source, and is a mess of different things that is hard to
2750 replicate. I default to the Pile, which covers some of it, but imperfectly.
2752 (Some models will have actually been trained on the data supplied here, for some it's from
2753 the validation set).
2754 """
2755 model_dataset_map = {
2756 "neel": "c4_code",
2757 "neel-solu-old": "pile",
2758 "GPT2LMHeadModel": "openwebtext",
2759 "GPTNeoForCausalLM": "pile",
2760 "GPTNeoXForCausalLM": "pile",
2761 "GPTJForCausalLM": "pile",
2762 "OPTForCausalLM": "pile",
2763 }
2764 if self.cfg.original_architecture in model_dataset_map:
2765 self.dataset = utils.get_dataset(
2766 model_dataset_map[self.cfg.original_architecture], **kwargs
2767 )
2768 else:
2769 raise ValueError(
2770 f"We do not have an available dataset for the relevant model: {self.cfg.original_architecture}"
2771 )
2772 return self.dataset
2774 def sample_datapoint(
2775 self,
2776 tokenize: bool = False,
2777 prepend_bos: Optional[Union[bool, None]] = USE_DEFAULT_VALUE,
2778 padding_side: Optional[Literal["left", "right"]] = USE_DEFAULT_VALUE,
2779 ) -> Union[str, Float[torch.Tensor, "1 pos"]]:
2780 """Sample Data Point from Dataset.
2782 Helper function to randomly sample a data point from self.dataset, a small dataset from the
2783 data distribution the model was trained on.
2785 Implicitly calls self.load_sample_training_dataset if it hasn't already been called. Only
2786 works for pretrained models with an associated dataset. But you can manually replace
2787 self.dataset with a dataset of your choice if you want.
2789 Args:
2790 tokenize (bool): Whether to return tokens (instead of text). Defaults to False. Note
2791 that the returned tokens will be automatically truncated to the model's max context
2792 size.
2793 prepend_bos (bool, optional): Overrides self.cfg.default_prepend_bos. Whether to prepend
2794 the BOS token to the input (applicable when input is a string). Defaults to None,
2795 implying usage of self.cfg.default_prepend_bos (default is True unless specified
2796 otherwise). Pass True or False to override the default.
2797 padding_side (Union[Literal["left", "right"], None], optional): Overrides
2798 self.tokenizer.padding_side. Specifies which side to pad when tokenizing multiple
2799 strings of different lengths.
2800 """
2801 if self.dataset is None:
2802 self.load_sample_training_dataset()
2803 assert self.dataset is not None # keep mypy happy
2804 sample_dataset_size = len(self.dataset)
2805 index = np.random.randint(0, sample_dataset_size)
2806 if not tokenize:
2807 return self.dataset[index]["text"]
2808 else:
2809 return self.to_tokens(
2810 self.dataset[index]["text"],
2811 prepend_bos=prepend_bos,
2812 padding_side=padding_side,
2813 truncate=True,
2814 )