Coverage for transformer_lens/config/transformer_bridge_config.py: 96%
124 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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 qk_norm_after_rope: bool = False,
87 rotary_interleaved_cos_sin: bool = False,
88 load_in_4bit: bool = False,
89 num_experts: Optional[int] = None,
90 experts_per_token: Optional[int] = None,
91 n_key_value_heads: Optional[int] = None,
92 # Heterogeneous attention geometry (e.g. Gemma 4): per-layer values when
93 # they vary across layers; d_head / n_key_value_heads then hold the
94 # majority-layer scalar and attention math is delegated to HF.
95 per_layer_head_dim: Optional[list] = None,
96 per_layer_num_key_value_heads: Optional[list] = None,
97 relative_attention_max_distance: Optional[int] = None,
98 relative_attention_num_buckets: Optional[int] = None,
99 decoder_start_token_id: Optional[int] = None,
100 scale_embedding: Optional[bool] = None,
101 tie_word_embeddings: bool = False,
102 use_normalization_before_and_after: bool = False,
103 attn_scores_soft_cap: float = SOFTCAP_DISABLED,
104 output_logits_soft_cap: float = SOFTCAP_DISABLED,
105 use_NTK_by_parts_rope: bool = False,
106 NTK_by_parts_low_freq_factor: float = 1.0,
107 NTK_by_parts_high_freq_factor: float = 4.0,
108 NTK_by_parts_factor: float = 8.0,
109 rmsnorm_uses_offset: bool = False,
110 attn_implementation: Optional[str] = None,
111 # Audio model configuration
112 is_audio_model: bool = False,
113 # Vision model (ViT, DeiT) configuration
114 is_visual_model: bool = False,
115 # Stateful model configuration (e.g., Mamba SSMs use cache_params,
116 # not past_key_values, so generation delegates to hf_generate)
117 is_stateful: bool = False,
118 # Multimodal configuration
119 is_multimodal: bool = False,
120 vision_hidden_size: Optional[int] = None,
121 vision_num_layers: Optional[int] = None,
122 vision_num_heads: Optional[int] = None,
123 mm_tokens_per_image: Optional[int] = None,
124 **kwargs,
125 ):
126 """Initialize TransformerBridgeConfig."""
127 object.__setattr__(self, "_bridge_ref", None)
128 super().__init__(
129 d_model=d_model,
130 d_head=d_head,
131 n_layers=n_layers,
132 n_ctx=n_ctx,
133 d_vocab=d_vocab,
134 n_heads=n_heads,
135 **kwargs,
136 )
138 # Architecture information for adapter selection
139 self.architecture = architecture
141 # Tokenizer configuration
142 self.tokenizer_prepends_bos = tokenizer_prepends_bos
143 self.tokenizer_appends_eos = tokenizer_appends_eos
144 self.default_padding_side = default_padding_side
146 # Attention weight processing configuration
147 self.split_attention_weights = False
149 # HookedTransformerConfig compatibility fields
150 self.model_name = model_name
151 self.act_fn = act_fn
152 self.eps = eps
153 self.use_attn_scale = use_attn_scale
154 self.attn_scale = attn_scale
155 self.use_hook_mlp_in = use_hook_mlp_in
156 self.use_attn_in = use_attn_in
157 self.use_qk_norm = use_qk_norm
158 self.use_local_attn = use_local_attn
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.qk_norm_after_rope = qk_norm_after_rope
189 self.rotary_interleaved_cos_sin = rotary_interleaved_cos_sin
190 self.load_in_4bit = load_in_4bit
191 self.num_experts = num_experts
192 self.experts_per_token = experts_per_token
193 self.n_key_value_heads = n_key_value_heads
194 self.per_layer_head_dim = per_layer_head_dim
195 self.per_layer_num_key_value_heads = per_layer_num_key_value_heads
196 self.relative_attention_max_distance = relative_attention_max_distance
197 self.relative_attention_num_buckets = relative_attention_num_buckets
198 self.decoder_start_token_id = decoder_start_token_id
199 # Seq2seq families (Bart, Marian, Pegasus, ...) scale token embeddings
200 # by sqrt(d_model) in the HF forward when set.
201 self.scale_embedding = scale_embedding
202 self.tie_word_embeddings = tie_word_embeddings
203 self.use_normalization_before_and_after = use_normalization_before_and_after
204 self.attn_scores_soft_cap = attn_scores_soft_cap
205 self.output_logits_soft_cap = output_logits_soft_cap
206 self.use_NTK_by_parts_rope = use_NTK_by_parts_rope
207 self.NTK_by_parts_low_freq_factor = NTK_by_parts_low_freq_factor
208 self.NTK_by_parts_high_freq_factor = NTK_by_parts_high_freq_factor
209 self.NTK_by_parts_factor = NTK_by_parts_factor
210 self.rmsnorm_uses_offset = rmsnorm_uses_offset
211 self.attn_implementation = attn_implementation
212 # Audio model configuration
213 self.is_audio_model = is_audio_model
214 # Vision model (ViT, DeiT) configuration
215 self.is_visual_model = is_visual_model
216 # Stateful model configuration
217 self.is_stateful = is_stateful
218 # Multimodal configuration
219 self.is_multimodal = is_multimodal
220 self.vision_hidden_size = vision_hidden_size
221 self.vision_num_layers = vision_num_layers
222 self.vision_num_heads = vision_num_heads
223 self.mm_tokens_per_image = mm_tokens_per_image
224 self.__post_init__()
226 def __setattr__(self, name: str, value: Any) -> None:
227 """Route live Bridge hook-flag assignments through their public setters."""
228 if name in self._BRIDGE_MANAGED_HOOK_FLAGS:
229 bridge_ref = getattr(self, "_bridge_ref", None)
230 bridge = bridge_ref() if bridge_ref is not None else None
231 if bridge is not None:
232 getattr(bridge, f"set_{name}")(value)
233 return
234 super().__setattr__(name, value)
236 def __getstate__(self) -> dict[str, Any]:
237 """Serialize config data without retaining its live Bridge binding."""
238 return self.__dict__.copy()
240 def __setstate__(self, state: dict[str, Any]) -> None:
241 """Restore an unbound config copy."""
242 self.__dict__.update(state)
243 object.__setattr__(self, "_bridge_ref", None)
245 def _bind_bridge(self, bridge: Any) -> None:
246 """Bind runtime hook-flag assignments to a constructed Bridge."""
247 bridge_ref = getattr(self, "_bridge_ref", None)
248 bound_bridge = bridge_ref() if bridge_ref is not None else None
249 if bound_bridge is None:
250 object.__setattr__(self, "_bridge_ref", weakref.ref(bridge))
251 elif bound_bridge is not bridge: 251 ↛ exitline 251 didn't return from function '_bind_bridge' because the condition on line 251 was always true
252 warnings.warn(
253 "TransformerBridgeConfig is already bound to another live "
254 "TransformerBridge; declining to bind it to this instance. "
255 "Direct assignments to Bridge-managed hook flags will continue "
256 "to configure the existing TransformerBridge.",
257 stacklevel=3,
258 )
260 def _set_bridge_managed_hook_flag(self, name: str, value: bool) -> None:
261 """Set a managed flag without re-entering the Bridge setter."""
262 if name not in self._BRIDGE_MANAGED_HOOK_FLAGS: 262 ↛ 263line 262 didn't jump to line 263 because the condition on line 262 was never true
263 raise ValueError(f"Unknown Bridge-managed hook flag: {name}")
264 object.__setattr__(self, name, value)
266 def __post_init__(self):
267 """Post-initialization processing."""
268 # dtype is guaranteed to be set at this point
270 # Validate architecture if provided before calling super()
271 if ( 271 ↛ 276line 271 didn't jump to line 276 because the condition on line 271 was never true
272 hasattr(self, "architecture")
273 and self.architecture is not None
274 and not isinstance(self.architecture, str)
275 ):
276 raise ValueError(f"architecture must be a string, got {type(self.architecture)}")
278 # Resolve the initializer_range sentinel (-1.0 means "not set by the user").
279 # Same rule the legacy HookedTransformerConfig applied.
280 # Guarded with getattr: this method also runs once from the dataclass
281 # parent's __init__, before self.initializer_range is assigned below.
282 if getattr(self, "initializer_range", None) is not None:
283 if self.initializer_range < 0 and self.init_mode == "gpt2":
284 # Roughly copy the GPT-2 value, but proportional to sqrt(1/d_model)
285 self.initializer_range = 0.8 / np.sqrt(self.d_model)
286 if self.initializer_range < 0 and self.init_mode != "gpt2":
287 # This is the gain parameter for the weight initialisation
288 self.initializer_range = 1.0
290 # Call parent's __post_init__ after our validation
291 if hasattr(super(), "__post_init__"): 291 ↛ exitline 291 didn't return from function '__post_init__' because the condition on line 291 was always true
292 super().__post_init__()
294 @property
295 def head_dim(self) -> int:
296 """Alias for d_head to match HuggingFace config naming convention."""
297 return self.d_head