transformer_lens.config.hooked_transformer_config module

Hooked Transformer Config.

Module with a dataclass for storing the configuration of a transformer_lens.HookedTransformer model.

class transformer_lens.config.hooked_transformer_config.HookedTransformerConfig(d_model: int, d_head: int, n_layers: int, n_ctx: int, n_heads: int = -1, d_mlp: int | None = None, d_vocab: int = -1, device: str | None = None, use_attn_result: bool = False, use_split_qkv_input: bool = False, default_prepend_bos: bool = True, positional_embedding_type: str = 'standard', n_key_value_heads: int | None = None, attn_only: bool = False, gated_mlp: bool = False, uses_rms_norm: bool = False, eps: float = 1e-05, layer_norm_folding: bool = False, act_fn: str = 'relu', normalization_type: str | None = 'LN', num_experts: int | None = None, experts_per_token: int | None = None, final_rms: bool = False, dtype: dtype = torch.float32, model_name: str = 'custom', use_attn_scale: bool = True, attn_scale: float = -1.0, use_hook_mlp_in: bool = False, use_attn_in: bool = False, use_qk_norm: bool = False, clip_qkv: float | None = None, use_logn_attn: bool = False, train_seq_length: int | None = None, use_dynamic_ntk_rope: bool = False, use_local_attn: bool = False, ungroup_grouped_query_attention: bool = False, original_architecture: str | None = None, from_checkpoint: bool = False, checkpoint_index: int | None = None, checkpoint_label_type: str | None = None, checkpoint_value: int | None = None, tokenizer_name: str | None = None, window_size: int | None = None, attn_types: List | None = None, init_mode: str = 'gpt2', n_devices: int = 1, attention_dir: str = 'causal', seed: int | None = None, initializer_range: float = -1.0, init_weights: bool = True, scale_attn_by_inverse_layer_idx: bool = False, d_vocab_out: int = -1, parallel_attn_mlp: bool = False, rotary_dim: int | None = None, n_params: int | None = None, use_hook_tokens: bool = False, tokenizer_prepends_bos: bool | None = None, post_embedding_ln: bool = False, rotary_base: float | int = 10000, rotary_base_local: float | int | None = None, rotary_scaling_factor: float = 1.0, trust_remote_code: bool = False, rotary_adjacent_pairs: bool = False, load_in_4bit: bool = False, quantization_method: str | None = None, relative_attention_max_distance: int | None = None, relative_attention_num_buckets: int | None = None, decoder_start_token_id: int | None = None, tie_word_embeddings: bool = False, use_normalization_before_and_after: bool = False, attn_scores_soft_cap: float = -1.0, output_logits_soft_cap: float = -1.0, use_NTK_by_parts_rope: bool = False, NTK_by_parts_low_freq_factor: float = 1.0, NTK_by_parts_high_freq_factor: float = 4.0, NTK_by_parts_factor: float = 8.0, NTK_original_ctx_len: int = 8192, use_yarn_rope: bool = False, yarn_factor: float = 1.0, yarn_attention_factor: float = 1.0, yarn_beta_fast: float = 32.0, yarn_beta_slow: float = 1.0, yarn_original_max_position_embeddings: int = 4096, yarn_truncate: bool = True, yarn_global_attn_only: bool = False, use_attention_sinks: bool = False, norm_topk_prob: bool = False)

Bases: TransformerLensConfig

Configuration class to store the configuration of a HookedTransformer model.

See further_comments.md for more details on the more complex arguments.

Parameters:

  • d_model (int) – The dimensionality of the embeddings.

  • d_head (int) – The dimensionality of each attention head.

  • n_layers (int) – The number of transformer blocks (one block = one attn layer AND one MLP layer).

  • n_ctx (int) – The maximum sequence length.

  • n_heads (int) – The number of attention heads. If not specified, will be set to d_model // d_head. (This is represented by a default value of -1)

  • d_mlp (int, optional) – The dimensionality of the feedforward mlp network. Defaults to 4 * d_model, and in an attn-only model is None.

  • d_vocab (int) – The size of the vocabulary. Defaults to -1, which means not set. If not set, will be automatically set from the tokenizer’s vocab size.

  • act_fn (str, optional) – The activation function to use. Always lowercase. Supports [‘relu’, ‘gelu’, ‘silu’, ‘gelu_new’, ‘solu_ln’, ‘gelu_fast’]. Must be set unless using an attn-only model.

  • eps (float) – The epsilon value to use for layer normalization. Defaults to 1e-5

  • use_attn_result (bool) – whether to explicitly calculate the amount each head adds to the residual stream (with a hook) and THEN add it up, vs just calculating the sum. This can be very memory intensive for large models, so defaults to False

  • use_split_qkv_input (bool) – whether to explicitly calculate the input of each head separately, with a hook. Defaults to false to save memory.

  • use_hook_mlp_in (bool) – whether to use a hook to get the input to the MLP layer. Defaults to false to save memory.

  • use_attn_in (bool) – whether to explicitly calculate the input of each attention head separately, with a hook. Defaults to false to save memory

  • use_attn_scale (bool) – whether to scale the attention weights by 1/sqrt(d_head)

  • ungroup_grouped_query_attention (bool) – whether to ungroup key and value heads, for models that use grouped query attention.

  • attn_scale (float) – The amount to divide attention scores by (if applicable). Defaults to sqrt(d_head)

  • model_name (str) – the name of the model, used to load weights from HuggingFace or initialized to “custom” if not passed

  • original_architecture (str, optional) – the family of the model, used

  • load (to help) – weights from HuggingFace or initialized to “custom” if not passed

  • from_checkpoint (bool) – Whether the model weights were loaded from a checkpoint (only applies to pretrained models)

  • checkpoint_index (int, optional) – The index of the checkpoint loaded (only applies to pretrained models).

  • checkpoint_label_type (str, optional) – Whether checkpoints are labelled by the number of steps or number of tokens.

  • checkpoint_value (int, optional) – The value of the checkpoint label (whether of steps or tokens).

  • tokenizer_name (str, optional) – the full name of the model, passed into HuggingFace to access the tokenizer. Only used when passing in custom config, if loading from pretrained then this is not needed.

  • use_local_attn (bool) – whether to use local attention - ie each destination token can only attend to source tokens a certain distance back.

  • window_size (int, optional) – the size of the window for local attention

  • attn_types (List[str], optional) – the types of attention to use for local attention

  • init_mode (str) – the initialization mode to use for the weights. Only relevant for custom models, ignored for pre-trained. We now support ‘gpt2’, ‘xavier_uniform’, ‘xavier_normal’, ‘kaiming_uniform’, ‘kaiming_normal’. MuP support to come. Defaults to ‘gpt2’.

  • normalization_type (str, optional) – the type of normalization to use. Options are None (no normalization), ‘LN’ (use LayerNorm, including weights & biases) and ‘LNPre’ (use LayerNorm, but no weights or biases), ‘RMS’ (use RMSNorm, including weights) and ‘RMSPre’ (use RMSNorm, but no weights or biases). Defaults to LN

  • device (str) – The device to use for the model. Defaults to ‘cuda’ if available, else ‘cpu’. Must be ‘cuda’ if n_devices > 1.

  • n_devices (int) – The number of devices to use for the model. Defaults to 1. Layers are loaded to support “pipeline parallelism”, where each device is responsible for a subset of the layers.

  • attention_dir (str) – Whether to use causal (aka unidirectional aka GPT-2 style) or bidirectional attention. Options are ‘causal’ and ‘bidirectional’. Defaults to ‘causal’

  • attn_only (bool) – Whether to only use attention layers, no feedforward layers. Defaults to False

  • seed (int, optional) – The seed to use for the model. Used to set sources of randomness (Python, PyTorch and NumPy) and to initialize weights. Defaults to None. We recommend setting a seed, so your experiments are reproducible.

  • initializer_range (float) – The standard deviation of the normal used to initialise the weights, initialized to 0.8 / sqrt(d_model). If init_mode is ‘xavier_uniform’ or ‘xavier_normal’, this value is instead treated as the gain parameter for the weight initialisation (a constant factor to scale the weights by). Defaults to -1.0, which means not set.

  • init_weights (bool) – Whether to initialize the weights. Defaults to True. If False, does not initialize weights.

  • scale_attn_by_inverse_layer_idx (bool) – Whether to scale the attention weights by 1/(layer_id+1), used by Mistral (Stanford) models for numerical stability when training in FP16. Defaults to False.

  • positional_embedding_type (str) – The positional embedding used. Options are ‘standard’ (ie GPT-2 style, absolute, randomly initialized learned positional embeddings, directly added to the residual stream), ‘rotary’ (described here: https://blog.eleuther.ai/rotary-embeddings/ ) and ‘shortformer’ (GPT-2 style absolute & learned, but rather than being added to the residual stream they’re only added to the inputs to the keys and the queries (ie key = W_K(res_stream + pos_embed), but values and MLPs don’t get any positional info)). Sinusoidal are not currently supported. Defaults to ‘standard’.

  • final_rms (bool) – Whether to replace the final normalization (just before the unembed) with RMSNorm (ie no centering or bias, just scaling + weights). Only included because of a dumb bug in my original SoLU code. Defaults to False.

  • d_vocab_out (int, optional) – The size of the output vocabulary. Defaults to -1, which means not set. If not set, will be equal to d_vocab. Mainly useful for algorithmic tasks where the input and output vocabularies may be different.

  • parallel_attn_mlp (bool) – Whether to parallelize the attention and MLP layers - a weird cursed thing done by GPT-J. Means that mlp_out=MLP(ln1(resid_pre)) and resid_post=resid_pre+attn_out+mlp_out. Defaults to False.

  • rotary_dim (int, optional) – The dimensionality of the rotary embeddings, may be d_head in which case only the first rotary_dim dimensions of each head are rotated. Defaults to None, if positional_embedding_type==”rotary” post-init then sets it to d_head, i.e. “rotate all dimensions of the query and key”.

  • n_params (int, optional) – The number of “hidden weight” parameters in the model, excluding embeddings, unembedding, biases, and layer norms. Counts only the attention projections (W_Q, W_K, W_V, W_O) and MLP weights (W_in, W_out, plus W_gate when gated_mlp=True). This matches the convention from the scaling laws paper, which found this to be the most meaningful number for predicting performance. Note: this is NOT the same as sum(p.numel() for p in model.parameters()) — that would include embeddings and biases and yield a larger number. Use the sum(p.numel() ...) form if you want the total parameter count (e.g. for memory-budget calculations). Automatically calculated; not intended to be set by the user.

  • use_hook_tokens (bool) – Will add a hook point on the token input to HookedTransformer.forward, which lets you cache or intervene on the tokens. Defaults to False.

  • gated_mlp (bool) – If True, the MLP layer uses a gated formulation (SwiGLU/GeGLU-style): mlp_out = W_out @ (act_fn(W_gate @ x) * (W_in @ x)), with an extra W_gate weight matrix alongside W_in and W_out. Used by LLaMA, Mistral, Gemma, Qwen and similar families. When False (default), the MLP is the plain mlp_out = W_out @ act_fn(W_in @ x) form. loading_from_pretrained sets this automatically per architecture; only set manually for a custom config.

  • default_prepend_bos (bool, optional) – Default behavior of whether to prepend the BOS token when the methods of HookedTransformer process input text to tokenize (only when input is a string). Defaults to True - even for models not explicitly trained with this, heads often use the first position as a resting position and accordingly lose information from the first token, so this empirically seems to give better results. To change the default behavior to False, pass in default_prepend_bos=False. Note that you can also locally override the default behavior by passing in prepend_bos=True/False when you call a method that processes the input string.

  • dtype (torch.dtype, optional) – The model’s dtype. Defaults to torch.float32.

  • tokenizer_prepends_bos (bool, optional) – This flag is set by set_tokenizer. It is set to True only when the tokenizer automatically prepends the BOS token if initialized with add_bos_token=True. We need this information to dynamically control bos prepending.

  • load_in_4bit (bool) – If this flag is set, then it’s assumed that parameters are 4-bit quantized with bitsandbytes. Currently only supported for Llama.

  • quantization_method (str, optional) – the quant_method declared by the checkpoint’s HF config (“mxfp4”, “bitsandbytes”, “gptq”, …), captured while that config is already in hand so later load steps need not refetch it. None when unquantized, and also when the config was never fetched (the llama/gemma name-based branches of convert_hf_model_config infer the architecture from the model name instead).

  • n_key_value_heads (int, optional) – The number of groups of heads that use the same key and value matrix. Only for models that use Grouped Query Attention.

  • post_embedding_ln (bool) – Whether to apply layer normalization after embedding the tokens. Defaults to False.

  • num_experts (int, optional) – The number of experts to use in the MoE layer. If set, experts_per_token must also be set. Set to None if not using MoE.

  • experts_per_token (int, optional) – The number of experts to use for each pass in the MoE layer. If set, num_experts must also be set. Set to None if not using MoE.

  • relative_attention_max_distance (int, optional) – The maximum distance between tokens for relative attention. If set, relative_attention_num_buckets must also be set.Only used in EncoderDecoder models, like T5.

  • relative_attention_num_buckets (int, optional) – The number of buckets to use for relative attention. If set, relative_attention_max_distance must also be set.Only used in EncoderDecoder models, like T5.

  • decoder_start_token_id (int, optional) – The start token id for the decoder. Only used in EncoderDecoder models, like T5.

  • 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.

  • use_normalization_before_and_after (bool) – Whether to apply normalization (LN/RMS/etc) to both the input of an attn/MLP block and the output (before adding back to the residual stream). Currently only used in Gemma-2. Defaults to False.

  • attn_scores_soft_cap (float) – An optional softcap for attention scores pre-softmax. If used, it will map attn_scores -> soft_cap * tanh(attn_scores / soft_cap). As tanh’s output is in [-1, 1], this maps attn_scores to [-soft_cap, soft_cap], with little effect on small values, but squashing large values into that interval. Currently only used in Gemma-2. Defaults to -1.0, which means not set.

  • output_logits_soft_cap (float) – An optional softcap for output logits, currently only used in Gemma-2 (see attn_scores_soft_cap for details). Defaults to -1.0, which means not set.

  • use_NTK_by_parts_rope (bool) – Whether to apply the “NTK-by-parts” method when using Rotary Positional Embedding. This method adjusts the interpolation based on frequency factors for different parts of the hidden dimensions. See Section 3.2 in https://arxiv.org/pdf/2309.00071 for details. Defaults to False.

  • NTK_by_parts_low_freq_factor (float) – The threshold applied to low-frequency hidden dimensions during interpolation when using the “NTK-by-parts” method. Defaults to 1.0.

  • NTK_by_parts_high_freq_factor (float) – The threshold applied to high-frequency hidden dimensions during interpolation in the “NTK-by-parts” method. Defaults to 4.0.

  • NTK_by_parts_factor (float) – The overall factor used in the “NTK-by-parts” method that affects the rate of change between low and high-frequency interpolation strategies. Defaults to 8.0.

  • use_yarn_rope (bool) – Whether to apply YARN (Yet Another RoPE extensioN) scaling to rotary positional embeddings. YARN blends interpolated and extrapolated frequencies per dimension using correction ranges. See https://arxiv.org/abs/2309.00071 for details. Used by OLMo 3. Defaults to False.

  • yarn_factor (float) – The interpolation factor for YARN RoPE scaling. Defaults to 1.0.

  • yarn_attention_factor (float) – Multiplicative scaling applied to sin/cos embeddings in YARN. Defaults to 1.0.

  • yarn_beta_fast (float) – Upper rotation threshold for YARN correction range. Defaults to 32.

  • yarn_beta_slow (float) – Lower rotation threshold for YARN correction range. Defaults to 1.

  • yarn_truncate (bool) – Whether to floor/ceil the YARN correction-range bounds (HF’s truncate). GPT-OSS ships truncate=False. Defaults to True.

  • yarn_global_attn_only (bool) – Whether YARN applies only to global-attention layers, with sliding/local layers keeping plain rope (Olmo-3’s per-layer-type rope). Defaults to False.

  • use_attention_sinks (bool) – Whether attention carries a learned per-head sink logit (GPT-OSS) that joins the softmax as an extra key column and is dropped afterward. Defaults to False.

  • yarn_original_max_position_embeddings (int) – The original max position embeddings before YARN extension. Defaults to 4096.

  • use_qk_norm (bool) – Whether to apply RMSNorm to the query and key projections before computing attention scores. Used by Gemma 3 models. Defaults to False.

  • rotary_base_local (float, optional) – The base for rotary positional embeddings in local attention layers. Used by models with hybrid local/global attention (e.g., Gemma 3) which use different RoPE bases for local (10k) and global (1M) attention. Defaults to None, which means the standard rotary_base is used for all layers.

  • norm_topk_prob (bool) – Whether to normalize the top-k probabilities in the MoE layer.

  • use_logn_attn (bool) – Qwen-1’s log-n attention: scale queries by log_{train_len}(position) past the training length (eval only).

  • train_seq_length (int, optional) – the length the model was trained at (Qwen-1’s seq_length). Both log-n scaling and dynamic-NTK RoPE threshold on it, and it stays fixed when n_ctx is overridden upward.

  • use_dynamic_ntk_rope (bool) – Qwen-1’s dynamic NTK: rescale the rotary base by alpha ** (rotary_dim / (rotary_dim - 2)) once the key length exceeds the training length (eval only).

  • clip_qkv (float, optional) – Clamp Q/K/V activations to [-clip_qkv, clip_qkv] after projection (and any qk-norm), as OLMo v1 and OLMoE do. Defaults to None (no clamp).

NTK_by_parts_factor: float = 8.0
NTK_by_parts_high_freq_factor: float = 4.0
NTK_by_parts_low_freq_factor: float = 1.0
NTK_original_ctx_len: int = 8192
act_fn: str = 'relu'
attention_dir: str = 'causal'
attn_only: bool = False
attn_scale: float = -1.0
attn_scores_soft_cap: float = -1.0
attn_types: List | None = None
checkpoint_index: int | None = None
checkpoint_label_type: str | None = None
checkpoint_value: int | None = None
clip_qkv: float | None = None
d_vocab_out: int = -1
decoder_start_token_id: int | None = None
dtype: dtype = torch.float32
eps: float = 1e-05
experts_per_token: int | None = None
final_rms: bool = False
from_checkpoint: bool = False
classmethod from_dict(config_dict: Dict[str, Any]) → HookedTransformerConfig

Instantiates a HookedTransformerConfig from a Python dictionary of parameters.

gated_mlp: bool = False
init_mode: str = 'gpt2'
init_weights: bool = True
initializer_range: float = -1.0
is_layer_norm_activation() → bool
load_in_4bit: bool = False
model_name: str = 'custom'
n_devices: int = 1
n_params: int | None = None
norm_topk_prob: bool = False
normalization_type: str | None = 'LN'
num_experts: int | None = None
original_architecture: str | None = None
output_logits_soft_cap: float = -1.0
parallel_attn_mlp: bool = False
post_embedding_ln: bool = False
quantization_method: str | None = None
relative_attention_max_distance: int | None = None
relative_attention_num_buckets: int | None = None
rotary_adjacent_pairs: bool = False
rotary_base: float | int = 10000
rotary_base_local: float | int | None = None
rotary_dim: int | None = None
rotary_scaling_factor: float = 1.0
scale_attn_by_inverse_layer_idx: bool = False
seed: int | None = None
set_seed_everywhere(seed: int)
tie_word_embeddings: bool = False
to_dict()

Convert the config to a dictionary.

tokenizer_name: str | None = None
tokenizer_prepends_bos: bool | None = None
train_seq_length: int | None = None
trust_remote_code: bool = False
ungroup_grouped_query_attention: bool = False
classmethod unwrap(config: Dict | TransformerLensConfig) → HookedTransformerConfig

Convenience function to avoid duplicate code from a common way config is passed to various components

use_NTK_by_parts_rope: bool = False
use_attention_sinks: bool = False
use_attn_in: bool = False
use_attn_scale: bool = True
use_dynamic_ntk_rope: bool = False
use_hook_mlp_in: bool = False
use_hook_tokens: bool = False
use_local_attn: bool = False
use_logn_attn: bool = False
use_normalization_before_and_after: bool = False
use_qk_norm: bool = False
use_yarn_rope: bool = False
window_size: int | None = None
yarn_attention_factor: float = 1.0
yarn_beta_fast: float = 32.0
yarn_beta_slow: float = 1.0
yarn_factor: float = 1.0
yarn_global_attn_only: bool = False
yarn_original_max_position_embeddings: int = 4096
yarn_truncate: bool = True