Skip to content

[Proposal] Migration recipes for the new bridge capabilities in migrating_to_v3.md #1557

Description

@jlarson4

Proposal

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)

Metadata

Metadata

Assignees

No one assigned

    Labels

    TransformerBridgeBug specific to the new TransformerBridge systemcomplexity-simpleSimple issues, which may be good for beginnersdocumentationImprovements or additions to documentationgood first issueGood for newcomershelp wantedExtra attention is neededv4.x

    Type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions