Coverage for transformer_lens/model_bridge/sources/_bridge_builder.py: 95%
100 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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.model_bridge.sources._hf_format import (
17 get_effective_text_config,
18 map_default_transformer_lens_config,
19 setup_tokenizer,
20)
21from transformer_lens.utilities.heterogeneous_config import (
22 het_safe_view,
23 per_layer_attr_names,
24 safe_config_get,
25)
27# Architecture-agnostic; do not extend per-architecture.
28_HF_PASSTHROUGH_ATTRS = [
29 # OPT
30 "is_gated_act",
31 # LongT5
32 "encoder_attention_type",
33 "word_embed_proj_dim",
34 "do_layer_norm_before",
35 # BART
36 "encoder_layers",
37 "decoder_layers",
38 "encoder_attention_heads",
39 "decoder_attention_heads",
40 "encoder_ffn_dim",
41 "decoder_ffn_dim",
42 # Marian
43 "scale_embedding",
44 # Granite
45 "position_embedding_type",
46 "logits_scaling",
47 "residual_multiplier",
48 # Falcon
49 "parallel_attn",
50 "multi_query",
51 "new_decoder_architecture",
52 "alibi",
53 "num_ln_in_parallel_attn",
54 # GPTNeoX
55 "use_parallel_residual",
56 # Mamba (SSM config)
57 "state_size",
58 "conv_kernel",
59 "expand",
60 "time_step_rank",
61 "intermediate_size",
62 # Mamba-2 (additional SSM config)
63 "n_groups",
64 "chunk_size",
65 # Falcon-H1 (parallel attn + Mamba-2 hybrid SSM config)
66 "mamba_d_ssm",
67 "mamba_n_heads",
68 "mamba_d_head",
69 "mamba_d_state",
70 "mamba_n_groups",
71 "mamba_d_conv",
72 "mamba_chunk_size",
73 "lm_head_multiplier",
74 # Multimodal
75 "vision_config",
76 # Cohere
77 "logit_scale",
78 "rope_parameters",
79 # HRM-Text
80 "H_cycles",
81 "L_cycles",
82 "L_bp_cycles",
83 "embedding_scale",
84 "prefix_lm",
85 "num_layers_per_stack",
86 "sliding_window_pattern",
87 "_sliding_window_pattern",
88 # Hybrid/MoE architectures
89 "layer_types",
90 "moe_intermediate_size",
91 "shared_expert_intermediate_size",
92 "norm_eps",
93 "attention_bias",
94 "lm_head_bias",
95 "router_jitter_noise",
96 "input_jitter_noise",
97 "eos_token_id",
98 # LLaDA remote-code model contract and tokenizer metadata
99 "block_type",
100 "block_group_size",
101 "rope",
102 "rope_full_precision",
103 "attention_layer_norm",
104 "include_bias",
105 "include_qkv_bias",
106 "scale_logits",
107 "input_emb_norm",
108 "layer_norm_type",
109 "embedding_size",
110 "mask_token_id",
111 "pad_token_id",
112 "bos_token_id",
113 # BD3LM
114 "model_length",
115 "block_size",
116 "cond_dim",
117 "adaln",
118 "cross_attn",
119 # Zamba2 (Mamba-2 + shared-attention hybrid)
120 "mamba_expand",
121 "mamba_ngroups",
122 "num_mem_blocks",
123 "layers_block_type",
124 "use_shared_attention_adapter",
125 # Jamba (attention + Mamba-1 hybrid; MoE schedule knobs)
126 "mamba_dt_rank",
127 "attn_layer_period",
128 "attn_layer_offset",
129 "expert_layer_period",
130 "expert_layer_offset",
131 # Ouro (LoopLM)
132 "total_ut_steps",
133 "early_exit_threshold",
134 # DeepSeek V4 (mHC + compressed attention)
135 "compress_rates",
136 "compress_rope_theta",
137 "hc_mult",
138 "hc_sinkhorn_iters",
139 "hc_eps",
140 "mlp_layer_types",
141 "swiglu_limit",
142 "o_groups",
143 "o_lora_rank",
144 "index_n_heads",
145 "index_head_dim",
146 "index_topk",
147 "q_lora_rank",
148 # Raven / Huginn (depth-recurrent)
149 "mean_recurrence",
150 "mean_backprop_depth",
151 "n_layers_in_prelude",
152 "n_layers_in_recurrent_block",
153 "n_layers_in_coda",
154 "injection_type",
155 "qk_bias",
156 # RWKV-7 (attention-free recurrent, generalized delta-rule time-mixing).
157 # head_dim is intentionally omitted: it is a read-only alias of d_head on
158 # TransformerBridgeConfig, so a passthrough setattr would raise.
159 "num_heads",
160 "value_dim",
161 "decay_low_rank_dim",
162 "gate_low_rank_dim",
163 "a_low_rank_dim",
164 "v_low_rank_dim",
165 "norm_first",
166 "norm_bias",
167 "fuse_norm",
168 "attn_mode",
169 "hidden_act",
170]
173def build_bridge_config_from_hf(
174 hf_config: Any,
175 architecture: str,
176 model_name: str,
177 dtype: torch.dtype,
178) -> TransformerBridgeConfig:
179 """Translate an HF config into a :class:`TransformerBridgeConfig`."""
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"):
223 encoder_subconfig = safe_config_get(encoder_subconfig, "text_config")
224 if encoder_subconfig is not None: 224 ↛ 226line 224 didn't jump to line 226 because the condition on line 224 was always true
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."""
239 # Non-empty test string — "" is unreliable with token aliasing.
240 encoded_test = tokenizer.encode("a")
241 # CLS counts: BERT-style tokenizers prepend [CLS], which the legacy stack
242 # treats as the BOS-like token; comparing only against bos_token_id (a
243 # fallback string on such tokenizers) concludes False and desyncs the stacks.
244 leading_special_ids = {
245 token_id
246 for token_id in (tokenizer.bos_token_id, getattr(tokenizer, "cls_token_id", None))
247 if token_id is not None
248 }
249 prepends_bos = len(encoded_test) > 1 and encoded_test[0] in leading_special_ids
250 appends_eos = (
251 len(encoded_test) > 1
252 and tokenizer.eos_token_id is not None
253 and encoded_test[-1] == tokenizer.eos_token_id
254 )
255 return prepends_bos, appends_eos
258def skip_tokenizer_for_modality(cfg: Any) -> bool:
259 """True when a source must not auto-load a text tokenizer: audio/vision models use
260 feature extractors or image processors, not text tokenizers (their repos ship none)."""
261 return bool(getattr(cfg, "is_audio_model", False) or getattr(cfg, "is_visual_model", False))
264def configure_tokenizer(tokenizer: Any, cfg: Any) -> Any:
265 """Shared boot step: normalize the tokenizer and record its BOS/EOS behavior on cfg.
267 Every source must run this — skipping it leaves the dataclass default
268 ``tokenizer_prepends_bos=True``, which position-shifts every activation on
269 non-BOS-prepending tokenizers (Qwen family)."""
270 tokenizer = setup_tokenizer(
271 tokenizer, default_padding_side=getattr(cfg, "default_padding_side", None)
272 )
273 cfg.tokenizer_prepends_bos, cfg.tokenizer_appends_eos = detect_tokenizer_bos_eos(tokenizer)
274 return tokenizer
277def build_bridge_from_module(
278 model: nn.Module,
279 architecture: str,
280 *,
281 hf_config: Optional[Any] = None,
282 tl_config: Optional[TransformerBridgeConfig] = None,
283 tokenizer: Optional[Any] = None,
284 dtype: Optional[torch.dtype] = None,
285 device: Optional[Any] = None,
286 model_name: str = "external",
287 post_adapter_hook: Optional[Callable[[ArchitectureAdapter], None]] = None,
288) -> TransformerBridge:
289 """Build a :class:`TransformerBridge` around a pre-loaded model.
291 The bridge never moves, casts, or mutates the supplied model.
293 Args:
294 model: Any ``nn.Module`` whose submodule tree matches the adapter's
295 expected dot-paths for ``architecture``.
296 architecture: Architecture identifier registered in the
297 ``ArchitectureAdapterFactory`` (e.g. ``"LlamaForCausalLM"``,
298 ``"TransformerLensNative"``).
299 hf_config: Optional HF-style config; translated via
300 :func:`build_bridge_config_from_hf`. Mutually exclusive with ``tl_config``.
301 tl_config: Optional pre-built :class:`TransformerBridgeConfig`; bypasses
302 HF translation. Mutually exclusive with ``hf_config``.
303 tokenizer: Optional tokenizer. If supplied, passes through
304 ``setup_tokenizer`` and detects BOS/EOS behavior.
305 dtype: Recorded on ``cfg.dtype``. Default ``None`` reads from the model's
306 first parameter; explicit values override.
307 device: Recorded on ``cfg.device``. Default ``None`` reads from the
308 model's first parameter.
309 model_name: Recorded on ``cfg.model_name``.
310 post_adapter_hook: Optional callback invoked after adapter selection and
311 before :meth:`adapter.prepare_model`. Source-specific overlays mutate
312 ``component_mapping`` here.
314 Returns:
315 A :class:`TransformerBridge` wrapping the supplied model.
316 """
317 if hf_config is None and tl_config is None:
318 raise ValueError(
319 "build_bridge_from_module requires exactly one of hf_config or "
320 "tl_config — the bridge needs config fields (d_model, n_heads, "
321 "n_layers, ...) that can't be inferred from the model alone."
322 )
323 if hf_config is not None and tl_config is not None:
324 raise ValueError(
325 "build_bridge_from_module got both hf_config and tl_config; supply "
326 "exactly one. hf_config triggers HF→bridge translation; tl_config "
327 "bypasses it."
328 )
330 # Reading dtype from the model avoids silently lying about a bf16 model.
331 if dtype is None:
332 try:
333 dtype = next(model.parameters()).dtype
334 except StopIteration:
335 dtype = torch.float32
337 if tl_config is not None:
338 # Defensive copy so adapter-init mutations (normalization_type, device,
339 # ...) don't leak between bridges built from the same config.
340 bridge_config = copy.deepcopy(tl_config)
341 bridge_config.architecture = architecture
342 # Explicit kwarg wins over whatever tl_config carries; default only fills a gap.
343 if model_name != "external" or not getattr(bridge_config, "model_name", None):
344 bridge_config.model_name = model_name
345 bridge_config.dtype = dtype
346 else:
347 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, dtype)
349 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
351 if post_adapter_hook is not None:
352 post_adapter_hook(adapter)
354 if device is not None:
355 adapter.cfg.device = str(device)
356 else:
357 try:
358 adapter.cfg.device = str(next(model.parameters()).device)
359 except StopIteration:
360 adapter.cfg.device = "cpu"
362 adapter.prepare_model(model)
364 if tokenizer is not None:
365 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)
367 from transformer_lens.model_bridge.sources.transformers_driver import (
368 TransformersDriver,
369 )
371 driver = TransformersDriver(model, adapter, tokenizer)
372 return TransformerBridge(model, adapter, tokenizer, driver=driver)