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, code_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, build the model with from_config instead of loading the checkpoint: real parameters with the architecture’s random initialisation, on the CPU, in the config’s dtype and in training mode. device is recorded in cfg.device but not applied, and neither are dtype nor device_map; move, cast and eval() the model before running it on another device or in another dtype.

  • 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.

  • code_revision – Optional HF revision string pinning the trust_remote_code modeling file separately from the weights (HF’s own code_revision). Forwarded alongside revision to config and model loading and to the adapters’ remote-code patching; weights and tokenizer keep revision.

  • 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.