Skip to content

migrate the othello off hooked transformer - #1667

Merged
jlarson4 merged 4 commits into
TransformerLensOrg:dev-4.xfrom
MdSadiqMd:sadiq/migrate-the-othello
Aug 13, 2026
Merged

migrate the othello off hooked transformer#1667
jlarson4 merged 4 commits into
TransformerLensOrg:dev-4.xfrom
MdSadiqMd:sadiq/migrate-the-othello

Conversation

@MdSadiqMd

Copy link
Copy Markdown
Contributor

Description

Migrate demos/Othello_GPT.ipynb from HookedTransformer(cfg) to TransformerBridge.boot_native(cfg), unblocking HookedTransformer deprecation at 4.0.

Motivation: The Othello demo is the public face of train-from-scratch and is in the active CI matrix. It instantiates HookedTransformer(cfg) directly, which will break at 4.0 when HT is deprecated.

Changes:

  1. New convert_tl_checkpoint() utility (transformer_lens/conversion_utils/tl_checkpoint.py):

    • Converts TL-format state_dicts (keys like embed.W_E, blocks.0.attn.W_Q) to native bridge format (keys like tok_embed.weight, layers.0.attn.q.weight)
    • Handles attention weight reshaping: [n_heads, d_model, d_head][n_heads*d_head, d_model]
    • Handles MLP/unembed transposition
  2. Updated demos/Othello_GPT.ipynb:

    • Changed imports to use TransformerBridgeConfig, TransformerBridge, convert_tl_checkpoint
    • Model creation now uses TransformerBridge.boot_native(cfg)
    • Checkpoint loading chains through convert_tl_checkpoint() before load_state_dict()
    • Both main path and LOAD_AND_CONVERT_CHECKPOINT path updated
  3. Unit tests (tests/unit/conversion_utils/test_tl_checkpoint.py):

    • test_convert_tl_checkpoint_produces_matching_outputs: Verifies HT → bridge conversion produces identical outputs
    • test_convert_tl_checkpoint_key_mapping: Verifies all keys are correctly mapped
    • test_convert_tl_checkpoint_shape_transformations: Verifies tensor reshaping is correct

Verification:

  • Notebook output matches expected (argmax identical to original)
  • All 13 notebook cells pass --nbval-lax

Closes #1663

Type of change

  • New feature (non-breaking change which adds functionality)

Screenshots

N/A

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

Copy link
Copy Markdown
Collaborator

Great work @MdSadiqMd, thanks for your help! Merging now

@jlarson4
jlarson4 merged commit 3d59a51 into TransformerLensOrg:dev-4.x Aug 13, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants