transformer_lens.model_bridge.sources.transformers.source module

boot — load a model via HuggingFace transformers and wrap it in a TransformerBridge.

transformer_lens.model_bridge.sources.transformers.source.boot(model_name: str, hf_config_overrides: dict | None = None, device: str | device | None = None, dtype: dtype = torch.float32, tokenizer: PreTrainedTokenizerBase | None = None, load_weights: bool = True, trust_remote_code: bool = False, model_class: Any | None = None, hf_model: Any | None = None, n_ctx: int | None = None, revision: str | None = None, checkpoint_index: int | None = None, checkpoint_value: int | None = None, device_map: str | dict[str, str | int] | None = None, n_devices: int | None = None, max_memory: dict[str | int, str | int] | None = None, offload_folder: str | None = None) → TransformerBridge

Boot a model from HuggingFace (exposed as TransformerBridge.boot_transformers).

Returns raw HF weights by default — logits/activations match HF, not legacy HookedTransformer (which folds LayerNorm + centers weights). Call enable_compatibility_mode() on the result for HookedTransformer- equivalent numerics. Generation, argmax, and CE loss are unaffected.

Attention implementation is forced to "eager" so hooks can capture scores and patterns. For an apples-to-apples HF comparison, load the HF model with attn_implementation="eager" too; comparing against the default "sdpa" shows ~1e-3 fp32 drift from kernel-level op reordering, not a bridge bug.

Parameters:

  • model_name – The name of the model to load.

  • hf_config_overrides – Optional overrides applied to the HuggingFace config before model load.

  • device – The device to use. If None, will be determined automatically. Mutually exclusive with device_map.

  • dtype – The dtype to use for the model.

  • tokenizer – Optional pre-initialized tokenizer to use; if not provided one will be created.

  • load_weights – If False, load model without weights (on meta device) for config inspection only.

  • model_class – Optional HuggingFace model class to use instead of the default auto-detected class. When the class name matches a key in SUPPORTED_ARCHITECTURES, the corresponding adapter is selected automatically (e.g., BertForNextSentencePrediction).

  • hf_model – Optional pre-loaded HuggingFace model to use instead of loading one. Useful for models loaded with custom configurations (e.g., quantization via BitsAndBytesConfig). When provided, load_weights is ignored.

  • device_map – HuggingFace-style device map ("auto", "balanced", dict, etc.) for dispatched inference. Explicit maps may include CPU and disk targets; meta targets are still rejected when load_weights=True (meta has no real data to offload from, unlike disk/cpu). Mixed CPU/disk + GPU maps are rejected too, not because they’re known to be broken but because CPU/disk offload has only been verified on CPU-only hardware — no GPU to mix in. bridge.enable_compatibility_mode() (with weight processing, i.e. not no_processing=True) is unsupported on a CPU/disk-offloaded bridge and raises immediately rather than mid-fold; the default (non-compat-mode) forward pass, run_with_cache, and hooks all work normally under offload. Mutually exclusive with device.

  • n_devices – Convenience: split the model across this many CUDA devices (translated to a max_memory dict internally). Requires CUDA with at least this many visible devices.

  • max_memory – Optional per-device memory budget for HF’s dispatcher.

  • offload_folder – Directory for disk-offloaded weight shards when device_map includes a "disk" target. Defaults to a temporary directory (HF’s own default) if omitted.

  • n_ctx – Optional context length override. The bridge normally uses the model’s documented max context from the HF config. Setting this writes to whichever HF field the model uses (n_positions / max_position_embeddings / etc.), so callers don’t need to know the field name. If larger than the model’s default, a warning is emitted — quality may degrade past the trained length for rotary models.

  • revision – Optional HF revision string (branch, tag, or commit). Forwarded to config, model, and tokenizer loading. Mutually exclusive with checkpoint_index and checkpoint_value.

  • checkpoint_index – Index into the available training checkpoints for the model family. Convenience over revision for checkpointed models like EleutherAI/pythia* and stanford-crfm/*. Resolved to a revision string via the known per-family naming conventions (step{value} for Pythia, checkpoint-{value} for stanford-crfm).

  • checkpoint_value – Training step or token count of the desired checkpoint. Alternative to checkpoint_index; must be one of the labels returned by get_checkpoint_labels.

Returns:

The bridge to the loaded model.