Coverage for transformer_lens/HookedEncoderDecoder.py: 70%
277 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Hooked EncoderDecoder
3Contains a T5 style model. This is separate from :class:`transformer_lens.HookedTransformer`
4because it has a significantly different architecture to e.g. GPT style transformers.
5"""
7from __future__ import annotations
9import logging
10import os
11from itertools import chain
12from pathlib import Path
13from typing import (
14 Any,
15 Dict,
16 List,
17 Optional,
18 Tuple,
19 Type,
20 TypeVar,
21 Union,
22 cast,
23 overload,
24)
26import torch
27import tqdm
28from einops import repeat
29from jaxtyping import Float, Int
30from transformers import AutoTokenizer, PreTrainedTokenizerBase
31from typing_extensions import Literal
33import transformer_lens.loading_from_pretrained as loading
34from transformer_lens.ActivationCache import ActivationCache
35from transformer_lens.components import MLP, Embed, GatedMLP, RMSNorm, T5Block, Unembed
36from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig
37from transformer_lens.FactoredMatrix import FactoredMatrix
38from transformer_lens.hook_points import HookPoint
39from transformer_lens.HookedRootModule import HookedRootModule
40from transformer_lens.utilities import TypedModuleList, sample_logits, warn_if_mps
41from transformer_lens.utilities.multi_gpu import get_device_for_block_index
43T = TypeVar("T", bound="HookedEncoderDecoder")
46class HookedEncoderDecoder(HookedRootModule):
47 """
48 This class implements a T5 encoder-decoder using the components in ./components.py, with HookPoints on every interesting activation. It inherits from HookedRootModule.
50 Limitations:
51 - Also note that model does not include dropouts, which may lead to inconsistent results from training or fine-tuning.
53 Like HookedTransformer, it can have a pretrained Transformer's weights loaded via `.from_pretrained`. There are a few features you might know from HookedTransformer which are not yet supported:
54 - There is no preprocessing (e.g. LayerNorm folding) when loading a pretrained model
55 - The model only accepts tokens as inputs, and not strings, or lists of strings
56 """
58 tokenizer: Optional[PreTrainedTokenizerBase]
59 encoder: TypedModuleList[T5Block]
60 decoder: TypedModuleList[T5Block]
62 def __init__(
63 self,
64 cfg: Union[HookedTransformerConfig, Dict],
65 tokenizer: Optional[PreTrainedTokenizerBase] = None,
66 move_to_device: bool = True,
67 **kwargs: Any,
68 ):
69 super().__init__()
70 if isinstance(cfg, Dict): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 cfg = HookedTransformerConfig(**cfg)
72 elif isinstance(cfg, str): 72 ↛ 73line 72 didn't jump to line 73 because the condition on line 72 was never true
73 raise ValueError(
74 "Please pass in a config dictionary or HookedTransformerConfig object. If you want to load a pretrained model, use HookedEncoderDecoder.from_pretrained() instead."
75 )
76 self.cfg: HookedTransformerConfig = cfg
78 if self.cfg.n_devices != 1: 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true
79 raise ValueError("Multiple devices not supported for HookedEncoderDecoder")
80 if tokenizer is not None: 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true
81 self.tokenizer = tokenizer
82 elif self.cfg.tokenizer_name is not None: 82 ↛ 89line 82 didn't jump to line 89 because the condition on line 82 was always true
83 huggingface_token = os.environ.get("HF_TOKEN", "")
84 self.tokenizer = AutoTokenizer.from_pretrained(
85 self.cfg.tokenizer_name,
86 token=huggingface_token if len(huggingface_token) > 0 else None,
87 )
88 else:
89 self.tokenizer = None
91 if self.cfg.d_vocab == -1: 91 ↛ 93line 91 didn't jump to line 93 because the condition on line 91 was never true
92 # If we have a tokenizer, vocab size can be inferred from it.
93 if self.tokenizer is None:
94 raise ValueError("Must provide a tokenizer if d_vocab is not provided")
96 self.cfg.d_vocab = len(self.tokenizer)
97 if self.cfg.d_vocab_out == -1: 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true
98 self.cfg.d_vocab_out = self.cfg.d_vocab
100 self.embed = Embed(self.cfg)
101 self.encoder = TypedModuleList(
102 [
103 T5Block(self.cfg, num_layer, is_decoder=False)
104 for num_layer in range(self.cfg.n_layers)
105 ]
106 )
107 self.encoder_final_ln = RMSNorm(self.cfg)
108 self.decoder = TypedModuleList(
109 [
110 T5Block(self.cfg, num_layer, is_decoder=True)
111 for num_layer in range(self.cfg.n_layers)
112 ]
113 )
114 self.decoder_final_ln = RMSNorm(self.cfg)
115 self.unembed = Unembed(self.cfg)
117 self.hook_embed = HookPoint()
119 if move_to_device and self.cfg.device is not None: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true
120 self.to(self.cfg.device)
122 self.setup()
124 def to_tokens(
125 self,
126 input: Union[str, List[str]],
127 move_to_device: bool = True,
128 truncate: bool = True,
129 ) -> Tuple[Int[torch.Tensor, "batch pos"], Int[torch.Tensor, "batch pos"]]:
130 """Converts a string to a tensor of tokens.
131 Taken mostly from the HookedTransformer implementation, but does not support default padding
132 sides or prepend_bos.
134 Args:
135 input (Union[str, List[str]]): The input to tokenize.
136 move_to_device (bool): Whether to move the output tensor of tokens to the device the
137 model lives on. Defaults to True
138 truncate (bool): If the output tokens are too long, whether to truncate the output
139 tokens to the model's max context window. Does nothing for shorter inputs.
140 Defaults to True.
141 """
143 assert self.tokenizer is not None, "Cannot use to_tokens without a tokenizer"
145 encodings = self.tokenizer(
146 input,
147 return_tensors="pt",
148 padding=True,
149 truncation=truncate,
150 max_length=self.cfg.n_ctx if truncate else None,
151 )
153 tokens = encodings.input_ids
154 attention_mask = encodings.attention_mask
156 if move_to_device: 156 ↛ 159line 156 didn't jump to line 159 because the condition on line 156 was always true
157 tokens = tokens.to(self.cfg.device)
158 attention_mask = attention_mask.to(self.cfg.device)
159 return tokens, attention_mask
161 @overload
162 def forward(
163 self,
164 input: Union[str, List[str], Int[torch.Tensor, "batch pos"]],
165 decoder_input: Optional[Int[torch.Tensor, "batch decoder_pos"]] = None,
166 return_type: Literal["logits"] = "logits",
167 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None,
168 ) -> Float[torch.Tensor, "batch pos d_vocab"]:
169 ...
171 @overload
172 def forward(
173 self,
174 input: Union[str, List[str], Int[torch.Tensor, "batch pos"]],
175 decoder_input: Optional[Int[torch.Tensor, "batch decoder_pos"]] = None,
176 return_type: Optional[Literal[None]] = None,
177 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None,
178 ) -> Optional[Float[torch.Tensor, "batch pos d_vocab"]]:
179 ...
181 def forward(
182 self,
183 input: Union[str, List[str], Int[torch.Tensor, "batch pos"]],
184 decoder_input: Optional[Int[torch.Tensor, "batch decoder_pos"]] = None,
185 return_type: Optional[str] = "logits",
186 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None,
187 ) -> Optional[Float[torch.Tensor, "batch decoder_pos d_vocab"]]:
188 """Forward pass of the T5 model.
190 Args:
191 input: Input to be processed. Can be one of:
192 - str: A single string input
193 - List[str]: A batch of string inputs
194 - Int[torch.Tensor, "batch pos"]: A batch of token IDs
195 decoder_input: Tensor of shape (batch, decoder_pos) containing the decoder input sequence.
196 If None and input is of type str or List[str], starts with batch of beginning-of-sequence (BOS) tokens.
197 return_type: Specifies the model output type:
198 - "logits": Return logits tensor
199 - None: Returns nothing
200 one_zero_attention_mask: A binary mask which indicates
201 which tokens should be attended to (1) and which should be ignored (0).
202 Primarily used for padding variable-length sentences in a batch.
203 For instance, in a batch with sentences of differing lengths, shorter
204 sentences are padded with 0s on the right. If not provided, the model
205 assumes all tokens should be attended to.
206 This parameter gets inferred from the tokenizer if input is a string or list of strings.
207 Shape is (batch_size, sequence_length).
209 Returns:
210 Optional[Float[torch.Tensor, "batch decoder_pos d_vocab"]]:
211 If return_type="logits": Returns logits tensor of shape (batch, decoder_pos, vocab_size)
212 If return_type=None: Returns None
213 """
215 if isinstance(input, (str, list)):
216 tokens, attention_mask = self.to_tokens(input)
218 # If attention mask is not provided, use the ones from the tokenizer
219 one_zero_attention_mask = (
220 attention_mask if one_zero_attention_mask is None else one_zero_attention_mask
221 )
223 # If decoder_input is not provided, start with tensor of PAD tokens of shape (batch, 1)
224 if decoder_input is None: 224 ↛ 244line 224 didn't jump to line 244 because the condition on line 224 was always true
225 assert self.tokenizer is not None
226 decoder_input = torch.full(
227 (tokens.shape[0], 1),
228 self.tokenizer.pad_token_id,
229 device=self.cfg.device,
230 )
231 else:
232 tokens = input
234 if one_zero_attention_mask is None:
235 logging.warning(
236 "No attention mask provided. Assuming all tokens should be attended to."
237 )
239 if decoder_input is None: 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true
240 raise ValueError(
241 "Must provide decoder_input if input is not a string or list of strings"
242 )
244 if tokens.device.type != self.cfg.device: 244 ↛ 245line 244 didn't jump to line 245 because the condition on line 244 was never true
245 tokens = tokens.to(self.cfg.device)
247 if one_zero_attention_mask is not None:
248 one_zero_attention_mask = one_zero_attention_mask.to(self.cfg.device)
250 resid = self.hook_embed(self.embed(tokens))
252 if one_zero_attention_mask is not None:
253 additive_attention_mask = (
254 repeat(1 - one_zero_attention_mask, "batch pos -> batch 1 1 pos")
255 ) * torch.finfo(self.cfg.dtype).min
256 else:
257 additive_attention_mask = None
259 query_len = key_len = tokens.shape[1]
261 encoder_positional_bias = self.encoder[0].attn.compute_relative_attention_bias(
262 query_len, key_len, device=self.cfg.device
263 )
265 for encoder_block in self.encoder:
266 resid = encoder_block(
267 resid_pre=resid,
268 additive_attention_mask=additive_attention_mask,
269 position_bias=encoder_positional_bias,
270 )
272 encoder_resid = self.encoder_final_ln(resid)
274 if decoder_input is None: 274 ↛ 275line 274 didn't jump to line 275 because the condition on line 274 was never true
275 raise ValueError("decoder_input cannot be None when input is not a string")
276 decoder_resid = self.embed(decoder_input)
277 decoder_query_len = decoder_key_len = decoder_input.shape[1]
278 decoder_positional_bias = self.decoder[0].attn.compute_relative_attention_bias(
279 decoder_query_len, decoder_key_len, device=self.cfg.device
280 )
282 for decoder_block in self.decoder:
283 decoder_resid = decoder_block(
284 resid_pre=decoder_resid,
285 position_bias=decoder_positional_bias,
286 encoder_hidden_states=encoder_resid,
287 encoder_additive_attention_mask=additive_attention_mask,
288 )
290 decoder_resid = self.decoder_final_ln(decoder_resid)
292 if self.cfg.tie_word_embeddings: 292 ↛ 297line 292 didn't jump to line 297 because the condition on line 292 was always true
293 # Rescale output before projecting on vocab
294 # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
295 decoder_resid *= self.cfg.d_model**-0.5
297 logits = self.unembed(decoder_resid)
298 if return_type is None: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 return None
300 return logits
302 @torch.inference_mode()
303 def generate(
304 self,
305 input: Union[str, Int[torch.Tensor, "batch pos"]] = "",
306 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None,
307 max_new_tokens: int = 10,
308 stop_at_eos: bool = True,
309 eos_token_id: Optional[Union[int, List[int]]] = None,
310 do_sample: bool = True,
311 top_k: Optional[int] = None,
312 top_p: Optional[float] = None,
313 temperature: float = 1.0,
314 freq_penalty: float = 0.0,
315 return_type: Optional[str] = "input",
316 verbose: bool = True,
317 ) -> Union[Int[torch.Tensor, "batch new_tokens"], str]:
318 """Sample tokens from the T5 encoder-decoder model.
320 Sample tokens from the model until the model outputs eos_token or max_new_tokens is reached.
321 This function is primarily taken from HookedTransformer but adjusted for the HookedEncoderDecoder
322 architecture.
323 This function does not support key value caching and no default padding sides or prepend_bos.
325 To avoid fiddling with ragged tensors, if we input a batch of text and some sequences finish
326 (by producing an EOT token), we keep running the model on the entire batch, but throw away
327 the output for a finished sequence and just keep adding EOTs to pad.
329 This supports entering a single string, but not a list of strings - if the strings don't
330 tokenize to exactly the same length, this gets messy. If that functionality is needed,
331 convert them to a batch of tokens and input that instead.
333 Args:
334 input (Union[str, Int[torch.Tensor, "batch pos"])]): Either a batch of tokens ([batch,
335 pos]) or a text string (this will be converted to a batch of tokens with batch size
336 1).
337 max_new_tokens (int): Maximum number of tokens to generate.
338 stop_at_eos (bool): If True, stop generating tokens when the model outputs eos_token.
339 eos_token_id (Optional[Union[int, Sequence]]): The token ID to use for end
340 of sentence. If None, use the tokenizer's eos_token_id - required if using
341 stop_at_eos. It's also possible to provide a list of token IDs (not just the
342 eos_token_id), in which case the generation will stop when any of them are output
343 (useful e.g. for stable_lm).
344 do_sample (bool): If True, sample from the model's output distribution. Otherwise, use
345 greedy search (take the max logit each time).
346 top_k (int): Number of tokens to sample from. If None, sample from all tokens.
347 top_p (float): Probability mass to sample from. If 1.0, sample from all tokens. If <1.0,
348 we take the top tokens with cumulative probability >= top_p.
349 temperature (float): Temperature for sampling. Higher values will make the model more
350 random (limit of temp -> 0 is just taking the top token, limit of temp -> inf is
351 sampling from a uniform distribution).
352 freq_penalty (float): Frequency penalty for sampling - how much to penalise previous
353 tokens. Higher values will make the model more random.
354 return_type (Optional[str]): The type of the output to return - either a string (str),
355 a tensor of tokens (tensor) or whatever the format of the input was (input).
356 verbose (bool): If True, show tqdm progress bars for generation.
358 Returns:
359 outputs (torch.Tensor): [batch, new_tokens], generated sequence of new tokens
360 (by default returns same type as input).
361 """
363 if isinstance(input, str): 363 ↛ 375line 363 didn't jump to line 375 because the condition on line 363 was always true
364 # If text, convert to tokens (batch_size=1)
365 assert (
366 self.tokenizer is not None
367 ), "Must provide a tokenizer if passing a string to the model"
368 encoder_input, attention_mask = self.to_tokens(input)
370 # If attention mask is not provided, use the one from the tokenizer
371 one_zero_attention_mask = (
372 attention_mask if one_zero_attention_mask is None else one_zero_attention_mask
373 )
374 else:
375 assert isinstance(input, torch.Tensor) # keep mypy happy
376 encoder_input = input
378 # If tokens are provided, user should be aware that attention mask will not be inferred
379 if one_zero_attention_mask is None:
380 logging.warning(
381 "No attention mask provided. Assuming all tokens should be attended to."
382 )
384 if return_type == "input": 384 ↛ 390line 384 didn't jump to line 390 because the condition on line 384 was always true
385 if isinstance(input, str): 385 ↛ 388line 385 didn't jump to line 388 because the condition on line 385 was always true
386 return_type = "str"
387 else:
388 return_type = "tensor"
390 assert isinstance(encoder_input, torch.Tensor)
391 batch_size = encoder_input.shape[0]
392 device = get_device_for_block_index(0, self.cfg)
394 # For the decoder input, we start with a tensor of PAD tokens of shape (batch, 1)
395 assert self.tokenizer is not None
396 decoder_input = torch.full((batch_size, 1), self.tokenizer.pad_token_id).to(device)
398 stop_tokens: List[int] = []
399 eos_token_for_padding = 0
400 if stop_at_eos: 400 ↛ 426line 400 didn't jump to line 426 because the condition on line 400 was always true
401 tokenizer_has_eos_token = self.tokenizer.eos_token_id is not None
403 local_eos_token_id: Optional[Union[int, List[int]]] = eos_token_id
404 if local_eos_token_id is None: 404 ↛ 411line 404 didn't jump to line 411 because the condition on line 404 was always true
405 assert (
406 tokenizer_has_eos_token
407 ), "Must pass a eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
409 local_eos_token_id = self.tokenizer.eos_token_id
411 if isinstance(local_eos_token_id, int): 411 ↛ 416line 411 didn't jump to line 416 because the condition on line 411 was always true
412 stop_tokens = [local_eos_token_id]
413 eos_token_for_padding = local_eos_token_id
414 else:
415 # eos_token_id is a Sequence (e.g. list or tuple)
416 if local_eos_token_id is None:
417 raise ValueError("eos_token_id cannot be None here")
418 stop_tokens = local_eos_token_id
419 eos_token_for_padding = (
420 self.tokenizer.eos_token_id
421 if tokenizer_has_eos_token
422 else local_eos_token_id[0]
423 )
425 # An array to track which sequences in the batch have finished.
426 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
428 # Currently nothing in HookedTransformer changes with eval, but this is here in case
429 # that changes in the future.
430 self.eval()
431 for _ in tqdm.tqdm(range(max_new_tokens), disable=not verbose): 431 ↛ 478line 431 didn't jump to line 478 because the loop on line 431 didn't complete
432 # While generating, we keep generating logits, throw away all but the final logits,
433 # and then use those logits to sample from the distribution We keep adding the
434 # sampled tokens to the end of tokens.
435 # We input the entire sequence, as a [batch, pos] tensor, since we aren't using
436 # the cache.
438 # Encoder input will be the same for all iterations
439 # Decoder input will be appended with the new token each iteration
440 logits = self.forward(
441 encoder_input,
442 decoder_input=decoder_input,
443 one_zero_attention_mask=one_zero_attention_mask,
444 )
445 assert logits is not None
446 final_logits = logits[:, -1, :]
448 if do_sample: 448 ↛ 449line 448 didn't jump to line 449 because the condition on line 448 was never true
449 sampled_tokens = sample_logits(
450 final_logits,
451 top_k=top_k,
452 top_p=top_p,
453 temperature=temperature,
454 freq_penalty=freq_penalty,
455 tokens=decoder_input,
456 ).to(get_device_for_block_index(0, self.cfg))
457 else:
458 sampled_tokens = final_logits.argmax(-1).to(get_device_for_block_index(0, self.cfg))
460 if stop_at_eos: 460 ↛ 473line 460 didn't jump to line 473 because the condition on line 460 was always true
461 # For all unfinished sequences, add on the next token. If a sequence was
462 # finished, throw away the generated token and add eos_token_for_padding
463 # instead.
464 sampled_tokens[finished_sequences] = eos_token_for_padding
465 finished_sequences.logical_or_(
466 torch.isin(
467 sampled_tokens.to(self.cfg.device),
468 torch.tensor(stop_tokens).to(self.cfg.device),
469 )
470 )
472 # Append new token to the decoder input
473 decoder_input = torch.cat([decoder_input, sampled_tokens.unsqueeze(-1)], dim=-1)
475 if stop_at_eos and finished_sequences.all():
476 break
478 if return_type == "str": 478 ↛ 484line 478 didn't jump to line 484 because the condition on line 478 was always true
479 assert self.tokenizer is not None
480 # Convert tokens to string
481 return cast(str, self.tokenizer.decode(decoder_input[0], skip_special_tokens=True))
483 else:
484 return decoder_input
486 @overload # type: ignore[overload-overlap]
487 def run_with_cache(
488 self, *model_args: Any, return_cache_object: Literal[True] = True, **kwargs: Any
489 ) -> Tuple[Float[torch.Tensor, "batch pos d_vocab"], ActivationCache]:
490 ...
492 @overload # type: ignore[overload-overlap]
493 def run_with_cache(
494 self, *model_args: Any, return_cache_object: Literal[False] = False, **kwargs: Any
495 ) -> Tuple[Float[torch.Tensor, "batch pos d_vocab"], Dict[str, torch.Tensor]]:
496 ...
498 def run_with_cache(
499 self,
500 *model_args: Any,
501 return_cache_object: bool = True,
502 remove_batch_dim: bool = False,
503 **kwargs: Any,
504 ) -> Tuple[
505 Float[torch.Tensor, "batch pos d_vocab"],
506 Union[ActivationCache, Dict[str, torch.Tensor]],
507 ]:
508 """
509 Wrapper around run_with_cache in HookedRootModule. If return_cache_object is True, this will return an ActivationCache object, with a bunch of useful HookedTransformer specific methods, otherwise it will return a dictionary of activations as in HookedRootModule. This function was copied directly from HookedTransformer.
510 """
511 out, cache_dict = super().run_with_cache(
512 *model_args, remove_batch_dim=remove_batch_dim, **kwargs
513 )
514 if return_cache_object: 514 ↛ 518line 514 didn't jump to line 518 because the condition on line 514 was always true
515 cache = ActivationCache(cache_dict, self, has_batch_dim=not remove_batch_dim)
516 return out, cache
517 else:
518 return out, cache_dict
520 def to(self: T, *args: Any, **kwargs: Any) -> T:
521 return super().to(*args, **kwargs)
523 def cuda(self: T, device: Optional[Union[int, torch.device]] = None) -> T:
524 if isinstance(device, int):
525 return self.to(f"cuda:{device}")
526 elif device is None:
527 return self.to("cuda")
528 else:
529 return self.to(device)
531 def cpu(self: T) -> T:
532 return self.to("cpu")
534 def mps(self: T) -> T:
535 """Warning: MPS may produce silently incorrect results. See #1178."""
536 warn_if_mps("mps")
537 return self.to(torch.device("mps"))
539 @classmethod
540 def from_pretrained(
541 cls: Type[T],
542 model_name: str,
543 checkpoint_index: Optional[int] = None,
544 checkpoint_value: Optional[int] = None,
545 hf_model: Optional[Any] = None,
546 device: Optional[str] = None,
547 tokenizer: Optional[Any] = None,
548 move_to_device: bool = True,
549 dtype: Optional[torch.dtype] = torch.float32,
550 **from_pretrained_kwargs: Any,
551 ) -> T:
552 """Loads in the pretrained weights from huggingface. Currently supports loading weight from HuggingFace BertForMaskedLM. Unlike HookedTransformer, this does not yet do any preprocessing on the model."""
553 import warnings
555 warnings.warn(
556 "HookedEncoderDecoder.from_pretrained is deprecated and will be removed in a "
557 "future major release. Use TransformerBridge.boot_transformers(...) instead — "
558 "the bridge supports T5-style encoder-decoder models. See "
559 "docs/source/content/migrating_to_v3.md.",
560 DeprecationWarning,
561 stacklevel=2,
562 )
564 logging.warning(
565 "Support for T5 in TransformerLens is currently experimental, until such a time when it has feature "
566 "parity with HookedTransformer and has been tested on real research tasks. Until then, backward "
567 "compatibility is not guaranteed. Please see the docs for information on the limitations of the current "
568 "implementation."
569 "\n"
570 "If using T5 for interpretability research, keep in mind that T5 has some significant architectural "
571 "differences to GPT. The major one is that T5 is an Encoder-Decoder model"
572 "Also, it uses relative positional embeddings, different types of Attention (without bias) and LayerNorm"
573 )
575 if from_pretrained_kwargs.get("load_in_8bit", False) or from_pretrained_kwargs.get( 575 ↛ 578line 575 didn't jump to line 578 because the condition on line 575 was never true
576 "load_in_4bit", False
577 ):
578 raise ValueError("Quantization not supported")
580 if "torch_dtype" in from_pretrained_kwargs: 580 ↛ 581line 580 didn't jump to line 581 because the condition on line 580 was never true
581 dtype = from_pretrained_kwargs["torch_dtype"]
583 if dtype is None: 583 ↛ 584line 583 didn't jump to line 584 because the condition on line 583 was never true
584 dtype = torch.float32
586 name_or_path = (
587 model_name if Path(model_name).exists() else loading.get_official_model_name(model_name)
588 )
590 cfg = loading.get_pretrained_model_config(
591 name_or_path,
592 checkpoint_index=checkpoint_index,
593 checkpoint_value=checkpoint_value,
594 fold_ln=False,
595 device=device,
596 n_devices=1,
597 dtype=dtype,
598 **from_pretrained_kwargs,
599 )
601 state_dict = loading.get_pretrained_state_dict(
602 name_or_path, cfg, hf_model, dtype=dtype, **from_pretrained_kwargs
603 )
605 model = cls(cfg, tokenizer, move_to_device=False)
607 model.load_state_dict(state_dict, strict=False)
609 if move_to_device and cfg.device is not None: 609 ↛ 612line 609 didn't jump to line 612 because the condition on line 609 was always true
610 model.to(cfg.device)
612 print(f"Loaded pretrained model {model_name} into HookedTransformer")
614 return model
616 @property
617 def W_U(self) -> Float[torch.Tensor, "d_model d_vocab"]:
618 """
619 Convenience to get the unembedding matrix (ie the linear map from the final residual stream to the output logits)
620 """
621 return self.unembed.W_U
623 @property
624 def b_U(self) -> Float[torch.Tensor, "d_vocab"]:
625 """
626 Convenience to get the unembedding bias
627 """
628 return self.unembed.b_U
630 @property
631 def W_E(self) -> Float[torch.Tensor, "d_vocab d_model"]:
632 """
633 Convenience to get the embedding matrix
634 """
635 return self.embed.W_E
637 @property
638 def W_pos(self) -> None:
639 """
640 Convenience function to get the positional embedding. Only works on models with absolute positional embeddings!
641 """
642 raise NotImplementedError(
643 "T5 does not have absolute positional embeddings. Uses relative positional embeddings instead."
644 )
646 @property
647 def W_K(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]:
648 """Stacks the key weights across all layers"""
649 return torch.stack(
650 [block.attn.W_K for block in chain(self.encoder, self.decoder)],
651 dim=0,
652 )
654 @property
655 def W_Q(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]:
656 """Stacks the query weights across all layers"""
657 return torch.stack(
658 [block.attn.W_Q for block in chain(self.encoder, self.decoder)],
659 dim=0,
660 )
662 @property
663 def W_V(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]:
664 """Stacks the value weights across all layers"""
665 return torch.stack(
666 [block.attn.W_V for block in chain(self.encoder, self.decoder)],
667 dim=0,
668 )
670 @property
671 def W_O(self) -> Float[torch.Tensor, "n_layers n_heads d_head d_model"]:
672 """Stacks the attn output weights across all layers"""
673 return torch.stack(
674 [block.attn.W_O for block in chain(self.encoder, self.decoder)],
675 dim=0,
676 )
678 @property
679 def W_in(self) -> Float[torch.Tensor, "n_layers d_model d_mlp"]:
680 """Stacks the MLP input weights across all layers"""
681 weights: List[torch.Tensor] = []
682 for block in chain(self.encoder, self.decoder):
683 mlp = block.mlp
684 if isinstance(mlp, (MLP, GatedMLP)):
685 weights.append(mlp.W_in)
686 else:
687 raise NotImplementedError(
688 f"W_in property is not supported for MLP of type {type(mlp).__name__}"
689 )
690 return torch.stack(weights, dim=0)
692 @property
693 def W_out(self) -> Float[torch.Tensor, "n_layers d_mlp d_model"]:
694 """Stacks the MLP output weights across all layers"""
695 weights: List[torch.Tensor] = []
696 for block in chain(self.encoder, self.decoder):
697 mlp = block.mlp
698 if isinstance(mlp, (MLP, GatedMLP)):
699 weights.append(mlp.W_out)
700 else:
701 raise NotImplementedError(
702 f"W_out property is not supported for MLP of type {type(mlp).__name__}"
703 )
704 return torch.stack(weights, dim=0)
706 @property
707 def b_K(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]:
708 """Stacks the key biases across all layers"""
709 return torch.stack(
710 [block.attn.b_K for block in chain(self.encoder, self.decoder)],
711 dim=0,
712 )
714 @property
715 def b_Q(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]:
716 """Stacks the query biases across all layers"""
717 return torch.stack(
718 [block.attn.b_Q for block in chain(self.encoder, self.decoder)],
719 dim=0,
720 )
722 @property
723 def b_V(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]:
724 """Stacks the value biases across all layers"""
725 return torch.stack(
726 [block.attn.b_V for block in chain(self.encoder, self.decoder)],
727 dim=0,
728 )
730 @property
731 def b_O(self) -> Float[torch.Tensor, "n_layers d_model"]:
732 """Stacks the attn output biases across all layers"""
733 return torch.stack(
734 [block.attn.b_O for block in chain(self.encoder, self.decoder)],
735 dim=0,
736 )
738 @property
739 def b_in(self) -> Float[torch.Tensor, "n_layers d_mlp"]:
740 """Stacks the MLP input biases across all layers"""
741 biases: List[torch.Tensor] = []
742 for block in chain(self.encoder, self.decoder):
743 mlp = block.mlp
744 if isinstance(mlp, (MLP, GatedMLP)):
745 biases.append(mlp.b_in)
746 else:
747 raise NotImplementedError(
748 f"b_in property is not supported for MLP of type {type(mlp).__name__}"
749 )
750 return torch.stack(biases, dim=0)
752 @property
753 def b_out(self) -> Float[torch.Tensor, "n_layers d_model"]:
754 """Stacks the MLP output biases across all layers"""
755 biases: List[torch.Tensor] = []
756 for block in chain(self.encoder, self.decoder):
757 mlp = block.mlp
758 if isinstance(mlp, (MLP, GatedMLP)):
759 biases.append(mlp.b_out)
760 else:
761 raise NotImplementedError(
762 f"b_out property is not supported for MLP of type {type(mlp).__name__}"
763 )
764 return torch.stack(biases, dim=0)
766 @property
767 def QK(self) -> FactoredMatrix: # [n_layers, n_heads, d_model, d_model]
768 """Returns a FactoredMatrix object with the product of the Q and K matrices for each layer and head.
769 Useful for visualizing attention patterns."""
770 return FactoredMatrix(self.W_Q, self.W_K.transpose(-2, -1))
772 @property
773 def OV(self) -> FactoredMatrix: # [n_layers, n_heads, d_model, d_model]
774 """Returns a FactoredMatrix object with the product of the O and V matrices for each layer and head."""
775 return FactoredMatrix(self.W_V, self.W_O)
777 def all_head_labels(self) -> List[str]:
778 """Returns a list of strings with the format "L{l}H{h}", where l is the layer index and h is the head index."""
779 return [f"EL{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)] + [
780 f"DL{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)
781 ]