You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Document the recently-added bridge capabilities in docs/source/content/migrating_to_v3.md, extending the existing ## Migrating specific HookedTransformer APIs section (heading at :211, table at :221-225): start_at_layer/input_to_embed, caching hooks (get_caching_hooks/add_caching_hooks), pos_slice, incl_bwd, HF-native past_key_values/logits_and_cache, mod_dict, reset_hooks parameters, and the model_protocol typing. Make sure you're based on dev-4.x, where these changes exist
Motivation
These primitives shipped on TransformerBridge with tests but zero docs. The migration guide's ## Migrating specific HookedTransformer APIs section explicitly invites new rows ("when you (or we) work out the equivalent, add a row below"), but currently holds only three rows (W_pos, W_E_pos, from_pretrained_no_processing).
Pitch
Add one table row per capability to the table at migrating_to_v3.md:221-225, and for the four that need real code, a short prose recipe beneath the table (line 225 is currently end-of-file): start_at_layer resume-from-layer, caching hooks (+pos_slice/incl_bwd), past_key_values/logits_and_cache incremental decoding, and model_protocol typing for downstream libraries.
Also add a config-construction row: HookedTransformerConfig(...) → TransformerBridgeConfig(...) — construct the bridge config directly with the same core fields; legacy config objects are deliberately not accepted by boot_native (it raises a clear TypeError — draft Attn only #20). This is the train-from-scratch on-ramp for the toy-model demos.
Also add a calling-convention row: model.all_head_labels() (HT method) → bridge.all_head_labels (property — drop the parens). The bridge deliberately does not add a callable shim; calling the property raises TypeError: 'list' object is not callable, which this row preempts.
Also add a tokenizer row: model.set_tokenizer(tok) → not ported; the tokenizer is fixed at boot — pass tokenizer= to boot_transformers(...) or re-boot. Warn that raw bridge.tokenizer = tok assignment skips the d_vocab/BOS wiring (see draft Fix bug in bwd hooks #17).
Copy snippets from the shipped integration tests rather than inventing them.
Write these four caveats verbatim next to the relevant recipe: start_at_layer excludes embedding-stage hooks and still runs blocks 0..k-1 on a discarded path; incl_bwd needs a scalar output (return_type="loss") and the transformers driver only; the KV-cache is HF's own cache object, not TransformerLensKeyValueCache; the caching-hooks recipe should steer users to get_caching_hooks/add_caching_hooks/run_with_cache — the bridge's cache_all/cache_some mirrors emit DeprecationWarning (bridge_core.py:705,727).
Cross-link existing anchors where relevant: ### Will my numbers match HookedTransformer? (:84), ## Hook names (:109), ## APIs that are unchanged (:132).
Acceptance:
Each capability has a table row; the four code-bearing ones have runnable HT-to-bridge recipes
The four caveats appear next to the relevant recipe
/build-docs succeeds; doctests pass
make unit-test passes
uv run mypy . passes
Mirror: transformer_lens/model_bridge/bridge_core.py:627 (get_caching_hooks; add_caching_hooks at :669); transformer_lens/model_bridge/transformer_bridge.py:1358 (start_at_layer param; implementation _setup_start_at_layer at :1656); transformer_bridge.py:1362 (past_key_values), :1618 (logits_and_cache); bridge_core.py:129 (mod_dict), :585 (reset_hooks); transformer_lens/model_protocol.py:24,:45. Tests to copy snippets from: tests/integration/model_bridge/test_bridge_{start_at_layer,caching_hooks,manual_kv_cache,kv_cache_parity}.py.
Checklist
I have checked that there is no similar issue in the repo (required)
Proposal
Document the recently-added bridge capabilities in
docs/source/content/migrating_to_v3.md, extending the existing## Migrating specific HookedTransformer APIssection (heading at:211, table at:221-225):start_at_layer/input_to_embed, caching hooks (get_caching_hooks/add_caching_hooks),pos_slice,incl_bwd, HF-nativepast_key_values/logits_and_cache,mod_dict,reset_hooksparameters, and themodel_protocoltyping. Make sure you're based ondev-4.x, where these changes existMotivation
These primitives shipped on
TransformerBridgewith tests but zero docs. The migration guide's## Migrating specific HookedTransformer APIssection explicitly invites new rows ("when you (or we) work out the equivalent, add a row below"), but currently holds only three rows (W_pos,W_E_pos,from_pretrained_no_processing).Pitch
migrating_to_v3.md:221-225, and for the four that need real code, a short prose recipe beneath the table (line 225 is currently end-of-file):start_at_layerresume-from-layer, caching hooks (+pos_slice/incl_bwd),past_key_values/logits_and_cacheincremental decoding, andmodel_protocoltyping for downstream libraries.HookedTransformerConfig(...)→TransformerBridgeConfig(...)— construct the bridge config directly with the same core fields; legacy config objects are deliberately not accepted byboot_native(it raises a clearTypeError— draft Attn only #20). This is the train-from-scratch on-ramp for the toy-model demos.model.all_head_labels()(HT method) →bridge.all_head_labels(property — drop the parens). The bridge deliberately does not add a callable shim; calling the property raisesTypeError: 'list' object is not callable, which this row preempts.model.set_tokenizer(tok)→ not ported; the tokenizer is fixed at boot — passtokenizer=toboot_transformers(...)or re-boot. Warn that rawbridge.tokenizer = tokassignment skips thed_vocab/BOS wiring (see draft Fix bug in bwd hooks #17).start_at_layerexcludes embedding-stage hooks and still runs blocks0..k-1on a discarded path;incl_bwdneeds a scalar output (return_type="loss") and the transformers driver only; the KV-cache is HF's own cache object, notTransformerLensKeyValueCache; the caching-hooks recipe should steer users toget_caching_hooks/add_caching_hooks/run_with_cache— the bridge'scache_all/cache_somemirrors emitDeprecationWarning(bridge_core.py:705,727).### Will my numbers match HookedTransformer?(:84),## Hook names(:109),## APIs that are unchanged(:132).Acceptance:
/build-docssucceeds; doctests passmake unit-testpassesuv run mypy .passesMirror:
transformer_lens/model_bridge/bridge_core.py:627(get_caching_hooks;add_caching_hooksat:669);transformer_lens/model_bridge/transformer_bridge.py:1358(start_at_layerparam; implementation_setup_start_at_layerat:1656);transformer_bridge.py:1362(past_key_values),:1618(logits_and_cache);bridge_core.py:129(mod_dict),:585(reset_hooks);transformer_lens/model_protocol.py:24,:45. Tests to copy snippets from:tests/integration/model_bridge/test_bridge_{start_at_layer,caching_hooks,manual_kv_cache,kv_cache_parity}.py.Checklist