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.

[jlarson4]

A couple file wide notes:

1. Can we rename the demo to Realtime_Training_Telemetry_Demo
2. Can you add some detailed text cells about what/how/why this notebook is doing? A majority of the demo notebooks include some level of explanation, see Main_Demo.ipynb for an example of what I'm looking for here
3. Can we update the setup cell to function similar to the setup cells in our other notebooks, so that this notebook can be run locally via Jupyter or in Colab? Specifically, this new demo uses matplotlib which is not installed by default in TransformerLens, and the demo cannot be run locally because of it.

Let me know if you have any questions

[jonathanrbelanger-lang]

Yes, absolutely. Let me get to work on that, probably tomorrow evening as I'm knee deep in another project tonight, if that's amenable.

[jlarson4]

Absolutely, no rush. Thank you!

[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.