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