Coverage for transformer_lens/model_bridge/sources/transformers/source.py: 90%
171 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""``boot`` — load a model via HuggingFace ``transformers`` and wrap it in a TransformerBridge."""
2from __future__ import annotations
4import contextlib
5import copy
6import logging
7import warnings
8from typing import Any
10import torch
11from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizerBase
13from transformer_lens.factories.architecture_adapter_factory import (
14 ArchitectureAdapterFactory,
15)
16from transformer_lens.model_bridge.bridge import TransformerBridge
17from transformer_lens.model_bridge.sources._bridge_builder import (
18 build_bridge_config_from_hf,
19 configure_tokenizer,
20 skip_tokenizer_for_modality,
21)
22from transformer_lens.model_bridge.sources._hf_format import (
23 determine_architecture_from_hf_config,
24)
25from transformer_lens.tools.model_registry.registry_io import resolve_model_alias
26from transformer_lens.utilities import get_device
27from transformer_lens.utilities.attn_implementation import force_eager_attention
29from .helpers import (
30 _resolve_checkpoint_to_revision,
31 get_hf_model_class_for_architecture,
32 load_modality_processor,
33)
35# Suppress transformers warnings that go to stderr; otherwise notebook tests fail
36# on unexpected stderr output.
37warnings.filterwarnings("ignore", message=".*generation flags.*not valid.*")
38logging.getLogger("transformers").setLevel(logging.ERROR)
41def boot(
42 model_name: str,
43 hf_config_overrides: dict | None = None,
44 device: str | torch.device | None = None,
45 dtype: torch.dtype = torch.float32,
46 tokenizer: PreTrainedTokenizerBase | None = None,
47 load_weights: bool = True,
48 trust_remote_code: bool = False,
49 model_class: Any | None = None,
50 hf_model: Any | None = None,
51 n_ctx: int | None = None,
52 revision: str | None = None,
53 checkpoint_index: int | None = None,
54 checkpoint_value: int | None = None,
55 # Multi-device placement (accelerate-dispatched). Mutually exclusive with device.
56 device_map: str | dict[str, str | int] | None = None,
57 n_devices: int | None = None,
58 max_memory: dict[str | int, str | int] | None = None,
59 offload_folder: str | None = None,
60) -> TransformerBridge:
61 """Boot a model from HuggingFace (exposed as ``TransformerBridge.boot_transformers``).
63 Returns raw HF weights by default — logits/activations match HF, *not*
64 legacy ``HookedTransformer`` (which folds LayerNorm + centers weights).
65 Call ``enable_compatibility_mode()`` on the result for HookedTransformer-
66 equivalent numerics. Generation, argmax, and CE loss are unaffected.
68 Attention implementation is forced to ``"eager"`` so hooks can capture scores
69 and patterns. For an apples-to-apples HF comparison, load the HF model with
70 ``attn_implementation="eager"`` too; comparing against the default ``"sdpa"``
71 shows ~1e-3 fp32 drift from kernel-level op reordering, not a bridge bug.
73 Args:
74 model_name: The name of the model to load.
75 hf_config_overrides: Optional overrides applied to the HuggingFace config before model load.
76 device: The device to use. If None, will be determined automatically. Mutually exclusive
77 with ``device_map``.
78 dtype: The dtype to use for the model.
79 tokenizer: Optional pre-initialized tokenizer to use; if not provided one will be created.
80 load_weights: If False, load model without weights (on meta device) for config inspection only.
81 model_class: Optional HuggingFace model class to use instead of the default auto-detected
82 class. When the class name matches a key in SUPPORTED_ARCHITECTURES, the corresponding
83 adapter is selected automatically (e.g., BertForNextSentencePrediction).
84 hf_model: Optional pre-loaded HuggingFace model to use instead of loading one. Useful for
85 models loaded with custom configurations (e.g., quantization via BitsAndBytesConfig).
86 When provided, load_weights is ignored.
87 device_map: HuggingFace-style device map (``"auto"``, ``"balanced"``, dict, etc.) for
88 dispatched inference. Explicit maps may include CPU and disk targets; meta targets
89 are still rejected when ``load_weights=True`` (meta has no real data to offload from,
90 unlike disk/cpu). Mixed CPU/disk + GPU maps are rejected too, not because they're
91 known to be broken but because CPU/disk offload has only been verified on CPU-only
92 hardware — no GPU to mix in. ``bridge.enable_compatibility_mode()`` (with weight
93 processing, i.e. not ``no_processing=True``) is unsupported on a CPU/disk-offloaded
94 bridge and raises immediately rather than mid-fold; the default (non-compat-mode)
95 forward pass, ``run_with_cache``, and hooks all work normally under offload. Mutually
96 exclusive with ``device``.
97 n_devices: Convenience: split the model across this many CUDA devices (translated to a
98 ``max_memory`` dict internally). Requires CUDA with at least this many visible devices.
99 max_memory: Optional per-device memory budget for HF's dispatcher.
100 offload_folder: Directory for disk-offloaded weight shards when ``device_map`` includes
101 a ``"disk"`` target. Defaults to a temporary directory (HF's own default) if omitted.
102 n_ctx: Optional context length override. The bridge normally uses the model's documented
103 max context from the HF config. Setting this writes to whichever HF field the model
104 uses (n_positions / max_position_embeddings / etc.), so callers don't need to know
105 the field name. If larger than the model's default, a warning is emitted — quality
106 may degrade past the trained length for rotary models.
107 revision: Optional HF revision string (branch, tag, or commit). Forwarded to
108 config, model, and tokenizer loading.
109 Mutually exclusive with ``checkpoint_index`` and ``checkpoint_value``.
110 checkpoint_index: Index into the available training checkpoints for the model family.
111 Convenience over ``revision`` for checkpointed models like EleutherAI/pythia* and
112 stanford-crfm/*. Resolved to a revision string via the known per-family naming
113 conventions (``step{value}`` for Pythia, ``checkpoint-{value}`` for stanford-crfm).
114 checkpoint_value: Training step or token count of the desired checkpoint. Alternative to
115 ``checkpoint_index``; must be one of the labels returned by ``get_checkpoint_labels``.
117 Returns:
118 The bridge to the loaded model.
119 """
120 official_name = resolve_model_alias(model_name)
121 if official_name is not None:
122 logging.warning(
123 f"DEPRECATED: You are using a deprecated, model_name alias '{model_name}'. TransformerLens will now load the official transformers model name, '{official_name}' instead.\n Please update your code to use the official name by changing model_name from '{model_name}' to '{official_name}'.\nSince TransformerLens v3, all model names should be the official transformers model names.\nThe aliases may be removed in a future version of TransformerLens, so please do the update now."
124 )
125 model_name = official_name
126 if checkpoint_index is not None or checkpoint_value is not None:
127 if revision is not None:
128 raise ValueError(
129 "Specify either revision= or checkpoint_index/checkpoint_value, not both."
130 )
131 revision = _resolve_checkpoint_to_revision(model_name, checkpoint_index, checkpoint_value)
132 # Pass HF token for gated model access (e.g. meta-llama/*)
133 from transformer_lens.utilities.hf_utils import (
134 autoconfig_with_remote_post_init_compat,
135 autotokenizer_with_special_token_compat,
136 get_hf_token,
137 )
139 _hf_token = get_hf_token()
140 if hf_model is not None:
141 # Reuse the pre-loaded model's config to avoid a Hub call when model_name
142 # is a Hub repo ID but the model is already loaded locally.
143 hf_config = copy.deepcopy(hf_model.config)
144 else:
145 # Compat wrapper: 4.x-era remote-code configs (OpenELM) define an
146 # argless __post_init__ that 5.x's dataclass machinery calls with the
147 # class's own fields as kwargs — unloadable without the shim.
148 hf_config = autoconfig_with_remote_post_init_compat(
149 model_name,
150 auto_config=AutoConfig,
151 output_attentions=True,
152 trust_remote_code=trust_remote_code,
153 token=_hf_token,
154 revision=revision,
155 )
156 _n_ctx_field: str | None = None
157 if n_ctx is not None:
158 if n_ctx <= 0:
159 raise ValueError(f"n_ctx must be a positive integer, got n_ctx={n_ctx}.")
160 # Resolve n_ctx to whichever HF config field this model uses. Mirrors the order in
161 # map_default_transformer_lens_config so the TL config derivation picks up the override.
162 for _field in ( 162 ↛ 173line 162 didn't jump to line 173 because the loop on line 162 didn't complete
163 "n_positions",
164 "max_position_embeddings",
165 "max_context_length",
166 "max_length",
167 "seq_length",
168 "max_sequence_length",
169 ):
170 if hasattr(hf_config, _field):
171 _n_ctx_field = _field
172 break
173 if _n_ctx_field is None: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true
174 raise ValueError(
175 f"Cannot apply n_ctx={n_ctx}: no recognized context-length field on "
176 f"HF config for {model_name}. Use hf_config_overrides instead."
177 )
178 _default_n_ctx = getattr(hf_config, _n_ctx_field)
179 if _default_n_ctx is not None and n_ctx > _default_n_ctx:
180 logging.warning(
181 "Setting n_ctx=%d which is larger than the model's default "
182 "context length of %d. The model was not trained on sequences "
183 "this long and may produce unreliable results (especially for "
184 "rotary models without RoPE scaling).",
185 n_ctx,
186 _default_n_ctx,
187 )
188 # Warn if the caller also set the same field via hf_config_overrides — explicit n_ctx wins.
189 if hf_config_overrides and _n_ctx_field in hf_config_overrides:
190 _conflicting_value = hf_config_overrides[_n_ctx_field]
191 if _conflicting_value != n_ctx:
192 logging.warning(
193 "Both n_ctx=%d and hf_config_overrides['%s']=%s were provided. "
194 "The explicit n_ctx takes precedence.",
195 n_ctx,
196 _n_ctx_field,
197 _conflicting_value,
198 )
199 hf_config_overrides = dict(hf_config_overrides or {})
200 hf_config_overrides[_n_ctx_field] = n_ctx
201 if hf_config_overrides:
202 hf_config.__dict__.update(hf_config_overrides)
203 architecture = determine_architecture_from_hf_config(hf_config)
204 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, dtype)
205 bridge_config.trust_remote_code = trust_remote_code
206 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
207 # Pre-loaded models carry their own weight placement (possibly set by the caller via
208 # device_map). Passing device_map / n_devices / max_memory alongside hf_model= is ambiguous
209 # and would silently be ignored, so fail loudly.
210 if hf_model is not None and (
211 device_map is not None or n_devices is not None or max_memory is not None
212 ):
213 raise ValueError(
214 "device_map / n_devices / max_memory are only supported when the bridge loads "
215 "the HF model itself. When passing hf_model=..., apply device_map via "
216 "AutoModel.from_pretrained before handing the model to the bridge."
217 )
218 # Stateful/SSM (Mamba) models keep a per-layer recurrent cache that must live on that
219 # layer's device. The bridge allocates the stateful cache on a single cfg.device, so
220 # cross-device splits would silently misplace the cache. Blocked until v2.
221 if (n_devices is not None and n_devices > 1) or device_map is not None:
222 if getattr(bridge_config, "is_stateful", False): 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 raise ValueError(
224 "Multi-device splits are not yet supported for stateful (SSM / Mamba) "
225 "architectures: the stateful cache allocation is single-device. "
226 "Load on one device, or wait for v2 support."
227 )
228 # Resolve device_map before defaulting `device` — the two are mutually exclusive and the
229 # resolver raises on conflict. If n_devices>1 is passed it's translated into a device_map +
230 # max_memory pair here so downstream code only needs to check the resolved values.
231 from transformer_lens.utilities.multi_gpu import (
232 MIXED_OFFLOAD_GPU_ERROR,
233 count_unique_devices,
234 find_embedding_device,
235 find_misplaced_modules,
236 is_mixed_offload_gpu,
237 maybe_cast_floating_params,
238 resolve_device_map,
239 )
241 resolved_device_map, resolved_max_memory = resolve_device_map(
242 n_devices, device_map, device, max_memory
243 )
244 if resolved_device_map is None:
245 if device is None:
246 device = get_device()
247 adapter.cfg.device = str(device)
248 else:
249 # cfg.device set from hf_device_map after the model is loaded; provisionally None.
250 adapter.cfg.device = None
251 if model_class is None:
252 model_class = get_hf_model_class_for_architecture(architecture)
253 # Ensure pad_token_id exists (v5 raises AttributeError if missing).
254 if not hasattr(hf_config, "pad_token_id") or "pad_token_id" not in hf_config.__dict__:
255 fallback_pad = getattr(hf_config, "eos_token_id", None)
256 # eos_token_id can be a list (Gemma3 uses [1, 106]); take the first.
257 if isinstance(fallback_pad, list):
258 fallback_pad = fallback_pad[0] if fallback_pad else None
259 hf_config.pad_token_id = fallback_pad
260 model_kwargs = {"config": hf_config, "torch_dtype": dtype}
261 if _hf_token: 261 ↛ 263line 261 didn't jump to line 263 because the condition on line 261 was always true
262 model_kwargs["token"] = _hf_token
263 if trust_remote_code:
264 model_kwargs["trust_remote_code"] = True
265 if revision is not None:
266 model_kwargs["revision"] = revision
267 if resolved_device_map is not None:
268 model_kwargs["device_map"] = resolved_device_map
269 if resolved_max_memory is not None: 269 ↛ 270line 269 didn't jump to line 270 because the condition on line 269 was never true
270 model_kwargs["max_memory"] = resolved_max_memory
271 if offload_folder is not None:
272 model_kwargs["offload_folder"] = offload_folder
273 if hasattr(adapter.cfg, "attn_implementation") and adapter.cfg.attn_implementation is not None:
274 model_kwargs["attn_implementation"] = adapter.cfg.attn_implementation
275 else:
276 # Eager is required for output_attentions hooks.
277 model_kwargs["attn_implementation"] = "eager"
278 adapter.prepare_loading(model_name, model_kwargs)
279 # Meta device_map targets crash at boot when loading weights
280 # (NotImplementedError in HF tie_weights, KeyError in Accelerate offload hooks).
281 # Only accepted with load_weights=False (config inspection; map not applied).
282 if load_weights and isinstance(resolved_device_map, dict):
283 _meta_targets = [
284 k
285 for k, v in resolved_device_map.items()
286 if isinstance(v, str) and v.strip().lower() == "meta"
287 ]
288 if _meta_targets:
289 raise ValueError(
290 f"device_map contains meta target(s): {_meta_targets}. "
291 "Meta device_map values crash at boot when loading weights. "
292 "Set load_weights=False for config inspection only "
293 "(the map is not applied; parameters load on CPU via from_config)."
294 )
295 if hf_model is not None:
296 # Pre-loaded weights stay untouched (quantized models with custom device_map are
297 # why this branch takes the model as-is), but attn-impl selection is runtime
298 # dispatch — safe to switch in place. Without it, transformers' default sdpa
299 # returns attn_weights=None and hook_pattern / hook_attn_scores silently never
300 # fire. An explicit adapter/user choice wins, matching the load path above.
301 if getattr(adapter.cfg, "attn_implementation", None) is None: 301 ↛ 348line 301 didn't jump to line 348 because the condition on line 301 was always true
302 force_eager_attention(hf_model)
303 elif not load_weights:
304 from_config_kwargs = {}
305 if trust_remote_code: 305 ↛ 306line 305 didn't jump to line 306 because the condition on line 305 was never true
306 from_config_kwargs["trust_remote_code"] = True
307 # adapter.prepare_loading may have replaced model_kwargs["config"] (e.g. Qwen3.5
308 # text-only extraction); honor that here so the no-weights path uses the
309 # same config the load-weights path would.
310 prepared_config = model_kwargs.get("config", hf_config)
311 with contextlib.redirect_stdout(None):
312 hf_model = model_class.from_config(prepared_config, **from_config_kwargs)
313 else:
314 try:
315 hf_model = model_class.from_pretrained(model_name, **model_kwargs)
316 except RuntimeError as e:
317 # HF refuses to load when positional-weight shapes don't match. If the user
318 # requested an n_ctx that conflicts with the saved weights (common for
319 # learned-pos-embed models like GPT-2), re-raise with a clearer message.
320 if n_ctx is not None and "ignore_mismatched_sizes" in str(e): 320 ↛ 331line 320 didn't jump to line 331 because the condition on line 320 was always true
321 raise RuntimeError(
322 f"Failed to load {model_name} with n_ctx={n_ctx}: the pretrained "
323 f"weights' positional-embedding shape does not match the requested "
324 f"context length. This affects models with learned positional "
325 f"embeddings (e.g. GPT-2, OPT). Options: (1) use the model's "
326 f"default n_ctx, (2) pass load_weights=False if you only need "
327 f"config inspection, or (3) choose a rotary-embedding model "
328 f"(e.g. Llama, Mistral) which supports n_ctx changes without "
329 f"weight mismatch."
330 ) from e
331 raise
332 # Skip explicit .to(device) when accelerate has placed weights via device_map.
333 if resolved_device_map is None and device is not None:
334 hf_model = hf_model.to(device)
335 # Cast params to dtype; preserve float32 buffers (e.g. RotaryEmbedding.inv_freq).
336 # Use module-level alignment so Accelerate can temporarily materialize offloaded
337 # parameters before we touch them.
338 # Skip dtype normalization entirely when the model has an active quantizer.
339 # from_pretrained settled the load dtype above, and on a quantized checkpoint
340 # that is the quantizer's effective dtype rather than necessarily the requested
341 # one (HfQuantizer.update_dtype lets AWQ, fbgemm-FP8 and FP-Quant override it;
342 # AWQ keys off CUDA/XPU being available at all, not off placement).
343 # Re-casting would overwrite those choices along with quantizer-owned storage,
344 # whose width is the quantizer's to choose (FP8 *or* float32 scales).
345 maybe_cast_floating_params(hf_model, dtype)
346 # Derive cfg.device / cfg.n_devices from hf_device_map when present. Covers fresh loads
347 # with a resolved device_map AND pre-loaded models with caller-dispatched device_map="auto".
348 hf_device_map_post = getattr(hf_model, "hf_device_map", None)
349 if hf_device_map_post:
350 # All-CPU placement, and disk offload, are supported (GeneralizedComponent.__call__
351 # wraps forward in Accelerate's align_module_device, so wrapped components reading
352 # raw params directly still see materialized data). meta remains rejected here: it
353 # has no real data to materialize even under offload/dispatch, unlike disk/cpu.
354 offload_values = {str(v).lower() for v in hf_device_map_post.values() if isinstance(v, str)}
355 unsupported = offload_values & {"meta"}
356 if unsupported: 356 ↛ 357line 356 didn't jump to line 357 because the condition on line 356 was never true
357 raise ValueError(
358 f"hf_device_map contains unsupported offload targets: {sorted(unsupported)}. "
359 "TransformerBridge currently supports CPU and disk device_map targets."
360 )
361 if is_mixed_offload_gpu(hf_device_map_post.values()): 361 ↛ 362line 361 didn't jump to line 362 because the condition on line 361 was never true
362 raise ValueError(f"Realized hf_device_map is unsupported: {MIXED_OFFLOAD_GPU_ERROR}")
363 if ( 363 ↛ 369line 363 didn't jump to line 369 because the condition on line 363 was never true
364 "cpu" in offload_values
365 and device_map is None
366 and n_devices is not None
367 and n_devices > 1
368 ):
369 raise ValueError(
370 "hf_device_map contains CPU targets. n_devices is GPU-only; pass device_map "
371 "explicitly for CPU placement."
372 )
373 misplaced = find_misplaced_modules(hf_model)
374 if misplaced: 374 ↛ 375line 374 didn't jump to line 375 because the condition on line 374 was never true
375 details = "; ".join(
376 f"{name!r} mapped to {mapped} but loaded on {actual}"
377 for name, mapped, actual in misplaced
378 )
379 raise ValueError(
380 f"device_map entries were not honored: {details}. This usually means the "
381 "map splits tied parameters (e.g. GPT-2's wte/lm_head share one tensor) "
382 "across devices — accelerate places a tied parameter once, leaving a module "
383 "executing on a device its weights aren't on, which crashes mid-forward. "
384 "Map tied modules to the same device."
385 )
386 embedding_device = find_embedding_device(hf_model)
387 if embedding_device is not None:
388 adapter.cfg.device = str(embedding_device)
389 adapter.cfg.n_devices = count_unique_devices(hf_model)
390 elif adapter.cfg.device is None:
391 # Pre-loaded single-device model with no hf_device_map — fall back to first param.
392 try:
393 adapter.cfg.device = str(next(hf_model.parameters()).device)
394 except StopIteration:
395 adapter.cfg.device = "cpu"
396 # Verify the n_ctx override actually took effect on the loaded model. If HF's config class
397 # silently dropped or normalized the value, warn so the user isn't misled.
398 if n_ctx is not None and _n_ctx_field is not None and hf_model is not None:
399 _actual = getattr(hf_model.config, _n_ctx_field, None)
400 if _actual != n_ctx:
401 logging.warning(
402 "n_ctx=%d was requested but hf_model.config.%s=%s after load. "
403 "The override may not have taken effect; the model may not "
404 "accept sequences longer than %s.",
405 n_ctx,
406 _n_ctx_field,
407 _actual,
408 _actual,
409 )
410 adapter.prepare_model(hf_model)
411 if tokenizer is not None:
412 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)
413 elif not skip_tokenizer_for_modality(adapter.cfg):
414 token_arg = get_hf_token()
415 use_fast = getattr(adapter.cfg, "use_fast", True)
416 # Some adapters override tokenizer source (e.g. OpenELM has no tokenizer of its own).
417 tokenizer_source = model_name
418 if hasattr(adapter.cfg, "tokenizer_name") and adapter.cfg.tokenizer_name is not None: 418 ↛ 419line 418 didn't jump to line 419 because the condition on line 418 was never true
419 tokenizer_source = adapter.cfg.tokenizer_name
420 # Encoder-decoder models like T5 don't have a BOS token and raise on add_bos_token=True.
421 try:
422 base_tokenizer = autotokenizer_with_special_token_compat(
423 tokenizer_source,
424 auto_tokenizer=AutoTokenizer,
425 add_bos_token=True,
426 use_fast=use_fast,
427 token=token_arg,
428 trust_remote_code=trust_remote_code,
429 revision=revision,
430 )
431 except ValueError:
432 base_tokenizer = AutoTokenizer.from_pretrained(
433 tokenizer_source,
434 use_fast=use_fast,
435 token=token_arg,
436 trust_remote_code=trust_remote_code,
437 revision=revision,
438 )
439 tokenizer = configure_tokenizer(base_tokenizer, adapter.cfg)
440 from transformer_lens.model_bridge.sources.transformers_driver import (
441 TransformersDriver,
442 )
444 driver = TransformersDriver(hf_model, adapter, tokenizer)
445 bridge = TransformerBridge(hf_model, adapter, tokenizer, driver=driver)
446 load_modality_processor(bridge, adapter.cfg, model_name, trust_remote_code, _hf_token)
447 # Mirror the legacy loader: record which training checkpoint this is, so
448 # checkpoint sweeps (e.g. induction-head formation) can read it back.
449 if checkpoint_index is not None or checkpoint_value is not None:
450 from transformer_lens.tools.model_registry.checkpoints import (
451 get_checkpoint_labels,
452 )
454 labels, _label_kind = get_checkpoint_labels(model_name)
455 if checkpoint_value is not None: 455 ↛ 456line 455 didn't jump to line 456 because the condition on line 455 was never true
456 resolved_value = checkpoint_value
457 else:
458 assert checkpoint_index is not None # narrowed by the enclosing if
459 resolved_value = labels[checkpoint_index]
460 bridge.cfg.checkpoint_value = resolved_value # type: ignore[attr-defined]
461 bridge.cfg.checkpoint_index = labels.index(resolved_value) # type: ignore[attr-defined]
462 return bridge