Skip to content

Feature/telemetry demo notebook#1308

Open
jonathanrbelanger-lang wants to merge 5 commits into
TransformerLensOrg:devfrom
jonathanrbelanger-lang:feature/telemetry-demo-notebook
Open

Feature/telemetry demo notebook#1308
jonathanrbelanger-lang wants to merge 5 commits into
TransformerLensOrg:devfrom
jonathanrbelanger-lang:feature/telemetry-demo-notebook

Conversation

@jonathanrbelanger-lang
Copy link
Copy Markdown

@jonathanrbelanger-lang jonathanrbelanger-lang commented May 17, 2026

Description

Adds a new educational demo notebook (demos/TL_Demo_RT_Viz.ipynb) that provides a lightweight, zero-dependency bridge to extract and visualize mechanistic telemetry (Attention Coherence and Head Agreement) during a training loop.

Motivation and Context:

  • Optimization: The training loop is intentionally branched so model.run_with_cache is only called at log intervals, saving roughly 10x memory/compute overhead compared to naive caching loops.
  • Scalability: The dynamic dictionary logging and visualization matrix automatically scale to adapt to n_layers, making it highly forkable for users experimenting with larger architectures.
  • Linting: The notebook has been run through ruff and passes all modern syntax and formatting checks cleanly.

Fixes # N/A

Type of change

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

Screenshots

RT_Viz_Example

Checklist:

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

brendanlong and others added 5 commits April 20, 2026 14:50
* Fix type of HookedTransformerConfig.device

This is typed as `Optional[str]` but sometimes returns `torch.device`.
Updated the code to just return the `str` instead of wrapping with a
device.

I'm not confident that every function which takes a device will
always be passed a string, so I didn't change functions like
warn_if_mps.

Found while working on TransformerLensOrg#1219

* more cleanup

* 3.0 CI Bugs (TransformerLensOrg#1261)

* Fixing `utils` imports

* skip gated notebooks on PR from forks

* Updating notebooks

* Ensure LLaMA only runs when HF_TOKEN is available

---------

Co-authored-by: jlarson4 <jonahalarson@comcast.net>
@jonathanrbelanger-lang
Copy link
Copy Markdown
Author

jonathanrbelanger-lang commented May 17, 2026

Hey @jlarson4 -- I've opened this PR to address the task (1148) assigned to me.

You'll notice a few minor changes from my initial concept and code. These updates focus specifically on streamlining the loop and eliminating caching overhead, but the final result fully aligns with the original submission goals.

I'll be available this week to tweak or refactor anything based on your critical review. Thanks!

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