Repository navigation
fix(bridge): apply dtype, device and eval() when load_weights=False - #2020
pablocs116 wants to merge 1 commit into
Conversation
boot_transformers(load_weights=False) built the model with from_config and left it on the CPU, in the config's dtype and in training mode. It now mirrors the weights path: eval(), maybe_cast_floating_params(dtype), and a move to device once the bridge is built (fused-QKV splitting at boot needs real values, so the move can't precede it). device_map is still not applied. random_init_bridge drops its cast/.to/.eval() workaround. Closes TransformerLensOrg#2017 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for the quick turnaround on this, the eval and device handling are solid. A couple comments, and you will need to pull in the latest dev elements before this can merge. Let me know if you have any questions!
| hf_model.eval() | ||
| # Initialise in the config's dtype, then cast like the weights path, so float32 | ||
| # buffers (e.g. RotaryEmbedding.inv_freq) are preserved | ||
| maybe_cast_floating_params(hf_model, dtype) |
There was a problem hiding this comment.
from_pretrained builds the model under the requested dtype, so buffers created at init follow it. Here, the model is built in the config's dtype and only the parameters are cast, which leaves some architectures mixed. For example, no-weights XGLM or Falcon-H1 booted with dtype=torch.float16 now fail in its first forward pass, where the weights path runs. Can the model be built in the requested dtype, mirroring the weights path?
| 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. | ||
| checkpoint: real parameters with the architecture's random initialisation, moved to | ||
| ``device``, cast to ``dtype`` and in evaluation mode, as with loaded weights. |
There was a problem hiding this comment.
benchmarks/main_benchmark.py:750 boots without weights only to read attn_implementation, and passes the run's device and dtype. With this change that probe casts a full random-init model and moves it to the accelerator on every benchmark run. Could you take a look at that call so it stays a cheap CPU probe?
| torch.default_generator.manual_seed(seed) | ||
| bridge = TransformerBridge.boot_transformers( | ||
| model_name, load_weights=False, device=device, **boot_kwargs | ||
| return TransformerBridge.boot_transformers( |
There was a problem hiding this comment.
Some of my work, PR #2018 landed on dev and rewrote this block. boot is no longer passed dtype, and cfg.dtype is set by hand after the cast. Please pull in the latest dev and adjust accordingly, apologies for the conflict I did not mean to overlap with you.
Description
TransformerBridge.boot_transformers(..., load_weights=False)built the model withfrom_configand left it on the CPU, in the config's dtype and in training mode.devicewas only recorded incfg.device, anddtypewas ignored, socfg.dtypeand the weights could disagree. Callers had to cast, move and.eval()the model themselves. With GPT-2's default dropout, two identical forward passes gave different outputs.The no-weights path now does what the weights path does:
hf_model.eval().from_pretrainedreturns an eval-mode model, butfrom_configdoesn't.maybe_cast_floating_params(hf_model, dtype), the same call the weights path uses. The model is still initialised in the config's dtype and then cast, so float32 buffers such asRotaryEmbedding.inv_freqare kept.deviceright after it's built. The weights path moves the HF model before building the bridge, but on the no-weights path I move it after. Splitting fused QKV at boot reads the weight values, so moving a GPT-2 tometafirst fails withNotImplementedError.metais whattest_random_init_bridge_puts_its_weights_on_the_device_it_recordsuses to stand in for a GPU, and it's whatrandom_init_bridgeitself did (bridge.to(bridge.cfg.device)).from_configalready builds on the CPU, so a CPU target is not moved, and the move doesn't print.The
load_weightsdocstring is updated to match.random_init_bridgeintools/analysis/interchange.pydrops its cast /.to/.eval()workaround and just seeds and boots. Its existing tests pass unchanged.Closes #2017
Type of change
Checklist:
Tests
New
TestBridgeWithoutWeightsintests/integration/model_bridge/test_bridge_creation_modes.py. It uses a tiny GPT-2 config throughhf_config_overrides, so no weights are downloaded:cfg.dtypehf-internal-testing/tiny-random-LlamaForCausalLMin bfloat16: theinv_freqbuffers stay float32device="meta":cfg.deviceand every parameter are onmetaLocal gates:
make check-formatanduv run mypy .: cleanmake unit-test: 8199 passed. The one failure in the first run (test_qwen3_5_adapter.py::TestQwen3_5LoadingGuards::test_load_weights_false_uses_prepared_text_config) came from this change: its stub bridge has nocfg, and.to()reads it. It's fixed by skipping the CPU move, and the test now passes.test_bridge_creation_modes.py,test_interchange.py,test_qwen3_5_adapter.py,test_multi_gpu_bridge.py,test_n_ctx_override.py,test_remote_bridge.py: 162 passed, 3 skippedHardware: CPU only (WSL2).
🤖 Generated with Claude Code