Coverage for transformer_lens/config/transformer_bridge_config.py: 96%
122 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Configuration class for TransformerBridge."""
3import warnings
4import weakref
5from typing import Any, Optional
7import numpy as np
8import torch
10from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED
12from .transformer_lens_config import TransformerLensConfig
15class TransformerBridgeConfig(TransformerLensConfig):
16 """
17 Configuration for TransformerBridge.
19 This extends TransformerLensConfig with bridge-specific properties,
20 particularly architecture information needed for adapter selection.
21 Also includes all HookedTransformerConfig fields for compatibility.
22 """
24 __slots__ = ("_bridge_ref",)
26 _BRIDGE_MANAGED_HOOK_FLAGS = frozenset(
27 {
28 "use_attn_result",
29 "use_attn_in",
30 "use_hook_mlp_in",
31 "use_split_qkv_input",
32 }
33 )
35 def __init__(
36 self,
37 d_model: int,
38 d_head: int,
39 n_layers: int,
40 n_ctx: int,
41 n_heads: int = -1, # Add n_heads to signature so it's not filtered out by from_dict
42 d_vocab: int = -1,
43 architecture: Optional[str] = None,
44 tokenizer_prepends_bos: bool = True,
45 tokenizer_appends_eos: bool = False,
46 default_padding_side: Optional[str] = None,
47 # HookedTransformerConfig compatibility fields
48 model_name: str = "custom",
49 act_fn: str = "relu",
50 eps: float = 1e-5,
51 use_attn_scale: bool = True,
52 attn_scale: float = -1.0,
53 use_hook_mlp_in: bool = False,
54 use_attn_in: bool = False,
55 use_qk_norm: bool = False,
56 use_local_attn: bool = False,
57 original_architecture: Optional[str] = None,
58 from_checkpoint: bool = False,
59 checkpoint_index: Optional[int] = None,
60 checkpoint_label_type: Optional[str] = None,
61 checkpoint_value: Optional[int] = None,
62 tokenizer_name: Optional[str] = None,
63 window_size: Optional[int] = None,
64 attn_types: Optional[list] = None,
65 init_mode: str = "gpt2",
66 normalization_type: Optional[str] = "LN",
67 n_devices: int = 1,
68 attention_dir: str = "causal",
69 attn_only: bool = False,
70 seed: Optional[int] = None,
71 initializer_range: float = -1.0,
72 init_weights: bool = True,
73 scale_attn_by_inverse_layer_idx: bool = False,
74 final_rms: bool = False,
75 d_vocab_out: int = -1,
76 parallel_attn_mlp: bool = False,
77 rotary_dim: Optional[int] = None,
78 n_params: Optional[int] = None,
79 use_hook_tokens: bool = False,
80 gated_mlp: bool = False,
81 dtype: Optional[torch.dtype] = torch.float32,
82 post_embedding_ln: bool = False,
83 rotary_base: int | float = 10000,
84 trust_remote_code: bool = False,
85 rotary_adjacent_pairs: bool = False,
86 load_in_4bit: bool = False,
87 num_experts: Optional[int] = None,
88 experts_per_token: Optional[int] = None,
89 n_key_value_heads: Optional[int] = None,
90 # Heterogeneous attention geometry (e.g. Gemma 4): per-layer values when
91 # they vary across layers; d_head / n_key_value_heads then hold the
92 # majority-layer scalar and attention math is delegated to HF.
93 per_layer_head_dim: Optional[list] = None,
94 per_layer_num_key_value_heads: Optional[list] = None,
95 relative_attention_max_distance: Optional[int] = None,
96 relative_attention_num_buckets: Optional[int] = None,
97 decoder_start_token_id: Optional[int] = None,
98 scale_embedding: Optional[bool] = None,
99 tie_word_embeddings: bool = False,
100 use_normalization_before_and_after: bool = False,
101 attn_scores_soft_cap: float = SOFTCAP_DISABLED,
102 output_logits_soft_cap: float = SOFTCAP_DISABLED,
103 use_NTK_by_parts_rope: bool = False,
104 NTK_by_parts_low_freq_factor: float = 1.0,
105 NTK_by_parts_high_freq_factor: float = 4.0,
106 NTK_by_parts_factor: float = 8.0,
107 rmsnorm_uses_offset: bool = False,
108 attn_implementation: Optional[str] = None,
109 # Audio model configuration
110 is_audio_model: bool = False,
111 # Vision model (ViT, DeiT) configuration
112 is_visual_model: bool = False,
113 # Stateful model configuration (e.g., Mamba SSMs use cache_params,
114 # not past_key_values, so generation delegates to hf_generate)
115 is_stateful: bool = False,
116 # Multimodal configuration
117 is_multimodal: bool = False,
118 vision_hidden_size: Optional[int] = None,
119 vision_num_layers: Optional[int] = None,
120 vision_num_heads: Optional[int] = None,
121 mm_tokens_per_image: Optional[int] = None,
122 **kwargs,
123 ):
124 """Initialize TransformerBridgeConfig."""
125 object.__setattr__(self, "_bridge_ref", None)
126 super().__init__(
127 d_model=d_model,
128 d_head=d_head,
129 n_layers=n_layers,
130 n_ctx=n_ctx,
131 d_vocab=d_vocab,
132 n_heads=n_heads,
133 **kwargs,
134 )
136 # Architecture information for adapter selection
137 self.architecture = architecture
139 # Tokenizer configuration
140 self.tokenizer_prepends_bos = tokenizer_prepends_bos
141 self.tokenizer_appends_eos = tokenizer_appends_eos
142 self.default_padding_side = default_padding_side
144 # Attention weight processing configuration
145 self.split_attention_weights = False
147 # HookedTransformerConfig compatibility fields
148 self.model_name = model_name
149 self.act_fn = act_fn
150 self.eps = eps
151 self.use_attn_scale = use_attn_scale
152 self.attn_scale = attn_scale
153 self.use_hook_mlp_in = use_hook_mlp_in
154 self.use_attn_in = use_attn_in
155 self.use_qk_norm = use_qk_norm
156 self.use_local_attn = use_local_attn
157 self.original_architecture = original_architecture
158 self.from_checkpoint = from_checkpoint
159 self.checkpoint_index = checkpoint_index
160 self.checkpoint_label_type = checkpoint_label_type
161 self.checkpoint_value = checkpoint_value
162 self.tokenizer_name = tokenizer_name
163 self.window_size = window_size
164 self.attn_types = attn_types
165 self.init_mode = init_mode
166 self.normalization_type = normalization_type
167 self.n_devices = n_devices
168 self.attention_dir = attention_dir
169 self.attn_only = attn_only
170 self.seed = seed
171 self.initializer_range = initializer_range
172 self.init_weights = init_weights
173 self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
174 self.final_rms = final_rms
175 self.d_vocab_out = d_vocab_out
176 self.parallel_attn_mlp = parallel_attn_mlp
177 self.rotary_dim = rotary_dim
178 self.n_params = n_params
179 self.use_hook_tokens = use_hook_tokens
180 self.gated_mlp = gated_mlp
181 self.dtype = dtype if dtype is not None else torch.float32
182 self.post_embedding_ln = post_embedding_ln
183 self.rotary_base = int(rotary_base)
184 self.trust_remote_code = trust_remote_code
185 self.rotary_adjacent_pairs = rotary_adjacent_pairs
186 self.load_in_4bit = load_in_4bit
187 self.num_experts = num_experts
188 self.experts_per_token = experts_per_token
189 self.n_key_value_heads = n_key_value_heads
190 self.per_layer_head_dim = per_layer_head_dim
191 self.per_layer_num_key_value_heads = per_layer_num_key_value_heads
192 self.relative_attention_max_distance = relative_attention_max_distance
193 self.relative_attention_num_buckets = relative_attention_num_buckets
194 self.decoder_start_token_id = decoder_start_token_id
195 # Seq2seq families (Bart, Marian, Pegasus, ...) scale token embeddings
196 # by sqrt(d_model) in the HF forward when set.
197 self.scale_embedding = scale_embedding
198 self.tie_word_embeddings = tie_word_embeddings
199 self.use_normalization_before_and_after = use_normalization_before_and_after
200 self.attn_scores_soft_cap = attn_scores_soft_cap
201 self.output_logits_soft_cap = output_logits_soft_cap
202 self.use_NTK_by_parts_rope = use_NTK_by_parts_rope
203 self.NTK_by_parts_low_freq_factor = NTK_by_parts_low_freq_factor
204 self.NTK_by_parts_high_freq_factor = NTK_by_parts_high_freq_factor
205 self.NTK_by_parts_factor = NTK_by_parts_factor
206 self.rmsnorm_uses_offset = rmsnorm_uses_offset
207 self.attn_implementation = attn_implementation
208 # Audio model configuration
209 self.is_audio_model = is_audio_model
210 # Vision model (ViT, DeiT) configuration
211 self.is_visual_model = is_visual_model
212 # Stateful model configuration
213 self.is_stateful = is_stateful
214 # Multimodal configuration
215 self.is_multimodal = is_multimodal
216 self.vision_hidden_size = vision_hidden_size
217 self.vision_num_layers = vision_num_layers
218 self.vision_num_heads = vision_num_heads
219 self.mm_tokens_per_image = mm_tokens_per_image
220 self.__post_init__()
222 def __setattr__(self, name: str, value: Any) -> None:
223 """Route live Bridge hook-flag assignments through their public setters."""
224 if name in self._BRIDGE_MANAGED_HOOK_FLAGS:
225 bridge_ref = getattr(self, "_bridge_ref", None)
226 bridge = bridge_ref() if bridge_ref is not None else None
227 if bridge is not None:
228 getattr(bridge, f"set_{name}")(value)
229 return
230 super().__setattr__(name, value)
232 def __getstate__(self) -> dict[str, Any]:
233 """Serialize config data without retaining its live Bridge binding."""
234 return self.__dict__.copy()
236 def __setstate__(self, state: dict[str, Any]) -> None:
237 """Restore an unbound config copy."""
238 self.__dict__.update(state)
239 object.__setattr__(self, "_bridge_ref", None)
241 def _bind_bridge(self, bridge: Any) -> None:
242 """Bind runtime hook-flag assignments to a constructed Bridge."""
243 bridge_ref = getattr(self, "_bridge_ref", None)
244 bound_bridge = bridge_ref() if bridge_ref is not None else None
245 if bound_bridge is None:
246 object.__setattr__(self, "_bridge_ref", weakref.ref(bridge))
247 elif bound_bridge is not bridge: 247 ↛ exitline 247 didn't return from function '_bind_bridge' because the condition on line 247 was always true
248 warnings.warn(
249 "TransformerBridgeConfig is already bound to another live "
250 "TransformerBridge; declining to bind it to this instance. "
251 "Direct assignments to Bridge-managed hook flags will continue "
252 "to configure the existing TransformerBridge.",
253 stacklevel=3,
254 )
256 def _set_bridge_managed_hook_flag(self, name: str, value: bool) -> None:
257 """Set a managed flag without re-entering the Bridge setter."""
258 if name not in self._BRIDGE_MANAGED_HOOK_FLAGS: 258 ↛ 259line 258 didn't jump to line 259 because the condition on line 258 was never true
259 raise ValueError(f"Unknown Bridge-managed hook flag: {name}")
260 object.__setattr__(self, name, value)
262 def __post_init__(self):
263 """Post-initialization processing."""
264 # dtype is guaranteed to be set at this point
266 # Validate architecture if provided before calling super()
267 if ( 267 ↛ 272line 267 didn't jump to line 272 because the condition on line 267 was never true
268 hasattr(self, "architecture")
269 and self.architecture is not None
270 and not isinstance(self.architecture, str)
271 ):
272 raise ValueError(f"architecture must be a string, got {type(self.architecture)}")
274 # Resolve the initializer_range sentinel (-1.0 means "not set by the user").
275 # Same rule the legacy HookedTransformerConfig applied.
276 # Guarded with getattr: this method also runs once from the dataclass
277 # parent's __init__, before self.initializer_range is assigned below.
278 if getattr(self, "initializer_range", None) is not None:
279 if self.initializer_range < 0 and self.init_mode == "gpt2":
280 # Roughly copy the GPT-2 value, but proportional to sqrt(1/d_model)
281 self.initializer_range = 0.8 / np.sqrt(self.d_model)
282 if self.initializer_range < 0 and self.init_mode != "gpt2":
283 # This is the gain parameter for the weight initialisation
284 self.initializer_range = 1.0
286 # Call parent's __post_init__ after our validation
287 if hasattr(super(), "__post_init__"): 287 ↛ exitline 287 didn't return from function '__post_init__' because the condition on line 287 was always true
288 super().__post_init__()
290 @property
291 def head_dim(self) -> int:
292 """Alias for d_head to match HuggingFace config naming convention."""
293 return self.d_head