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