Feature/telemetry demo notebook

[jonathanrbelanger-lang]

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

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

[jonathanrbelanger-lang]

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!

[jlarson4]

Thank you for putting this together @jonathanrbelanger-lang, it looks awesome. I should have time today to give it a thorough review & send over any comments if I have them.

[jonathanrbelanger-lang]

Hey - I've pushed the revised Realtime_Training_Telemetry_Demo.ipynb to align with your review.

I know it is a busy weekend for everyone, so please no rush at all on reviewing this, take your time. I am completely prepared to make additional adjustments if you find them necessary.

Here is a quick breakdown of what was updated:

Changelog: Real-time Telemetry Demo Refactor
Migrated Visualizations to Plotly: Completely removed matplotlib dependencies, replacing them with a dual-trace Plotly subplot (Loss/Coherence Line Graph + Layer Depth Heatmap) to perfectly align with the TransformerLens ecosystem.

Cross-Platform Environment Detection: Implemented the standard IN_COLAB try/except block to dynamically assign the correct Plotly renderer ("colab" vs "notebook_connected") and handle pip installations seamlessly across standard Jupyter, VS Code, and Google Colab.

Real-Time Rendering Optimization: Replaced static plotting with a highly optimized real-time loop. To prevent browser DOM crashes and memory leaks, the loop pre-allocates numpy arrays and directly mutates the fig.data traces in-place, relying on IPython.display.clear_output(wait=True) for smooth frame redraws.

Compute Efficiency (Selective Caching): Modified the training loop to only invoke model.run_with_cache during defined logging intervals (every 10 steps), ensuring the extraction process does not suffocate local CPU/GPU memory bandwidth.

Self-Contained Execution: Added a localized synthetic induction data generator (get_batch()) directly above the training loop to guarantee the cell runs cleanly without state leakage or variable dependency issues from upstream cells.

Pedagogical Markdown Overhaul: Rewrote all cell headers to match the rigorous, educational tone of Main_Demo.ipynb. Added specific architectural context detailing:

Why the model uses 2 layers/2 heads (minimum depth for induction).

Why special tokens (BOS) act as attention sinks and how they skew coherence metrics.

Explicit instructions on how to revert to static rendering for users prioritizing raw compute speed.

[jlarson4]

Hey @jonathanrbelanger-lang! Just a couple small additional notes:

1. Can you remove the old notebook file (TL_Demo_RT_Viz)?
2. Can you move all of your commented # @markdown comments out of the code cells and into their own text/markdown cells?

Other than that, it looks great! Once those edits come through I will merge and get this released.

[jonathanrbelanger-lang]

Absolutely, this will be done by EOD.

[jlarson4]

Excellent, thank you for the update @jonathanrbelanger-lang

[jonathanrbelanger-lang]

@jlarson4 - separated the instructional portions, now shows explanation and clean code cells separately. One open question regarding the HookedTransformer change.

Changelog

• Removed unused imports (os, torch.nn.functional, torch.optim, collections) from setup cell
• Added # noqa: F811 to numpy re-import in training cell (intentional standalone-cell portability pattern)
• Removed deprecated demo notebooks (TL_Demo_RT_Viz.ipynb, prior Realtime_Training_Telemetry_Demo.ipynb)
• Note: HookedTransformer/HookedTransformerConfig flagged as deprecated in TL 3.0 — TransformerBridge does not appear to cover from-scratch toy model instantiation; flagging for maintainer confirmation before addressing

[jlarson4]

Hi @jonathanrbelanger-lang! Thanks for getting this wrapped up.

Good question on toy models. This is a known TransformerBridge coverage gap, it doesn't currently have a from-scratch / config-only constructor, so HookedTransformer(cfg) is still the right tool for toy models in training demos. Please keep it as-is. From-scratch instantiation is on my roadmap for future bridge work; in the meantime, no action needed here.