diff --git a/docs/source/content/analysis_tools.md b/docs/source/content/analysis_tools.md index b9683ecf5..bb0c8c1e5 100644 --- a/docs/source/content/analysis_tools.md +++ b/docs/source/content/analysis_tools.md @@ -20,7 +20,7 @@ the [4.0 migration guide](migrating_to_v4.md) for existing `HookedTransformer` c | What does each layer "predict", and where does a token's rank change? | **Logit Lens** | A prompt or activation cache, optional target tokens, top-k or a vocabulary subset → per-layer logits, log-probabilities, ranks and entropy. | What the residual stream reads as through the real final norm and unembedding; the final entry reproduces the model's own logits. | | What vocabulary directions appear in an MLP's gradient factors? | **Backward Lens** | A prompt, one target token, and selected layers → gradient factors and vocabulary rankings. | A diagnostic of the forward inputs and backward signals composing a weight gradient. | | How can I read or edit residuals through a fitted transport map? | **Jacobian Lens** | A matching lens artifact and model, plus prompts or activations → vocabulary readouts, decompositions, or interventions. | Readouts under the fitted map; causal effects require running and measuring an intervention. | -| Does a subspace at one hook point carry a variable the output depends on? | **Interchange Interventions** | A hook point, base and source inputs, a basis (fixed, or learned by DAS from pairs labelled by a high-level model), and a task readout → task logits with the base's subspace coordinates set to the source's. | The effect of that intervention on the chosen pairs; compare with a norm-matched placebo and random bases before reading it as use. | +| Does a subspace at one hook point carry a variable the output depends on? | **Interchange Interventions** | A hook point, base and source inputs, a basis (fixed, or learned by DAS from pairs labelled by a high-level model), and a task readout → task logits with the base's subspace coordinates set to the source's. | The effect of that intervention on the chosen pairs; compare a learned subspace with its trained controls (arbitrary labels, an untrained model), a norm-matched placebo and random bases before reading it as use. | | How much do two subspaces overlap? | **Projection Kernel** | Two subspace bases, or Bridge attention-head weight spaces → overlap scores and principal-angle information. | Shared geometric support, independent of the choice of basis within each subspace. | | What are a head's QK/OV singular directions? | **SVD Circuits: head decomposition** | Model weights and a layer/head index → singular values, vectors, and degeneracy reports. | The linear structure of the head's weight maps. | diff --git a/docs/source/content/interchange_interventions.md b/docs/source/content/interchange_interventions.md index 0564378f8..f0abf118f 100644 --- a/docs/source/content/interchange_interventions.md +++ b/docs/source/content/interchange_interventions.md @@ -36,18 +36,19 @@ task_logits = lambda tokens: bridge(tokens)[:, -1, answers] hook = "blocks.1.hook_out" -def pairs(nouns): +def pairs(nouns, place): """Bases and sources differ only in the subject's number; the target is the source's verb.""" - singular = bridge.to_tokens([f"The {n} near the cabinet" for n in nouns]) - plural = bridge.to_tokens([f"The {n}s near the cabinet" for n in nouns]) + singular = bridge.to_tokens([f"The {n} {place}" for n in nouns]) + plural = bridge.to_tokens([f"The {n}s {place}" for n in nouns]) base, source = torch.cat([singular, plural]), torch.cat([plural, singular]) targets = (torch.arange(len(base)) < len(nouns)).long() # " are" when the source is plural subject = torch.full((len(base),), 2) # the subject token, after the BOS token and "The" return base, source, targets, subject -fit_base, fit_source, fit_targets, fit_at = pairs(["key", "dog", "cat", "book", "car", "boy"]) -test_base, test_source, test_targets, test_at = pairs(["girl", "bird", "tree", "house"]) +nouns = ["key", "dog", "cat", "book", "car", "boy", "girl", "bird"] +fit_base, fit_source, fit_targets, fit_at = pairs(nouns, "near the cabinet") +test_base, test_source, test_targets, test_at = pairs(nouns, "behind the table") # Fit the basis on the fit pairs' sources, labelled by subject number ... fit_acts = collect_activations(bridge, hook, fit_source, task_logits, positions=fit_at) @@ -65,7 +66,9 @@ print((logits.argmax(-1) == test_targets).float().mean()) # tensor(1.) tools do not assume next-token prediction; the hook must fire exactly once in that pass. Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position index per example. Fit a basis on one set of pairs and evaluate it on pairs it was not fit on, as -here. +here: the held-out pairs put the same nouns in a new frame. Reusing the nouns gives the +arbitrary-label control below keys it can learn; holding out nouns as well is a stronger test of +the real alignment. `collect_activations` and `intervened_logits` run in batches of `batch_size` without gradients and return CPU tensors, so memory does not grow with the number of pairs. `run_with_intervention` is @@ -113,6 +116,59 @@ together. The model's parameters are frozen during training and restored afterwa DAS is expressive enough to fit arbitrary label mappings and to find effects in randomly initialised models (Arora et al., 2024), so its IIA needs controls trained the same way. +## Trained controls + +Train each control with the same pairs, hook point, rank, budget and seeds as the real alignment, +and evaluate it on the same held-out pairs: + +- **Arbitrary labels.** `arbitrary_targets(keys, n_classes, seed)` gives each key a fixed + pseudo-random class. DAS trained on these shows how much IIA the method's expressivity alone + produces. Key the labels by something that recurs between the fit and held-out pairs, such as + the source's subject word: on keys the fit pairs never contained, the labels are fresh coin + flips, every subspace scores chance, and the control can reject nothing. +- **An untrained model.** `random_init_bridge(model_name, seed=...)` builds the architecture with + its HuggingFace initialisation and no training; `reinitialized(model, seed)` does the same for a + `HookedRootModule`. An effect DAS finds there does not come from anything the model learned. + +```python +from transformer_lens.tools.analysis import arbitrary_targets, random_init_bridge + +# Arbitrary labels: a fixed pseudo-random class per source subject word, which recurs between the +# fit and held-out pairs +labels = lambda sources: arbitrary_targets([bridge.to_string(s[2]) for s in sources], 2, seed=0) +arbitrary = fit_das( + bridge, hook, fit_base, fit_source, labels(fit_source), task_logits, rank=1, + base_positions=fit_at, source_positions=fit_at, epochs=10, +) +logits = intervened_logits( + bridge, hook, SubspaceInterchange(arbitrary.subspace), test_base, test_acts, task_logits, + positions=test_at, +) +print((logits.argmax(-1) == labels(test_source)).float().mean()) # tensor(0.6250) + +# An untrained model: the same pairs and targets, its own activations and readout +untrained = random_init_bridge("gpt2", seed=0) +untrained_logits = lambda tokens: untrained(tokens)[:, -1, answers] +random_init = fit_das( + untrained, hook, fit_base, fit_source, fit_targets, untrained_logits, rank=1, + base_positions=fit_at, source_positions=fit_at, epochs=10, +) +untrained_acts = collect_activations( + untrained, hook, test_source, untrained_logits, positions=test_at +) +logits = intervened_logits( + untrained, hook, SubspaceInterchange(random_init.subspace), test_base, untrained_acts, + untrained_logits, positions=test_at, +) +print((logits.argmax(-1) == test_targets).float().mean()) # tensor(0.8125) +``` + +Here the arbitrary labels reach 0.62 and the untrained model 0.81, against 1.0 for the real +alignment: DAS finds a subspace that moves even an untrained model's output, which is why the real +alignment's IIA has to be read against both. Report the real alignment's IIA minus each control's +(its selectivity), with uncertainty over the evaluation pairs. A subspace that does not clearly +beat both controls is not evidence that the model carries the variable there. + ## Placebo `NormMatchedPlacebo(reference, basis)` moves each base exactly as far as `reference` does, but diff --git a/tests/integration/model_bridge/test_interchange.py b/tests/integration/model_bridge/test_interchange.py index 74ad81abc..215cc20ec 100644 --- a/tests/integration/model_bridge/test_interchange.py +++ b/tests/integration/model_bridge/test_interchange.py @@ -1,7 +1,10 @@ """Real-Bridge integration checks for interchange interventions.""" +import pytest import torch +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.tools.analysis._model_state import require_eval_mode from transformer_lens.tools.analysis.interchange import ( NormMatchedPlacebo, SubspaceInterchange, @@ -9,6 +12,7 @@ fit_das, intervened_logits, random_basis, + random_init_bridge, ) PROMPTS = ["The key near the cabinet", "The keys near the cabinet"] @@ -119,3 +123,39 @@ def accuracy(basis): assert accuracy(result.subspace) >= 0.9 assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2 + + +@pytest.mark.parametrize("name", ["gpt2", "tiny-stories-1M"]) +def test_random_init_bridge_matches_the_trained_architecture_without_the_training(name): + trained = TransformerBridge.boot_transformers(name, device="cpu") + before = torch.random.get_rng_state() + untrained = random_init_bridge(name, seed=0, device="cpu") + assert torch.equal(torch.random.get_rng_state(), before) # the caller's stream is untouched + assert type(untrained.original_model) is type(trained.original_model) + assert untrained.cfg.n_layers == trained.cfg.n_layers + assert untrained.cfg.d_model == trained.cfg.d_model + assert untrained.W_E.dtype == torch.float32 + assert not torch.allclose(untrained.W_E, trained.W_E) + assert torch.equal(untrained.W_E, random_init_bridge(name, seed=0, device="cpu").W_E) + require_eval_mode(untrained, operation="the random-init control") # ready for fit_das + tokens = untrained.to_tokens(PROMPTS) + with torch.no_grad(): + assert untrained(tokens).shape == trained(tokens).shape + + +def test_random_init_bridge_casts_to_a_dtype_other_than_the_configs(): + untrained = random_init_bridge( + "gpt2", seed=0, dtype=torch.bfloat16, device="cpu" + ) # config: float32 + floating = [p for p in untrained.original_model.parameters() if p.is_floating_point()] + assert floating and all(p.dtype == torch.bfloat16 for p in floating) + + +def test_random_init_bridge_puts_its_weights_on_the_device_it_records(monkeypatch): + # Stand in for a GPU machine, where boot's default device is not the CPU + import transformer_lens.model_bridge.sources.transformers.source as source + + monkeypatch.setattr(source, "get_device", lambda: "meta") + untrained = random_init_bridge("gpt2", seed=0) + assert str(untrained.cfg.device) == "meta" + assert all(p.device.type == "meta" for p in untrained.original_model.parameters()) diff --git a/tests/unit/tools/test_interchange.py b/tests/unit/tools/test_interchange.py index f037d556f..c1c6f81af 100644 --- a/tests/unit/tools/test_interchange.py +++ b/tests/unit/tools/test_interchange.py @@ -12,12 +12,14 @@ NormMatchedPlacebo, SubspaceInterchange, _model_device, + arbitrary_targets, collect_activations, fit_das, intervened_logits, mean_difference_basis, probe_basis, random_basis, + reinitialized, run_with_intervention, ) @@ -322,9 +324,9 @@ def test_das_leaves_the_callers_random_stream_alone(): assert torch.equal(first.basis, again.basis) # the seed still fixes the result -def test_das_never_reseeds_a_device_generator(monkeypatch): +def _forbid_reseeding(monkeypatch): def reseeded(*args, **kwargs): - raise AssertionError("fit_das reseeded a global or device generator") + raise AssertionError("a global or device generator was reseeded") for module, name in ( (torch, "manual_seed"), @@ -334,6 +336,10 @@ def reseeded(*args, **kwargs): (torch.mps, "manual_seed"), ): monkeypatch.setattr(module, name, reseeded) + + +def test_das_never_reseeds_a_device_generator(monkeypatch): + _forbid_reseeding(monkeypatch) model, base, source, targets = _planted_pairs(n=32) fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=1, seed=5) @@ -366,3 +372,126 @@ def test_das_leaves_the_model_untouched(): assert model.scale.requires_grad and torch.equal(model.scale.detach(), before) with pytest.raises(ValueError, match="same length"): fit_das(model, "hook_mid", base, source[:-1], targets, model, rank=1) + + +def test_arbitrary_targets_are_fixed_per_key_and_seed(): + keys = [f"source {i}" for i in range(400)] + targets = arbitrary_targets(keys, n_classes=3, seed=0) + assert targets.dtype == torch.long and set(targets.tolist()) == {0, 1, 2} + assert torch.equal(targets, arbitrary_targets(keys, n_classes=3, seed=0)) + assert not torch.equal(targets, arbitrary_targets(keys, n_classes=3, seed=1)) + assert torch.equal(arbitrary_targets(keys[:5] * 2, n_classes=3, seed=0)[5:], targets[:5]) + assert all(abs(c - 400 / 3) < 50 for c in torch.bincount(targets).tolist()) + + +class _Layers(HookedRootModule): + """Only standard layers, each with reset_parameters, and a fixed buffer.""" + + def __init__(self): + super().__init__() + self.inp = nn.Linear(D_MODEL, D_MODEL) + self.norm = nn.LayerNorm(D_MODEL) + self.register_buffer("fixed", torch.randn(D_MODEL)) + self.hook_mid = HookPoint() + self.setup() + self.eval() + + +def test_reinitialized_resets_every_parameter_and_keeps_buffers(): + model = _Layers() + with torch.no_grad(): + model.norm.weight.add_(1.0) # as training would, away from LayerNorm's ones and zeros + model.norm.bias.add_(1.0) + before = {name: p.detach().clone() for name, p in model.named_parameters()} + fresh = reinitialized(model, seed=0) + for name, parameter in fresh.named_parameters(): + assert not torch.equal(parameter, before[name]), name + assert torch.equal(dict(model.named_parameters())[name], before[name]) # original untouched + assert torch.equal(fresh.fixed, model.fixed) and not fresh.training + assert torch.equal(reinitialized(model, seed=0).inp.weight, fresh.inp.weight) + + +def test_reinitialized_resets_a_root_that_resets_its_own_parameters(): + class _SelfResetting(_Layers): # the root owns a parameter and knows how to reset it + def __init__(self): + super().__init__() + self.scale = nn.Parameter(torch.full((D_MODEL,), 3.0)) + + def reset_parameters(self): + nn.init.normal_(self.scale) + + model = _SelfResetting() + fresh = reinitialized(model, seed=0) + assert not torch.equal(fresh.scale, model.scale) + + +def test_reinitialized_refuses_parameters_it_cannot_reset(): + with pytest.raises(ValueError, match="cannot reset scale"): + reinitialized(_Planted(torch.randn(D_MODEL)), seed=0) # a parameter on the root itself + + class _RawWeight(nn.Module): # like HuggingFace's Conv1D: a parameter and no reset_parameters + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(D_MODEL, D_MODEL)) + + model = _Layers() + model.conv = _RawWeight() + with pytest.raises(ValueError, match="cannot reset conv.weight"): + reinitialized(model, seed=0) + + +def test_reinitialized_leaves_every_random_stream_alone(monkeypatch): + model = _Layers() + torch.manual_seed(123) + expected = torch.rand(3) + torch.manual_seed(123) + _forbid_reseeding(monkeypatch) + reinitialized(model, seed=0) + assert torch.equal(torch.rand(3), expected) + + +def _typed_pairs(prototypes: torch.Tensor, n: int, generator: torch.Generator): + """Random bases; sources drawn from a few recurring types, so the types act as keys.""" + kinds = torch.randint(len(prototypes), (n,), generator=generator) + sources = prototypes[kinds] + 0.05 * torch.randn(n, D_MODEL, generator=generator) + return torch.randn(n, D_MODEL, generator=generator), sources, kinds + + +def _held_out_iia(model, fit, test, fit_targets, test_targets) -> float: + result = fit_das( + model, "hook_mid", fit[0], fit[1], fit_targets, model, rank=1, epochs=40, lr=5e-2 + ) + acts = collect_activations(model, "hook_mid", test[1], model) + logits = intervened_logits( + model, "hook_mid", SubspaceInterchange(result.subspace), test[0], acts, model + ) + return (logits.argmax(-1) == test_targets).float().mean().item() + + +def test_das_beats_the_arbitrary_label_control_on_recurring_keys(): + generator = torch.Generator().manual_seed(0) + model = _Planted(torch.randn(D_MODEL, generator=generator)) + prototypes = torch.randn(64, D_MODEL, generator=generator) + fit, test = _typed_pairs(prototypes, 256, generator), _typed_pairs(prototypes, 256, generator) + # Held-out sources share their keys with training sources, so the control could be learned + assert set(test[2].tolist()) <= set(fit[2].tolist()) + real = lambda sources: (sources @ model.direction > 0).long() + labels = lambda kinds: arbitrary_targets(kinds.tolist(), 2, seed=0) + assert _held_out_iia(model, fit, test, real(fit[1]), real(test[1])) > 0.9 + assert _held_out_iia(model, fit, test, labels(fit[2]), labels(test[2])) < 0.7 + + +def test_the_arbitrary_label_control_can_reject_a_subspace(): + # Types placed so that one direction fits their arbitrary labels: the control then does as + # well as the real alignment, and DAS does not beat it + generator = torch.Generator().manual_seed(0) + model = _Planted(torch.randn(D_MODEL, generator=generator)) + prototypes = torch.randn(8, D_MODEL, generator=generator) + wanted = arbitrary_targets(list(range(8)), 2, seed=0) + prototypes[(prototypes @ model.direction > 0).long() != wanted] *= -1 + fit, test = _typed_pairs(prototypes, 256, generator), _typed_pairs(prototypes, 256, generator) + real = lambda sources: (sources @ model.direction > 0).long() + labels = lambda kinds: arbitrary_targets(kinds.tolist(), 2, seed=0) + control = _held_out_iia(model, fit, test, labels(fit[2]), labels(test[2])) + assert control > 0.9 + assert _held_out_iia(model, fit, test, real(fit[1]), real(test[1])) - control < 0.1 diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index db9b450bf..82f94960f 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -20,8 +20,10 @@ - interchange: Interchange interventions on linear subspaces: set a base's activation inside a subspace to its value on a source input, with difference-of-means, probe and random bases, a norm-matched placebo, - batched gradient-free evaluation, and distributed alignment search (DAS) - to learn the subspace. + difference-of-means, probe and random bases, a norm-matched placebo, + batched gradient-free evaluation, distributed alignment search (DAS) to + learn the subspace, and DAS's trained controls (arbitrary labels and a + randomly initialised model). - jacobian_lens: The Jacobian lens (J-lens) -- per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, J-space sparse decomposition, and @@ -83,12 +85,15 @@ LearnedSubspace, NormMatchedPlacebo, SubspaceInterchange, + arbitrary_targets, collect_activations, fit_das, intervened_logits, mean_difference_basis, probe_basis, random_basis, + random_init_bridge, + reinitialized, run_with_intervention, ) from transformer_lens.tools.analysis.jacobian_lens import ( @@ -209,6 +214,7 @@ "SubspaceInterchange", "VocabularyRanking", "WeightLayout", + "arbitrary_targets", "attention_head_subspace_affinity", "attribution_patch", "collect_activations", @@ -233,7 +239,9 @@ "project_activations", "projection_kernel", "random_basis", + "random_init_bridge", "random_projection_kernel_moments", + "reinitialized", "run_causal_swap_benchmark", "run_with_intervention", "select_displacement_matched_control_token", diff --git a/transformer_lens/tools/analysis/interchange.py b/transformer_lens/tools/analysis/interchange.py index 2182d3093..54c234637 100644 --- a/transformer_lens/tools/analysis/interchange.py +++ b/transformer_lens/tools/analysis/interchange.py @@ -15,8 +15,13 @@ class means, a linear probe's weights, a random basis), and a norm-matched random-subspace placebo. :func:`fit_das` learns the basis instead, by distributed alignment search (DAS; Geiger et al., 2024): it trains an orthonormal basis so that the interchange produces the high-level model's -predicted outputs. An interchange accuracy is not evidence on its own: compare it with the placebo -and the baselines. +predicted outputs. + +An interchange accuracy is not evidence on its own. DAS fits arbitrary label mappings and finds +effects in randomly initialised models (Arora et al., 2024), so the module also provides the +controls to train the same way: :func:`arbitrary_targets` for an arbitrary-label control, and +:func:`random_init_bridge` or :func:`reinitialized` for an untrained copy of the model. Compare a +learned subspace with both, and with the placebo and the baselines. Works with any model exposing the hook API (``TransformerBridge`` or a ``HookedRootModule``). Activations at the hook must be ``[batch, d_model]``, or ``[batch, pos, d_model]`` with one @@ -25,6 +30,8 @@ class means, a linear probe's weights, a random basis), and a norm-matched rando from __future__ import annotations +import copy +import hashlib from dataclasses import dataclass from typing import Any, Callable, Hashable, Iterator, Optional, Sequence, Union, cast @@ -494,3 +501,115 @@ def fit_das( for p, flag in zip(model.parameters(), flags): p.requires_grad_(flag) return DASResult(subspace.eval(), tuple(losses), hook_name, rank) + + +def arbitrary_targets( + keys: Sequence[Hashable], n_classes: int, seed: int +) -> Int[torch.Tensor, "batch"]: + """A fixed pseudo-random class per key: targets for the arbitrary-label control. + + Train :func:`fit_das` on these instead of the high-level model's targets, with the same pairs, + rank and budget. The interchange accuracy it reaches shows how much accuracy the method's + expressivity alone produces; a learned subspace should clearly exceed it. + + Args: + keys: One key per pair identifying its source, e.g. the source text. + n_classes: Number of output classes. + seed: Seed for the mapping. + + Returns: + A class in ``[0, n_classes)`` per key; equal keys always get equal classes. + """ + + def draw(key: Hashable) -> int: + digest = hashlib.sha256(f"{seed}:{key!r}".encode()).digest() + return int.from_bytes(digest[:8], "big") % n_classes + + return torch.tensor([draw(k) for k in keys], dtype=torch.long) + + +def reinitialized(model: nn.Module, seed: int) -> nn.Module: + """A copy of ``model`` with freshly initialised parameters: the random-init control. + + For models built from standard ``nn`` layers, such as a ``HookedRootModule``: every module with + ``reset_parameters`` is reset, so every parameter must belong to one. Buffers, such as fixed + architectural matrices, are kept. For HuggingFace models use :func:`random_init_bridge`, which + applies the architecture's own initialisation scheme. + + Args: + model: The trained model; it is not modified. + seed: Seed for the new parameters. + + Returns: + The untrained copy, on the device of the model's parameters, in evaluation mode. + + Raises: + ValueError: If a parameter belongs to no module with ``reset_parameters``, so it would keep + its trained value. + """ + covered = { + id(parameter) + for module in model.modules() + if callable(getattr(module, "reset_parameters", None)) + for parameter in module.parameters(recurse=False) + } + missing = [name for name, parameter in model.named_parameters() if id(parameter) not in covered] + if missing: + more = f" and {len(missing) - 5} more" if len(missing) > 5 else "" + raise ValueError( + f"reinitialized cannot reset {', '.join(missing[:5])}{more}: no module's " + "reset_parameters covers them, so they would keep their trained values. Reset them " + "yourself, or use random_init_bridge for a HuggingFace model." + ) + device = _model_device(model) + clone = copy.deepcopy(model).cpu() + # Reset on the CPU, seeding only the CPU generator inside a fork that restores it; + # torch.manual_seed would also reseed every CUDA, MPS and XPU generator + with torch.random.fork_rng(devices=[]): + torch.default_generator.manual_seed(seed) + for module in clone.modules(): + reset = getattr(module, "reset_parameters", None) + if callable(reset): + reset() + return clone.to(device).eval() + + +def random_init_bridge( + model_name: str, + *, + seed: int, + dtype: torch.dtype = torch.float32, + device: Optional[Union[str, torch.device]] = None, + **boot_kwargs: Any, +) -> Any: + """A ``TransformerBridge`` with the architecture's HuggingFace initialisation and no training. + + The random-init control: train :func:`fit_das` on it with the same pairs, rank and budget as + on the trained model. Effects DAS finds here come from the method and the architecture, not + from anything the model learned. The model is built by ``boot_transformers`` with + ``load_weights=False``, so it has the same model class and prepared configuration as the + trained bridge. + + Args: + model_name: Model name or alias, as for ``TransformerBridge.boot_transformers``. + seed: Seed for the initialisation. + dtype: Parameter dtype. + device: Device for the bridge; by default the device ``boot_transformers`` would choose. + **boot_kwargs: Passed on to ``TransformerBridge.boot_transformers``, e.g. ``revision``. + + Returns: + The untrained ``TransformerBridge``, in evaluation mode. + """ + from transformer_lens.model_bridge import TransformerBridge + from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params + + # The model is built on the CPU, so seed only the CPU generator, inside a fork that restores it + with torch.random.fork_rng(devices=[]): + torch.default_generator.manual_seed(seed) + bridge = TransformerBridge.boot_transformers( + model_name, load_weights=False, dtype=dtype, device=device, **boot_kwargs + ) + # Building from a config follows the config's dtype and leaves the weights on the CPU, + # whatever device boot recorded; HuggingFace also leaves the model in training mode + maybe_cast_floating_params(bridge.original_model, dtype) + return bridge.to(bridge.cfg.device).eval()