Coverage for transformer_lens/model_bridge/sources/_bridge_builder.py: 88%
94 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Loader-agnostic helpers for building a TransformerBridge around a pre-loaded model."""
2from __future__ import annotations
4import copy
5from typing import Any, Callable, Optional
7import torch
8from torch import nn
10from transformer_lens.config import TransformerBridgeConfig
11from transformer_lens.factories.architecture_adapter_factory import (
12 ArchitectureAdapterFactory,
13)
14from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
15from transformer_lens.model_bridge.bridge import TransformerBridge
16from transformer_lens.utilities.heterogeneous_config import (
17 het_safe_view,
18 per_layer_attr_names,
19 safe_config_get,
20)
22# Architecture-agnostic; do not extend per-architecture.
23_HF_PASSTHROUGH_ATTRS = [
24 # OPT
25 "is_gated_act",
26 # LongT5
27 "encoder_attention_type",
28 "word_embed_proj_dim",
29 "do_layer_norm_before",
30 # BART
31 "encoder_layers",
32 "decoder_layers",
33 "encoder_attention_heads",
34 "decoder_attention_heads",
35 "encoder_ffn_dim",
36 "decoder_ffn_dim",
37 # Marian
38 "scale_embedding",
39 # Granite
40 "position_embedding_type",
41 "logits_scaling",
42 "residual_multiplier",
43 # Falcon
44 "parallel_attn",
45 "multi_query",
46 "new_decoder_architecture",
47 "alibi",
48 "num_ln_in_parallel_attn",
49 # GPTNeoX
50 "use_parallel_residual",
51 # Mamba (SSM config)
52 "state_size",
53 "conv_kernel",
54 "expand",
55 "time_step_rank",
56 "intermediate_size",
57 # Mamba-2 (additional SSM config)
58 "n_groups",
59 "chunk_size",
60 # Falcon-H1 (parallel attn + Mamba-2 hybrid SSM config)
61 "mamba_d_ssm",
62 "mamba_n_heads",
63 "mamba_d_head",
64 "mamba_d_state",
65 "mamba_n_groups",
66 "mamba_d_conv",
67 "mamba_chunk_size",
68 "lm_head_multiplier",
69 # Multimodal
70 "vision_config",
71 # Cohere
72 "logit_scale",
73 "rope_parameters",
74 # HRM-Text
75 "H_cycles",
76 "L_cycles",
77 "L_bp_cycles",
78 "embedding_scale",
79 "prefix_lm",
80 "num_layers_per_stack",
81 "sliding_window_pattern",
82 "_sliding_window_pattern",
83 # Hybrid/MoE architectures
84 "layer_types",
85 "moe_intermediate_size",
86 "shared_expert_intermediate_size",
87 "norm_eps",
88 "attention_bias",
89 "lm_head_bias",
90 "router_jitter_noise",
91 "input_jitter_noise",
92 "eos_token_id",
93 # LLaDA remote-code model contract and tokenizer metadata
94 "block_type",
95 "block_group_size",
96 "rope",
97 "rope_full_precision",
98 "attention_layer_norm",
99 "include_bias",
100 "include_qkv_bias",
101 "scale_logits",
102 "input_emb_norm",
103 "layer_norm_type",
104 "embedding_size",
105 "mask_token_id",
106 "pad_token_id",
107 "bos_token_id",
108 # BD3LM
109 "model_length",
110 "block_size",
111 "cond_dim",
112 "adaln",
113 "cross_attn",
114 # Zamba2 (Mamba-2 + shared-attention hybrid)
115 "mamba_expand",
116 "mamba_ngroups",
117 "num_mem_blocks",
118 "layers_block_type",
119 "use_shared_attention_adapter",
120 # Jamba (attention + Mamba-1 hybrid; MoE schedule knobs)
121 "mamba_dt_rank",
122 "attn_layer_period",
123 "attn_layer_offset",
124 "expert_layer_period",
125 "expert_layer_offset",
126 # Ouro (LoopLM)
127 "total_ut_steps",
128 "early_exit_threshold",
129 # DeepSeek V4 (mHC + compressed attention)
130 "compress_rates",
131 "compress_rope_theta",
132 "hc_mult",
133 "hc_sinkhorn_iters",
134 "hc_eps",
135 "mlp_layer_types",
136 "swiglu_limit",
137 "o_groups",
138 "o_lora_rank",
139 "index_n_heads",
140 "index_head_dim",
141 "index_topk",
142 "q_lora_rank",
143 # Raven / Huginn (depth-recurrent)
144 "mean_recurrence",
145 "mean_backprop_depth",
146 "n_layers_in_prelude",
147 "n_layers_in_recurrent_block",
148 "n_layers_in_coda",
149 "injection_type",
150 "qk_bias",
151 # RWKV-7 (attention-free recurrent, generalized delta-rule time-mixing).
152 # head_dim is intentionally omitted: it is a read-only alias of d_head on
153 # TransformerBridgeConfig, so a passthrough setattr would raise.
154 "num_heads",
155 "value_dim",
156 "decay_low_rank_dim",
157 "gate_low_rank_dim",
158 "a_low_rank_dim",
159 "v_low_rank_dim",
160 "norm_first",
161 "norm_bias",
162 "fuse_norm",
163 "attn_mode",
164 "hidden_act",
165]
168def build_bridge_config_from_hf(
169 hf_config: Any,
170 architecture: str,
171 model_name: str,
172 dtype: torch.dtype,
173) -> TransformerBridgeConfig:
174 """Translate an HF config into a :class:`TransformerBridgeConfig`."""
175 from transformer_lens.model_bridge.sources.transformers import (
176 get_effective_text_config,
177 map_default_transformer_lens_config,
178 )
180 tl_config = map_default_transformer_lens_config(hf_config)
181 config_dict = dict(tl_config.__dict__)
182 # HF's attribute_map remaps num_experts → num_local_experts; restore the TL name.
183 if "num_local_experts" in config_dict and "num_experts" not in config_dict:
184 config_dict["num_experts"] = config_dict["num_local_experts"]
185 bridge_config = TransformerBridgeConfig.from_dict(config_dict)
186 bridge_config.architecture = architecture
187 bridge_config.model_name = model_name
188 bridge_config.dtype = dtype
190 effective_config = get_effective_text_config(hf_config)
191 # Per-layer-registered attrs would raise on global access (transformers>=5.15).
192 _het_attrs = per_layer_attr_names(effective_config) | per_layer_attr_names(hf_config)
193 for attr in _HF_PASSTHROUGH_ATTRS:
194 if attr in _het_attrs:
195 continue
196 val = getattr(effective_config, attr, None)
197 if val is None and effective_config is not hf_config:
198 val = getattr(hf_config, attr, None)
199 if val is not None:
200 setattr(bridge_config, attr, val)
202 # Gemma2: HF softcap field names differ from TL's. Read through the het
203 # view: a per-layer-registered field raises (not AttributeError) on raw
204 # getattr, so the default would not save us.
205 effective_config = het_safe_view(effective_config)
206 final_logit_softcapping = getattr(effective_config, "final_logit_softcapping", None)
207 if final_logit_softcapping is not None:
208 bridge_config.output_logits_soft_cap = float(final_logit_softcapping)
209 logits_soft_cap = getattr(effective_config, "logits_soft_cap", None)
210 if logits_soft_cap is not None:
211 bridge_config.output_logits_soft_cap = float(logits_soft_cap)
212 attn_logit_softcapping = getattr(effective_config, "attn_logit_softcapping", None)
213 if attn_logit_softcapping is not None:
214 bridge_config.attn_scores_soft_cap = float(attn_logit_softcapping)
216 # Nested encoder sub-configs (T5Gemma family): n_heads/n_key_value_heads are
217 # decoder-effective, so expose encoder head counts for per-side conversions.
218 # T5Gemma2 nests them one level deeper (encoder.text_config).
219 encoder_subconfig = safe_config_get(hf_config, "encoder")
220 if encoder_subconfig is not None:
221 encoder_subconfig = het_safe_view(encoder_subconfig)
222 if encoder_subconfig is not None and not hasattr(encoder_subconfig, "num_attention_heads"): 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 encoder_subconfig = safe_config_get(encoder_subconfig, "text_config")
224 if encoder_subconfig is not None:
225 encoder_subconfig = het_safe_view(encoder_subconfig)
226 if encoder_subconfig is not None:
227 enc_heads = getattr(encoder_subconfig, "num_attention_heads", None)
228 if enc_heads is not None: 228 ↛ 230line 228 didn't jump to line 230 because the condition on line 228 was always true
229 bridge_config.encoder_attention_heads = enc_heads
230 enc_kv = getattr(encoder_subconfig, "num_key_value_heads", None)
231 if enc_kv is not None: 231 ↛ 234line 231 didn't jump to line 234 because the condition on line 231 was always true
232 bridge_config.encoder_key_value_heads = enc_kv
234 return bridge_config
237def detect_tokenizer_bos_eos(tokenizer: Any) -> tuple[bool, bool]:
238 """Detect whether the tokenizer prepends BOS and/or appends EOS.
240 Non-empty test string — "" is unreliable with token aliasing.
241 """
242 encoded_test = tokenizer.encode("a")
243 prepends_bos = (
244 len(encoded_test) > 1
245 and tokenizer.bos_token_id is not None
246 and encoded_test[0] == tokenizer.bos_token_id
247 )
248 appends_eos = (
249 len(encoded_test) > 1
250 and tokenizer.eos_token_id is not None
251 and encoded_test[-1] == tokenizer.eos_token_id
252 )
253 return prepends_bos, appends_eos
256def build_bridge_from_module(
257 model: nn.Module,
258 architecture: str,
259 *,
260 hf_config: Optional[Any] = None,
261 tl_config: Optional[TransformerBridgeConfig] = None,
262 tokenizer: Optional[Any] = None,
263 dtype: Optional[torch.dtype] = None,
264 device: Optional[Any] = None,
265 model_name: str = "external",
266 post_adapter_hook: Optional[Callable[[ArchitectureAdapter], None]] = None,
267) -> TransformerBridge:
268 """Build a :class:`TransformerBridge` around a pre-loaded model.
270 The bridge never moves, casts, or mutates the supplied model.
272 Args:
273 model: Any ``nn.Module`` whose submodule tree matches the adapter's
274 expected dot-paths for ``architecture``.
275 architecture: Architecture identifier registered in the
276 ``ArchitectureAdapterFactory`` (e.g. ``"LlamaForCausalLM"``,
277 ``"TransformerLensNative"``).
278 hf_config: Optional HF-style config; translated via
279 :func:`build_bridge_config_from_hf`. Mutually exclusive with ``tl_config``.
280 tl_config: Optional pre-built :class:`TransformerBridgeConfig`; bypasses
281 HF translation. Mutually exclusive with ``hf_config``.
282 tokenizer: Optional tokenizer. If supplied, passes through
283 ``setup_tokenizer`` and detects BOS/EOS behavior.
284 dtype: Recorded on ``cfg.dtype``. Default ``None`` reads from the model's
285 first parameter; explicit values override.
286 device: Recorded on ``cfg.device``. Default ``None`` reads from the
287 model's first parameter.
288 model_name: Recorded on ``cfg.model_name``.
289 post_adapter_hook: Optional callback invoked after adapter selection and
290 before :meth:`adapter.prepare_model`. Source-specific overlays mutate
291 ``component_mapping`` here.
293 Returns:
294 A :class:`TransformerBridge` wrapping the supplied model.
295 """
296 if hf_config is None and tl_config is None:
297 raise ValueError(
298 "build_bridge_from_module requires exactly one of hf_config or "
299 "tl_config — the bridge needs config fields (d_model, n_heads, "
300 "n_layers, ...) that can't be inferred from the model alone."
301 )
302 if hf_config is not None and tl_config is not None:
303 raise ValueError(
304 "build_bridge_from_module got both hf_config and tl_config; supply "
305 "exactly one. hf_config triggers HF→bridge translation; tl_config "
306 "bypasses it."
307 )
309 # Reading dtype from the model avoids silently lying about a bf16 model.
310 if dtype is None:
311 try:
312 dtype = next(model.parameters()).dtype
313 except StopIteration:
314 dtype = torch.float32
316 if tl_config is not None:
317 # Defensive copy so adapter-init mutations (normalization_type, device,
318 # ...) don't leak between bridges built from the same config.
319 bridge_config = copy.deepcopy(tl_config)
320 bridge_config.architecture = architecture
321 if model_name != "external" or not getattr(bridge_config, "model_name", None):
322 bridge_config.model_name = model_name
323 bridge_config.dtype = dtype
324 else:
325 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, dtype)
327 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
329 if post_adapter_hook is not None:
330 post_adapter_hook(adapter)
332 if device is not None:
333 adapter.cfg.device = str(device)
334 else:
335 try:
336 adapter.cfg.device = str(next(model.parameters()).device)
337 except StopIteration:
338 adapter.cfg.device = "cpu"
340 adapter.prepare_model(model)
342 if tokenizer is not None: 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true
343 from transformer_lens.model_bridge.sources.transformers import setup_tokenizer
345 default_padding_side = getattr(adapter.cfg, "default_padding_side", None)
346 tokenizer = setup_tokenizer(tokenizer, default_padding_side=default_padding_side)
347 (
348 adapter.cfg.tokenizer_prepends_bos,
349 adapter.cfg.tokenizer_appends_eos,
350 ) = detect_tokenizer_bos_eos(tokenizer)
352 return TransformerBridge(model, adapter, tokenizer)