Skip to content

Docs: Jlens Qwen3.5-4b Demo - #1547

Open
KYinXu wants to merge 4 commits into
TransformerLensOrg:devfrom
KYinXu:docs/jlens-Qwen-3.5-4b-demo
Open

Docs: Jlens Qwen3.5-4b Demo#1547
KYinXu wants to merge 4 commits into
TransformerLensOrg:devfrom
KYinXu:docs/jlens-Qwen-3.5-4b-demo

Conversation

@KYinXu

@KYinXu KYinXu commented Jul 27, 2026

Copy link
Copy Markdown

Description

Added a new walkthrough section using Qwen3.5-4b in demos/Jacobian_Lens_Demo.ipynb.

Key additions:

  • Demonstrates support for bridge-only model architectures within the Jacobian Lens pipeline.
  • Reproduces the multi-hop (twohop) reasoning experiment previously demonstrated on gemma-2-2b.

Fixes #1539 Tier 2 - second model demo

Type of change

  • Documentation update

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

@KYinXu
KYinXu marked this pull request as draft July 27, 2026 16:58
@KYinXu

KYinXu commented Jul 27, 2026

Copy link
Copy Markdown
Author

When loading Qwen-3.5 I received the following warning (x24+):

[/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:494](https://file+.vscode-resource.vscode-cdn.net/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:494): UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.
  getattr(module, "_register_aliases")()

The warning is harmless since we're not performing any hook actions on the hybrid layers, but it is something of note for clarity in the notebook cell as it becomes clutter and also as a potential ticket to implement hook support on Gated DeltaNets.

@KYinXu

KYinXu commented Jul 27, 2026

Copy link
Copy Markdown
Author

Also sidenote, notebook cell runs were done locally without a CUDA GPU or Colab's TPUs, so if default cell output is a desired clarity fix then I can look into rerunning remote.

@KYinXu
KYinXu marked this pull request as ready for review July 27, 2026 17:59
@emerardd

Copy link
Copy Markdown

Hi @KYinXu — thanks for adding the Qwen3.5 walkthrough. I tested it locally in a combined tree with current dev and #1545. The full notebook passed nbval (21 passed in 659.80s), and the Qwen readout worked with all 32 block-output hooks available.

I noticed two reproducibility/documentation points:

  1. The notebook describes QWEN_LAYERS as full-attention layers, but the live Bridge reports layers 8, 16, 24, and 30 as GatedDeltaNet linear-attention layers; only 31 is full attention. The full-attention layers are [3, 7, 11, 15, 19, 23, 27, 31]. Would it make sense either to update the wording or select layers such as [7, 15, 23, 31]?

  2. Could the Qwen model revision also be pinned? My run resolved to 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a. The lens revision is pinned, but its metadata does not contain a model revision, so it cannot catch future model-repository drift.

One integration note: #1545 and this PR both modify the same notebook from the same earlier base, so whichever lands second will need a cell-level conflict resolution and another full nbval run.

Overall, the Qwen path worked successfully in my local validation.

@jlarson4

jlarson4 commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Thanks for putting this together @KYinXu! And thank you @emerardd for mentioning the conflict potential with your PR, valid point that I was going to address.

We will want the default cell output, the goal is to add this demo to CI as validation, so we will need the output to match.

#1545 has already merged we will want to rebase this branch onto dev to avoid conflicts. In that pass please leave the earlier gemma cells' outputs, the base branch recorded them.

A few asks for the rebase:

  1. The final comparison cell has no committed output at all, as you noted. Ideally, we'd like every cell executed and nbval-clean.
  2. +1 to both of @emerardd's points: pin the model revision, and fix the full-attention wording / QWEN_LAYERS choice.
  3. Add a one-line comment on why the lens load bypasses the registry short name (the registry entry points at the non-n1000 artifact).
  4. Update the header memory claim. Qwen3.5-4B bf16 weights alone are ~8 GB, so it won't fit an 8 GB card; say ~10 GB GPU or a CPU run.
  5. Two small requests for the comparison cell: there's trailing whitespace inside the QWEN_LAYERS list literal (first line), and the final print hardcodes top_j[31]. Let's use model.cfg.n_layers - 1 (or reuse the last entry of QWEN_LAYERS) so the cell doesn't silently break if this gets reused for a different model.

The GatedDeltaNet hook-support ticket is a good idea, would you be willing to open an issue?

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.

3 participants