Repository navigation
TransformerBridge.boot_transformers(load_weights=False) ignores dtype and returns a train-mode model #2017
Description
Bug
TransformerBridge.boot_transformers(..., load_weights=False) ignores dtype, and it returns the model in train mode. With load_weights=True, the same call honors dtype and returns an eval-mode model.
import torch
from transformer_lens.model_bridge import TransformerBridge
tiny = {"n_layer": 1, "n_embd": 8, "n_head": 2, "n_positions": 4, "vocab_size": 20}
b = TransformerBridge.boot_transformers(
"gpt2", hf_config_overrides=tiny, load_weights=False, dtype=torch.float16, device="cpu"
)
print(b.cfg.dtype, b.W_E.dtype) # torch.float16 torch.float32 <- cfg and weights disagree
print(b.original_model.training) # True <- dropout active
t = torch.tensor([[0, 2, 3, 4]])
print((b(t) - b(t)).abs().max()) # ~0.1, two identical forward passes differSeen on transformer-lens 4.3.0, and the code is the same on dev (f4f7d22).
Cause
In model_bridge/sources/transformers/source.py, the load-weights path passes torch_dtype through model_kwargs ({"config": hf_config, "torch_dtype": dtype}). The load_weights=False branch builds its own from_config_kwargs and calls model_class.from_config(prepared_config, **from_config_kwargs) without the dtype. HF's from_config also leaves the model in train mode, while from_pretrained calls .eval().
Suggested fix
In the elif not load_weights: branch:
from_config_kwargs = {"torch_dtype": dtype}
...
hf_model = model_class.from_config(prepared_config, **from_config_kwargs)
hf_model.eval()This came up while building tiny random-init bridges for SAELens tests. The workaround there is .to(dtype) plus zeroing the dropout config. I'm happy to open a PR with a test for both.
Activity
jlarson4 commentedon Nov 12, 2026
Thanks for reporting this bug! The current
load_weightsdocstring documents that this behaviour occurs, the resolution to this issue will need to make sure to update that as well. A couple notes on the fix:dtype,deviceandeval()on the no-weights path by mirroring how way the weights path accomplishes the same taskrandom_init_bridgeintools/analysis/interchange.pyhas a workaround for this issue (cast,.to,.eval()), please clean that up as wellAre you planning to work on this yourself @pablocs116? Let me know and I will get you assigned, before someone else jumps on it
pablocs116 commentedon Nov 12, 2026
Yes, I'd like to work on this. Please assign me.
You're right, I missed that the
load_weightsdocstring already documents this ("on the CPU, in the config's dtype and in training mode"). So it's a documented limitation more than a bug, anddeviceisn't applied either. Here's what I'm planning:load_weights=Falsepath, applydtype,deviceand.eval()the same way the weights path does, then update the docstring to match..to/.eval()workaround fromrandom_init_bridgeintools/analysis/interchange.py.cfg.dtype, the model is in eval mode, and two identical forward passes give the same output.random_init_bridgekeeps its current behaviour.One question: I'd leave
device_mapas it is, so it stays only recorded and not applied on the no-weights path, since dispatching a randomly initialised model across devices doesn't really make sense. Does that sound right to you, or should it follow the weights path too?jlarson4 commentedon Nov 12, 2026
Leaving
device_mapas it is is the correct call. Sounds good on the scope! Let me know if you have any questions or concerns as you work on it.