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