Coverage for transformer_lens/config/hooked_transformer_config.py: 93%
142 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Hooked Transformer Config.
3Module with a dataclass for storing the configuration of a
4:class:`transformer_lens.HookedTransformer` model.
5"""
7from __future__ import annotations
9import pprint
10import random
11from dataclasses import dataclass
12from typing import Any, Dict, List, Optional, Union
14import numpy as np
15import torch
17from transformer_lens.utilities.activation_functions import (
18 SOFTCAP_DISABLED,
19 SUPPORTED_ACTIVATIONS,
20)
21from transformer_lens.utilities.devices import get_device
23from .transformer_lens_config import TransformerLensConfig
26@dataclass
27class HookedTransformerConfig(TransformerLensConfig):
28 """
29 Configuration class to store the configuration of a HookedTransformer model.
31 See further_comments.md for more details on the more complex arguments.
33 Args:
34 d_model (int): The dimensionality of the embeddings.
35 d_head (int): The dimensionality of each attention head.
36 n_layers (int): The number of transformer blocks (one block = one attn layer AND one MLP layer).
37 n_ctx (int): The maximum sequence length.
38 n_heads (int): The number of attention heads. If not
39 specified, will be set to d_model // d_head. (This is represented by a default value of -1)
40 d_mlp (int, *optional*): The dimensionality of the feedforward mlp
41 network. Defaults to 4 * d_model, and in an attn-only model is None.
42 d_vocab (int): The size of the vocabulary. Defaults to -1, which means not set. If not set, will be
43 automatically set from the tokenizer's vocab size.
44 act_fn (str, *optional*): The activation function to use. Always
45 lowercase. Supports ['relu', 'gelu', 'silu', 'gelu_new', 'solu_ln',
46 'gelu_fast']. Must be set unless using an attn-only model.
47 eps (float): The epsilon value to use for layer normalization. Defaults
48 to 1e-5
49 use_attn_result (bool): whether to explicitly calculate the amount
50 each head adds to the residual stream (with a hook) and THEN add it
51 up, vs just calculating the sum. This can be very memory intensive
52 for large models, so defaults to False
53 use_split_qkv_input (bool): whether to explicitly calculate the input of
54 each head separately, with a hook. Defaults to false to save memory.
55 use_hook_mlp_in (bool): whether to use a hook to get the input to the
56 MLP layer. Defaults to false to save memory.
57 use_attn_in (bool): whether to explicitly calculate the input of each
58 attention head separately, with a hook. Defaults to false to save memory
59 use_attn_scale (bool): whether to scale the attention weights by
60 1/sqrt(d_head)
61 ungroup_grouped_query_attention (bool): whether to ungroup key and value heads, for models that use
62 grouped query attention.
63 attn_scale (float): The amount to divide attention scores by (if applicable). Defaults to
64 sqrt(d_head)
65 model_name (str): the name of the model, used to load
66 weights from HuggingFace or initialized to "custom" if not passed
67 original_architecture (str, *optional*): the family of the model, used
68 to help load
69 weights from HuggingFace or initialized to "custom" if not passed
70 from_checkpoint (bool): Whether the model weights were
71 loaded from a checkpoint (only applies to pretrained models)
72 checkpoint_index (int, *optional*): The index of the
73 checkpoint loaded (only applies to pretrained models).
74 checkpoint_label_type (str, *optional*): Whether
75 checkpoints are labelled by the number of steps or number of tokens.
76 checkpoint_value (int, *optional*): The value of the
77 checkpoint label (whether of steps or tokens).
78 tokenizer_name (str, *optional*): the full name of the model, passed into
79 HuggingFace to access the tokenizer. Only used when passing in
80 custom config, if loading from pretrained then this is not needed.
81 use_local_attn (bool): whether to use local attention - ie each
82 destination token can only attend to source tokens a certain distance back.
83 window_size (int, *optional*): the size of the window for local
84 attention
85 attn_types (List[str], *optional*): the types of attention to use for
86 local attention
87 init_mode (str): the initialization mode to use for the
88 weights. Only relevant for custom models, ignored for pre-trained.
89 We now support 'gpt2', 'xavier_uniform', 'xavier_normal', 'kaiming_uniform',
90 'kaiming_normal'. MuP support to come. Defaults to 'gpt2'.
91 normalization_type (str, *optional*): the type of normalization to use.
92 Options are None (no normalization), 'LN' (use LayerNorm, including weights
93 & biases) and 'LNPre' (use LayerNorm, but no weights or biases), 'RMS'
94 (use RMSNorm, including weights) and 'RMSPre' (use RMSNorm, but no weights or biases).
95 Defaults to LN
96 device(str): The device to use for the model. Defaults to 'cuda' if
97 available, else 'cpu'. Must be 'cuda' if `n_devices` > 1.
98 n_devices (int): The number of devices to use for the model. Defaults to 1. Layers are loaded
99 to support "pipeline parallelism", where each device is responsible for a subset of the layers.
100 attention_dir (str): Whether to use causal (aka unidirectional aka GPT-2
101 style) or bidirectional attention. Options are 'causal' and
102 'bidirectional'. Defaults to 'causal'
103 attn_only (bool): Whether to only use attention layers, no feedforward
104 layers. Defaults to False
105 seed (int, *optional*): The seed to use for the model.
106 Used to set sources of randomness (Python, PyTorch and NumPy) and to initialize weights.
107 Defaults to None. We recommend setting a seed, so your experiments are reproducible.
108 initializer_range (float): The standard deviation of the normal used to
109 initialise the weights, initialized to 0.8 / sqrt(d_model). If init_mode is
110 'xavier_uniform' or 'xavier_normal', this value is instead treated as the `gain` parameter for the weight
111 initialisation (a constant factor to scale the weights by). Defaults to -1.0, which means not set.
112 init_weights (bool): Whether to initialize the weights. Defaults to
113 True. If False, does not initialize weights.
114 scale_attn_by_inverse_layer_idx (bool): Whether to scale the attention
115 weights by 1/(layer_id+1), used by Mistral (Stanford) models for numerical stability when
116 training in FP16. Defaults to False.
117 positional_embedding_type (str): The positional embedding used. Options
118 are 'standard' (ie GPT-2 style, absolute, randomly initialized learned positional
119 embeddings, directly added to the residual stream), 'rotary'
120 (described here: https://blog.eleuther.ai/rotary-embeddings/ ) and
121 'shortformer' (GPT-2 style absolute & learned, but rather than being
122 added to the residual stream they're only added to the inputs to the
123 keys and the queries (ie key = W_K(res_stream + pos_embed), but
124 values and MLPs don't get any positional info)). Sinusoidal are not
125 currently supported. Defaults to 'standard'.
126 final_rms (bool): Whether to replace the final normalization (just
127 before the unembed) with RMSNorm (ie no centering or bias, just
128 scaling + weights). Only included because of a dumb bug in my
129 original SoLU code. Defaults to False.
130 d_vocab_out (int, *optional*): The size of the output vocabulary. Defaults to -1, which means not set. If not
131 set, will be equal to d_vocab. Mainly useful for algorithmic tasks
132 where the input and output vocabularies may be different.
133 parallel_attn_mlp (bool): Whether to parallelize the attention and MLP
134 layers - a weird cursed thing done by GPT-J. Means that
135 mlp_out=MLP(ln1(resid_pre)) and resid_post=resid_pre+attn_out+mlp_out. Defaults to False.
136 rotary_dim (int, *optional*): The dimensionality of the rotary
137 embeddings, may be d_head in which case only the first rotary_dim
138 dimensions of each head are rotated. Defaults to None, if
139 positional_embedding_type=="rotary" post-init then sets it to d_head, i.e. "rotate all
140 dimensions of the query and key".
141 n_params (int, *optional*): The number of "hidden weight" parameters
142 in the model, **excluding** embeddings, unembedding, biases, and
143 layer norms. Counts only the attention projections (W_Q, W_K, W_V,
144 W_O) and MLP weights (W_in, W_out, plus W_gate when ``gated_mlp=True``).
145 This matches the convention from the
146 `scaling laws paper <https://arxiv.org/pdf/2001.08361.pdf>`_,
147 which found this to be the most meaningful number for predicting
148 performance. **Note:** this is NOT the same as
149 ``sum(p.numel() for p in model.parameters())`` — that would
150 include embeddings and biases and yield a larger number. Use the
151 ``sum(p.numel() ...)`` form if you want the total parameter count
152 (e.g. for memory-budget calculations). Automatically calculated;
153 not intended to be set by the user.
154 use_hook_tokens (bool): Will add a hook point on the token input to
155 HookedTransformer.forward, which lets you cache or intervene on the tokens.
156 Defaults to False.
157 gated_mlp (bool): If True, the MLP layer uses a gated formulation
158 (SwiGLU/GeGLU-style): ``mlp_out = W_out @ (act_fn(W_gate @ x) * (W_in @ x))``,
159 with an extra ``W_gate`` weight matrix alongside ``W_in`` and ``W_out``. Used by
160 LLaMA, Mistral, Gemma, Qwen and similar families. When False (default), the MLP
161 is the plain ``mlp_out = W_out @ act_fn(W_in @ x)`` form. ``loading_from_pretrained``
162 sets this automatically per architecture; only set manually for a custom config.
163 default_prepend_bos (bool, optional): Default behavior of whether to prepend the BOS token when the
164 methods of HookedTransformer process input text to tokenize (only when input is a string).
165 Defaults to True - even for models not explicitly trained with this, heads often use the
166 first position as a resting position and accordingly lose information from the first token,
167 so this empirically seems to give better results. To change the default behavior to False, pass in
168 default_prepend_bos=False. Note that you can also locally override the default behavior by passing
169 in prepend_bos=True/False when you call a method that processes the input string.
170 dtype (torch.dtype, *optional*): The model's dtype. Defaults to torch.float32.
171 tokenizer_prepends_bos (bool, *optional*): This flag is set by set_tokenizer. It is set to True only
172 when the tokenizer automatically prepends the BOS token if initialized with add_bos_token=True.
173 We need this information to dynamically control bos prepending.
174 load_in_4bit(bool): If this flag is set, then it's assumed that parameters are 4-bit quantized
175 with bitsandbytes. Currently only supported for Llama.
176 quantization_method (str, *optional*): the ``quant_method`` declared by the checkpoint's
177 HF config ("mxfp4", "bitsandbytes", "gptq", ...), captured while that config is already
178 in hand so later load steps need not refetch it. None when unquantized, and also when
179 the config was never fetched (the llama/gemma name-based branches of
180 ``convert_hf_model_config`` infer the architecture from the model name instead).
181 n_key_value_heads (int, *optional*): The number of groups of heads that use the same key and value matrix.
182 Only for models that use Grouped Query Attention.
183 post_embedding_ln (bool): Whether to apply layer normalization after embedding the tokens. Defaults
184 to False.
185 num_experts (int, *optional*): The number of experts to use in the MoE layer. If set, experts_per_token
186 must also be set. Set to None if not using MoE.
187 experts_per_token (int, *optional*): The number of experts to use for each pass in the MoE layer. If set,
188 num_experts must also be set. Set to None if not using MoE.
189 relative_attention_max_distance (int, *optional*): The maximum distance between tokens for relative
190 attention. If set, relative_attention_num_buckets must also be set.Only used in EncoderDecoder models, like T5.
191 relative_attention_num_buckets (int, *optional*): The number of buckets to use for relative attention.
192 If set, relative_attention_max_distance must also be set.Only used in EncoderDecoder models, like T5.
193 decoder_start_token_id (int, *optional*): The start token id for the decoder. Only used in EncoderDecoder models, like T5.
194 tie_word_embeddings (bool): Whether to tie the word embeddings and the output layer weights. Defaults to False. Only used in EncoderDecoder (T5) by now.
195 use_normalization_before_and_after (bool): Whether to apply normalization (LN/RMS/etc)
196 to both the input of an attn/MLP block *and* the output (before adding back to the
197 residual stream). Currently only used in Gemma-2. Defaults to False.
198 attn_scores_soft_cap (float): An optional softcap for attention scores pre-softmax. If
199 used, it will map attn_scores -> soft_cap * tanh(attn_scores / soft_cap). As tanh's
200 output is in [-1, 1], this maps attn_scores to [-soft_cap, soft_cap], with little
201 effect on small values, but squashing large values into that interval. Currently only
202 used in Gemma-2. Defaults to -1.0, which means not set.
203 output_logits_soft_cap (float): An optional softcap for output logits, currently only used
204 in Gemma-2 (see attn_scores_soft_cap for details). Defaults to -1.0, which means not
205 set.
206 use_NTK_by_parts_rope (bool): Whether to apply the "NTK-by-parts" method when using Rotary
207 Positional Embedding. This method adjusts the interpolation based on frequency factors
208 for different parts of the hidden dimensions. See Section 3.2 in
209 https://arxiv.org/pdf/2309.00071 for details. Defaults to False.
210 NTK_by_parts_low_freq_factor (float): The threshold applied to low-frequency hidden
211 dimensions during interpolation when using the "NTK-by-parts" method. Defaults to 1.0.
212 NTK_by_parts_high_freq_factor (float): The threshold applied to high-frequency hidden
213 dimensions during interpolation in the "NTK-by-parts" method. Defaults to 4.0.
214 NTK_by_parts_factor (float): The overall factor used in the "NTK-by-parts" method that
215 affects the rate of change between low and high-frequency interpolation strategies.
216 Defaults to 8.0.
217 use_yarn_rope (bool): Whether to apply YARN (Yet Another RoPE extensioN) scaling to
218 rotary positional embeddings. YARN blends interpolated and extrapolated frequencies
219 per dimension using correction ranges. See https://arxiv.org/abs/2309.00071 for
220 details. Used by OLMo 3. Defaults to False.
221 yarn_factor (float): The interpolation factor for YARN RoPE scaling. Defaults to 1.0.
222 yarn_attention_factor (float): Multiplicative scaling applied to sin/cos embeddings in
223 YARN. Defaults to 1.0.
224 yarn_beta_fast (float): Upper rotation threshold for YARN correction range. Defaults to 32.
225 yarn_beta_slow (float): Lower rotation threshold for YARN correction range. Defaults to 1.
226 yarn_truncate (bool): Whether to floor/ceil the YARN correction-range bounds
227 (HF's `truncate`). GPT-OSS ships truncate=False. Defaults to True.
228 yarn_global_attn_only (bool): Whether YARN applies only to global-attention
229 layers, with sliding/local layers keeping plain rope (Olmo-3's
230 per-layer-type rope). Defaults to False.
231 use_attention_sinks (bool): Whether attention carries a learned per-head sink
232 logit (GPT-OSS) that joins the softmax as an extra key column and is
233 dropped afterward. Defaults to False.
234 yarn_original_max_position_embeddings (int): The original max position embeddings before
235 YARN extension. Defaults to 4096.
236 use_qk_norm (bool): Whether to apply RMSNorm to the query and key projections before
237 computing attention scores. Used by Gemma 3 models. Defaults to False.
238 rotary_base_local (float, *optional*): The base for rotary positional embeddings in local
239 attention layers. Used by models with hybrid local/global attention (e.g., Gemma 3)
240 which use different RoPE bases for local (10k) and global (1M) attention. Defaults
241 to None, which means the standard rotary_base is used for all layers.
242 norm_topk_prob (bool): Whether to normalize the top-k probabilities in the MoE layer.
243 use_logn_attn (bool): Qwen-1's log-n attention: scale queries by
244 log_{train_len}(position) past the training length (eval only).
245 train_seq_length (int, *optional*): the length the model was trained at
246 (Qwen-1's ``seq_length``). Both log-n scaling and dynamic-NTK RoPE
247 threshold on it, and it stays fixed when n_ctx is overridden upward.
248 use_dynamic_ntk_rope (bool): Qwen-1's dynamic NTK: rescale the rotary
249 base by ``alpha ** (rotary_dim / (rotary_dim - 2))`` once the key
250 length exceeds the training length (eval only).
251 clip_qkv (float, *optional*): Clamp Q/K/V activations to [-clip_qkv, clip_qkv] after
252 projection (and any qk-norm), as OLMo v1 and OLMoE do. Defaults to None (no clamp).
253 """
255 model_name: str = "custom"
256 act_fn: str = "relu"
257 eps: float = 1e-5
258 use_attn_scale: bool = True
259 attn_scale: float = -1.0
260 use_hook_mlp_in: bool = False
261 use_attn_in: bool = False
262 use_qk_norm: bool = False
263 clip_qkv: Optional[float] = None
264 use_logn_attn: bool = False
265 train_seq_length: Optional[int] = None
266 use_dynamic_ntk_rope: bool = False
267 use_local_attn: bool = False
268 ungroup_grouped_query_attention: bool = False
269 original_architecture: Optional[str] = None
270 from_checkpoint: bool = False
271 checkpoint_index: Optional[int] = None
272 checkpoint_label_type: Optional[str] = None
273 checkpoint_value: Optional[int] = None
274 tokenizer_name: Optional[str] = None
275 window_size: Optional[int] = None
276 attn_types: Optional[List] = None
277 init_mode: str = "gpt2"
278 normalization_type: Optional[str] = "LN"
279 n_devices: int = 1
280 attention_dir: str = "causal"
281 attn_only: bool = False
282 seed: Optional[int] = None
283 initializer_range: float = -1.0
284 init_weights: bool = True
285 scale_attn_by_inverse_layer_idx: bool = False
286 final_rms: bool = False
287 d_vocab_out: int = -1
288 parallel_attn_mlp: bool = False
289 rotary_dim: Optional[int] = None
290 n_params: Optional[int] = None
291 use_hook_tokens: bool = False
292 gated_mlp: bool = False
293 dtype: torch.dtype = torch.float32
294 tokenizer_prepends_bos: Optional[bool] = None
295 post_embedding_ln: bool = False
296 rotary_base: Union[float, int] = 10000
297 rotary_base_local: Optional[
298 Union[float, int]
299 ] = None # For models with different RoPE bases per attention type (e.g., Gemma 3)
300 rotary_scaling_factor: float = (
301 1.0 # Linear RoPE scaling factor for global attention (e.g., 8.0 for Gemma 3 4B)
302 )
303 trust_remote_code: bool = False
304 rotary_adjacent_pairs: bool = False
305 load_in_4bit: bool = False
306 quantization_method: Optional[str] = None
307 num_experts: Optional[int] = None
308 experts_per_token: Optional[int] = None
309 relative_attention_max_distance: Optional[int] = None
310 relative_attention_num_buckets: Optional[int] = None
311 decoder_start_token_id: Optional[int] = None
312 tie_word_embeddings: bool = False
313 use_normalization_before_and_after: bool = False
314 attn_scores_soft_cap: float = SOFTCAP_DISABLED
315 output_logits_soft_cap: float = SOFTCAP_DISABLED
316 use_NTK_by_parts_rope: bool = False
317 NTK_by_parts_low_freq_factor: float = 1.0
318 NTK_by_parts_high_freq_factor: float = 4.0
319 NTK_by_parts_factor: float = 8.0
320 NTK_original_ctx_len: int = 8192
321 use_yarn_rope: bool = False
322 yarn_factor: float = 1.0
323 yarn_attention_factor: float = 1.0
324 yarn_beta_fast: float = 32.0
325 yarn_beta_slow: float = 1.0
326 yarn_original_max_position_embeddings: int = 4096
327 # HF yarn's `truncate` option: floor/ceil the correction range bounds.
328 # GPT-OSS ships truncate=False, keeping the bounds fractional.
329 yarn_truncate: bool = True
330 # Per-layer-type rope (Olmo-3): YARN applies only on global-attention
331 # layers; sliding/local layers keep plain rope.
332 yarn_global_attn_only: bool = False
333 # GPT-OSS: learned per-head sink logit that joins the attention softmax as
334 # an extra key column and is dropped afterward, so real positions share
335 # probability mass with the sink.
336 use_attention_sinks: bool = False
337 norm_topk_prob: bool = False
339 def __post_init__(self):
340 super().__post_init__()
342 if self.seed is not None: 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true
343 self.set_seed_everywhere(self.seed)
344 if self.use_local_attn:
345 assert self.window_size is not None, "window_size must be specified for local attention"
346 assert self.attn_types is not None, "attn_types must be specified for local attention"
347 if not self.attn_only:
348 assert self.act_fn is not None, "act_fn must be specified for non-attn-only models"
349 assert (
350 self.act_fn in SUPPORTED_ACTIVATIONS
351 ), f"act_fn={self.act_fn} must be one of {SUPPORTED_ACTIVATIONS}"
352 if self.initializer_range < 0 and self.init_mode == "gpt2":
353 # Roughly copy the GPT-2 value, but proportional to sqrt(1/d_model)
354 self.initializer_range = 0.8 / np.sqrt(self.d_model)
355 if self.initializer_range < 0 and self.init_mode != "gpt2": 355 ↛ 357line 355 didn't jump to line 357 because the condition on line 355 was never true
356 # This is the gain parameter for the weight initialisation
357 self.initializer_range = 1.0
359 if self.d_vocab_out == -1:
360 # d_vocab_out defaults to d_vocab, unless there's an algorithmic task
361 # If d_vocab is not set, it'll be inferred from tokenizer_name or from a tokenizer
362 # explicitly passed to HookedTransformer initialisation.
363 self.d_vocab_out = self.d_vocab
365 if self.positional_embedding_type == "rotary" and self.rotary_dim is None:
366 self.rotary_dim = self.d_head
368 if self.num_experts is not None:
369 assert (
370 self.experts_per_token is not None
371 ), "experts_per_token must be set if num_experts is set"
372 if self.experts_per_token is not None:
373 assert (
374 self.num_experts is not None
375 ), "num_experts must be set if experts_per_token is set"
377 # Attention params (W_Q, W_K, W_V, W_O), ignoring biases/LN
378 self.n_params = self.n_layers * ((self.d_model * self.d_head * self.n_heads * 4))
379 if not self.attn_only:
380 assert self.d_mlp is not None # mypy
381 # MLP params (W_in, W_out), ignoring biases/LN
382 mlp_params_per_layer = self.d_model * self.d_mlp * (2 + self.gated_mlp)
384 if self.num_experts:
385 # Scale by num_experts and add gate params
386 mlp_params_per_layer = (mlp_params_per_layer + self.d_model) * self.num_experts
387 self.n_params += self.n_layers * mlp_params_per_layer
389 if self.device is None: 389 ↛ 390line 389 didn't jump to line 390 because the condition on line 389 was never true
390 self.device = get_device()
391 else:
392 from transformer_lens.utilities import warn_if_mps
394 warn_if_mps(self.device)
396 if self.n_devices > 1: 396 ↛ 397line 396 didn't jump to line 397 because the condition on line 396 was never true
397 assert (
398 torch.cuda.device_count() >= self.n_devices
399 ), f"Not enough CUDA devices to support n_devices {self.n_devices}"
401 if self.use_attn_scale and self.attn_scale == -1.0:
402 self.attn_scale = np.sqrt(self.d_head)
404 assert self.default_prepend_bos in [
405 True,
406 False,
407 ], f"default_prepend_bos must be either True or False, but {self.default_prepend_bos} is given"
409 @classmethod
410 def unwrap(cls, config: Union[Dict, "TransformerLensConfig"]) -> HookedTransformerConfig:
411 """
412 Convenience function to avoid duplicate code from a common way config is passed to various components
413 """
414 if isinstance(config, Dict):
415 return cls.from_dict(config)
416 elif isinstance(config, cls): 416 ↛ 420line 416 didn't jump to line 420 because the condition on line 416 was always true
417 return config
418 else:
419 # Convert from TransformerLensConfig to HookedTransformerConfig
420 return cls.from_dict(config.to_dict())
422 @classmethod
423 def from_dict(cls, config_dict: Dict[str, Any]) -> HookedTransformerConfig:
424 """
425 Instantiates a `HookedTransformerConfig` from a Python dictionary of
426 parameters.
427 """
428 return cls(**config_dict)
430 def to_dict(self):
431 return self.__dict__
433 def __repr__(self):
434 return "HookedTransformerConfig:\n" + pprint.pformat(self.to_dict())
436 def set_seed_everywhere(self, seed: int):
437 torch.manual_seed(seed)
438 random.seed(seed)
439 np.random.seed(seed)
441 def is_layer_norm_activation(self) -> bool:
442 return self.act_fn is not None and self.act_fn.endswith("_ln")