Interpretune Circuit Tracer Tutorial#
Intro#
Interpretune is a flexible framework for exploring, analyzing and tuning llm world models. In this tutorial, we’ll walk through a simple example of using Interpretune to pursue interpretability research with Circuit Tracer. As we’ll see, Interpretune handles the required execution context composition, allowing us to use the same code in a variety of contexts, depending upon the level of abstraction required.
As a long-time PyTorch and PyTorch Lightning contributor, I’ve found the PyTorch Lightning framework is the right level
of abstraction for a large variety of ML research contexts, but some contexts benefit from using core PyTorch directly.
Additionally, some users may prefer to use the core PyTorch framework directly for a wide variety of reasons including
maximizing portability. As will be demonstrated here, Interpretune maximizes flexibility and portability by adhering to
a well-defined protocol that allows auto-composition of our research module with the adapters required for execution in
a wide variety of contexts. In this example, we’ll be executing the same module with core PyTorch and PyTorch Lightning,
demonstrating the use of Circuit Tracer w/ Interpretune for circuit discovery and interpretability research.
Note - this is a WIP, but this is the core idea. If you have any feedback, please let me know!
A note on memory usage#
In these exercises, we’ll be loading language models into memory for circuit analysis. It’s useful to have functions which can help profile memory usage for you, so that if you encounter OOM errors you can try and clear out unnecessary models. For example, we’ve found that with the right memory handling (i.e. deleting models and objects when you’re not using them any more) it should be possible to run all the exercises in this material on a Colab Pro notebook.
See this dropdown for some functions which you might find helpful, and how to use them.
First, we can run some code to inspect our current memory usage. Here’s an example of running this code during circuit analysis exercises.
# Profile memory usage
import torch
import gc
if torch.cuda.is_available():
print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
print(f"GPU Memory Free: {(torch.cuda.memory_reserved() - torch.cuda.memory_allocated()) / 1024**3:.2f} GB")
If you need to free up memory, you can delete large objects and run garbage collection:
# Delete large objects if needed
# del model
# del circuit_tracer_session
# Move objects to CPU if needed
THRESHOLD = 0.1 # GB
for obj in gc.get_objects():
try:
if isinstance(obj, torch.nn.Module):
# Calculate approximate size
total_params = sum(p.numel() for p in obj.parameters())
if total_params * 4 / 1024**3 > THRESHOLD: # Assuming float32
if hasattr(obj, "cpu"):
obj.cpu()
except:
pass
# Force garbage collection
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
This approach helps manage memory when working with large language models during circuit analysis.
Notebook Parameters#
This cell contains parameters that can be modified for different test configurations using papermill.
# Parameters - These will be injected by papermill during parameterized test runs
use_baseline_salient_logits = True # logits computation mode: True->salient logits, False->specific logits
use_baseline_transcoder_arch = (
False # transcoder architecture: True->SingleLayerTranscoder, False->CrossLayerTranscoder
)
# `circuit-tracer` backend configuration parameters
backend = "transformerlens" # Options: "transformerlens", "nnsight"
use_remote_execution = False # Only applicable for backend="nnsight"
core_log_dir = None # Directory to save analysis logs (if None, a temp directory will be created)
Imports#
# Core imports
import interpretune as it # registered analysis ops will be available as it.<op> when analysis is imported
# Import circuit tracer and required modules
from transformer_lens import ActivationCache # noqa: F401
from pprint import pformat
from datetime import datetime
from it_examples import _ACTIVE_PATCHES # noqa: F401 # TODO: add note about this unless patched in SL before release
from it_examples.example_module_registry import MODULE_EXAMPLE_REGISTRY # TODO: move to hub once implemented
from interpretune import ITSessionConfig, ITSession
from interpretune.base.call import it_init
from it_examples.utils.example_helpers import required_os_env
env_path: str | None = None # set to '/full/path/to/.env' to override
# Maybe load environment variables from an env file.
os_env_reqs = None
assert required_os_env(env_path=env_path, env_reqs=os_env_reqs)
Configure our IT Session#
Here we define or customize our session configuration, which includes:
Experiment/task module and datamodule (in this case, ‘rte’ for the RTE task)
We can customize any module, datamodule, or adapter-specific configuration options we want to use. In this case, we set target
circuit_tracer_cfgthat we want to use for our analysis. We also could customize generation parameters, tokenization, the pretrained/config-based model we want to use (in this case, GPT2) etc.
The adapter context we want to use. In this case,
corePyTorch (vs e.g. Lightning) andcircuit_tracer(vs e.g.transformer_lensorsae_lens).
When an ITSession is created, the selected adapter context will trigger composition of the relevant adapters with our experiment/task module and datamodule. The intention of this abstraction is to enable the same experiment/task logic to be used unchanged across a broad variety of PyTorch framework and analytical package contexts.
# Load our demo config (this will be done from the hub once that is available)
base_itdm_cfg, base_it_cfg, dm_cls, m_cls = MODULE_EXAMPLE_REGISTRY.get("gemma2.rte_demo.circuit_tracer")
# Configure backend selection
base_it_cfg.circuit_tracer_cfg.backend = backend
if backend == "nnsight":
base_it_cfg.circuit_tracer_cfg.nnsight_remote = use_remote_execution
print(f"Backend: {backend}, Remote execution: {use_remote_execution}")
if use_remote_execution:
print(" NDIF_API_KEY will be resolved from environment variable")
else:
print(f"Backend: {backend}")
# Optionally override base_it_cfg.core_log_dir with the notebook parameter if provided
if core_log_dir:
base_it_cfg.core_log_dir = core_log_dir
# If the user requests the baseline salient-logits path, clear the explicit target
# token configuration so the adapter will fall back to the default compute_salient_logits
# implementation. This forces usage of the baseline salient-logits computation.
if use_baseline_salient_logits:
# Clear any explicit token selection so compute_salient_logits() runs its default path
base_it_cfg.circuit_tracer_cfg.analysis_target_tokens = None
base_it_cfg.circuit_tracer_cfg.target_token_ids = None
print(
"use_baseline_salient_logits=True: cleared analysis_target_tokens and "
"target_token_ids -> using default compute_salient_logits path"
)
else:
print("use_baseline_salient_logits=False: keeping configured analysis_target_tokens / target_token_ids (if any)")
# Configure transcoder architecture selection based on the toggle. When True, use the 'gemma'
# CrossLayerTranscoder (demo). When False, point to the HF SingleLayerTranscoder checkpoint URL.
if use_baseline_transcoder_arch:
base_it_cfg.circuit_tracer_cfg.transcoder_set = "gemma"
print(
"use_baseline_transcoder_arch=True: set transcoder_set='gemma' -> "
"set transcoder_set to HF URL -> using the SingleLayerTranscoder checkpoint"
)
else:
base_it_cfg.circuit_tracer_cfg.transcoder_set = "mntss/clt-gemma-2-2b-426k"
print(
"use_baseline_transcoder_arch=False: demo will use CrossLayerTranscoder "
"instead of the default TranscoderSet of `SingleLayerTranscoder`s"
)
print(pformat(base_it_cfg.circuit_tracer_cfg))
# configure our session with our desired adapter composition, core and circuit_tracer in this case
session_cfg = ITSessionConfig(
adapter_ctx=(it.Adapter.core, it.Adapter.circuit_tracer),
datamodule_cfg=base_itdm_cfg,
module_cfg=base_it_cfg,
datamodule_cls=dm_cls,
module_cls=m_cls,
)
# start our session
it_session = ITSession(session_cfg)
print("\nIT Session created successfully!")
Backend: transformerlens
use_baseline_salient_logits=True: cleared analysis_target_tokens and target_token_ids -> using default compute_salient_logits path
use_baseline_transcoder_arch=False: demo will use CrossLayerTranscoder instead of the default TranscoderSet of `SingleLayerTranscoder`s
CircuitTracerConfig(backend='transformerlens',
model_name=None,
transcoder_set='mntss/clt-gemma-2-2b-426k',
dtype=torch.bfloat16,
max_n_logits=10,
desired_logit_prob=0.95,
batch_size=256,
max_feature_nodes=8192,
offload='cpu',
lazy_encoder=None,
lazy_decoder=True,
verbose=True,
default_node_threshold=0.8,
default_edge_threshold=0.98,
save_graphs=True,
graph_output_dir=None,
analysis_target_tokens=None,
target_token_ids=None,
use_neuronpedia=False,
intervention_scale_factor=1.0,
intervention_max_influence_norm_scale=False,
intervention_sign_aware_scale=True,
intervention_value=None,
intervention_value_source='top_feature_scores',
intervention_constrained_layers=None,
intervention_freeze_attention=None,
intervention_apply_activation_function=None,
intervention_sparse=False,
intervention_return_activations=False,
nnsight_remote=False,
ndif_api_key=None)
IT Session created successfully!
[INFO] interpretune.utils.logging: Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute 'quantization_config'
INFO:interpretune.utils.logging:Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute 'quantization_config'
[INFO] interpretune.utils.logging: Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute '_pre_quantization_dtype'
INFO:interpretune.utils.logging:Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute '_pre_quantization_dtype'
# manual init for now
it_init(**it_session)
print("\nIT Session initialized successfully!")
[INFO] interpretune.utils.logging: Preparing data: InterpretunableDataModule
INFO:interpretune.utils.logging:Preparing data: InterpretunableDataModule
[INFO] interpretune.utils.logging: The following columns don't have a corresponding argument in `Gemma2ForCausalLM.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `Gemma2ForCausalLM.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns don't have a corresponding argument in `Gemma2ForCausalLM.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `Gemma2ForCausalLM.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: The following columns don't have a corresponding argument in `Gemma2ForCausalLM.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `Gemma2ForCausalLM.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns don't have a corresponding argument in `Gemma2ForCausalLM.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `Gemma2ForCausalLM.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: The following columns don't have a corresponding argument in `Gemma2ForCausalLM.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `Gemma2ForCausalLM.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns don't have a corresponding argument in `Gemma2ForCausalLM.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `Gemma2ForCausalLM.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: Setting up datamodule: InterpretunableDataModule
INFO:interpretune.utils.logging:Setting up datamodule: InterpretunableDataModule
[INFO] interpretune.utils.logging: Setting up model: InterpretunableModule
INFO:interpretune.utils.logging:Setting up model: InterpretunableModule
[INFO] interpretune.utils.logging: initializing optimizers and schedulers: InterpretunableModule
INFO:interpretune.utils.logging:initializing optimizers and schedulers: InterpretunableModule
[INFO] interpretune.utils.logging: Input gradient requirements handled by circuit tracer internally.
INFO:interpretune.utils.logging:Input gradient requirements handled by circuit tracer internally.
IT Session initialized successfully!
Basic Attribution Graph#
from tqdm.auto import tqdm
limit_analysis_batches = 1
test_token_limit = -1
force_manual_debug_prompts = True # Set to True to use manual debug prompts instead of random samples
# specific tokens to analyze, will use tokens associated with top `max_n_logits` if `None`
# analysis_target_tokens: Optional[torch.Tensor] = None
example_prompts = []
ct_module = it_session.module
if not force_manual_debug_prompts:
dataloader = it_session.datamodule.test_dataloader()
for epoch_idx in range(1): # Run for a single epoch for simplicity
ct_module.current_epoch = epoch_idx
for batch_idx, batch in tqdm(enumerate(dataloader)):
if batch_idx >= limit_analysis_batches >= 0:
break
# fetch the first test_token_limit from the first example in the batch
first_ex_in_batch = batch[:1]
first_ex_in_batch = first_ex_in_batch["input"]
first_ex_in_batch.squeeze_()
if test_token_limit > 0:
first_ex_in_batch = first_ex_in_batch[-test_token_limit:]
first_ex_in_batch = first_ex_in_batch[first_ex_in_batch != 0]
example_prompts.append(first_ex_in_batch)
else:
# Generate attribution graphs for a few example prompts
example_prompts = [
# "The capital of France is",
"The capital of the state containing Dallas is",
# "When I look at the sky, I see",
]
print("Generating attribution graphs for example prompts...")
slug_base = "it_circuit_tracer_compute_specific_logits_demo"
results = []
for i, prompt in enumerate(example_prompts):
print(f"\nProcessing prompt {i + 1}: '{prompt}'")
slug = f"{slug_base}_{i + 1}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
# Process the batch using the session, the adapter will handle tokenization and graph generation
try:
graph, local_graph_path, _ = ct_module.generate_graph(prompt=prompt, slug=slug)
results.append(local_graph_path)
except Exception as e:
print(f" - Error processing prompt: {e}")
print(f"\nProcessed {len(results)} prompts successfully")
Generating attribution graphs for example prompts...
Processing prompt 1: 'The capital of the state containing Dallas is'
- Error processing prompt: ReplacementModel not loaded. Call _load_replacement_model() first.
Processed 0 prompts successfully
Saving and Visualizing Attribution Graphs#
In this section, we’ll demonstrate how to save the generated attribution graphs and prepare them for visualization. The CircuitTracerAdapter integrates with Interpretune’s AnalysisStore to persistently store graph data.
import os
import socket
from circuit_tracer.frontend.local_server import serve
enable_iframe = False # whether to enable the IFrame display or not
port = 8046
server = serve(data_dir=ct_module.circuit_tracer_cfg.graph_output_dir, port=port)
port_forwarding = False # whether to use port forwarding or not
# Host used to build the graph URL. Defaults to this machine's hostname, which is what you want when
# viewing from another machine on the same network. Override with IT_GRAPH_SERVER_HOST, or set
# port_forwarding=True above to use localhost (e.g. over an SSH tunnel).
hostname = os.environ.get("IT_GRAPH_SERVER_HOST") or socket.gethostname()
if port_forwarding:
hostname = "localhost" # use localhost for port forwarding
print(
f"Using port forwarding (ensure it is configured) and localhost."
f" Open your graph here at http://{hostname}:{port}/index.html"
)
else:
print(
f"Not using port forwarding. Use the IFrame below, or"
f" open your graph here directly at http://{hostname}:{port}/index.html"
)
if enable_iframe:
from IPython.display import IFrame
# Display the IFrame with the graph visualization
print(f"Displaying graph visualization in IFrame at http://{hostname}:{port}/index.html")
display(IFrame(src=f"http://{hostname}:{port}/index.html", width="100%", height="800px"))
Not using port forwarding. Use the IFrame below, or open your graph here directly at http://speediedl:8046/index.html
server.stop()
Next Steps#
This notebook demonstrates the basic CircuitTracerAdapter integration with Interpretune, including session management, configuration, graph visualization, and adapter composition.
For more advanced circuit analysis workflows, see:
CT Analysis Backend Demo — Full circuit-tracer analysis ops pipeline using the DISPATCHER: concept direction → attribution graph → node influence → top features → feature intervention
Concept-Direction Steering Demo — Store- and embed-based sign-aware multi-feature steering (orange example, local 262k Monology dashboards) composition combining GPT-2 SAE analysis (TransformerBridge) with Gemma-2 circuit analysis (CT NNsight), demonstrating AnalysisStore persistence and backend-agnostic result composition