interpretune.config#

class interpretune.config.AnalysisArtifactCfg(*, latent_effects_graphs=True, latent_effects_graphs_per_batch=False, table_per_latent_model=True, top_k_latents_table=2, top_k_latent_dashboards=1, top_k_clean_logit_diffs=10)[source]#

Configuration for analysis artifacts and visualizations.

Parameters:

  • latent_effects_graphs (bool)

  • latent_effects_graphs_per_batch (bool)

  • table_per_latent_model (bool)

  • top_k_latents_table (int)

  • top_k_latent_dashboards (int)

  • top_k_clean_logit_diffs (int)

class interpretune.config.AnalysisCfg(*, output_store=None, input_store=None, batch_inputs=<factory>, run_inputs=<factory>, target_op=None, output_schema=None, name=None, fwd_hooks=<factory>, bwd_hooks=<factory>, cache_dict=<factory>, names_filter=None, save_prompts=False, save_tokens=False, decode_kwargs=<factory>, latent_analysis_targets=None, ignore_manual=False, step_fn='analysis_step', auto_prune_batch_encoding=True, _op=None)[source]#
Parameters:
add_default_cache_hooks(include_backward=True)[source]#

Add default caching hooks for forward and optionally backward passes.

Parameters:

  • names_filter – Filter to determine which layers to hook

  • cache_dict – Dictionary to store activation values

  • include_backward (bool) – Whether to include backward hooks

Return type:

None

Returns:

Tuple of (forward hooks, backward hooks)

applied_to(module)[source]#

Check if this configuration has been applied to a specific module.

Parameters:

module – The module to check.

Returns:

True if this configuration has been applied to the module, False otherwise.

Return type:

bool

apply(module, cache_dir=None, op_output_dataset_path=None, fallback_sae_targets=None)[source]#

Set up analysis configuration and configure for the given module.

This method handles both setting up the analysis store and configuring the names_filter and hooks for the module. It also injects an analysis_step method if one doesn’t exist in the module.

Parameters:

  • module – The module to configure for.

  • cache_dir (str | None) – Optional cache directory.

  • op_output_dataset_path (str | None) – Optional output path.

  • fallback_sae_targets (LatentAnalysisTargets | None) – Optional fallback LatentAnalysisTargets to use if this config doesn’t have one.

check_add_default_hooks()[source]#

Construct forward and backward hooks based on analysis operation.

Return type:

tuple[list, list] | None

materialize_names_filter(module, fallback_sae_targets=None)[source]#

Set names_filter using latent_analysis_targets if not already set.

For TransformerBridge models the constructed filter is automatically extended with canonical-name equivalents so that both run_with_cache (which iterates canonical hook_dict keys) and downstream filtering (e.g. get_alive_latents_impl) can match activations regardless of whether the cache uses alias or canonical naming.

Parameters:

  • module – The module to construct the names_filter for.

  • fallback_sae_targets (LatentAnalysisTargets | None) – Optional fallback LatentAnalysisTargets to use if this config doesn’t have one.

Return type:

None

maybe_set_hooks()[source]#

Set hooks if they’re not already set.

Return type:

None

prepare_model_ctx(module, fallback_sae_targets=None)[source]#

Configure names_filter and hooks for a specific module.

Parameters:

  • module – The module to configure for.

  • fallback_sae_targets (LatentAnalysisTargets | None) – Optional fallback LatentAnalysisTargets to use if this config doesn’t have one.

Return type:

None

reset_applied_state(module=None)[source]#

Reset the applied state tracking.

Parameters:

module – Optional specific module to reset. If None, reset for all modules.

Return type:

None

save_batch(analysis_batch, batch, tokenizer=None)[source]#

Process and yield analysis batch results.

Uses AnalysisOp.process_batch for consistent processing regardless of whether an operation is defined or using a manual analysis step.

Parameters:

  • analysis_batch (BaseAnalysisBatchProtocol) – The analysis batch to process

  • batch (BatchEncoding) – The raw batch data

  • tokenizer (PreTrainedTokenizerBase | None) – Optional tokenizer for decoding prompts

Yields:

Processed analysis batch

update(**kwargs)[source]#

Update multiple fields of the dataclass at once.

property op: str | AnalysisOp | Callable | list[AnalysisOp | Callable] | None#

Get the operation, unwrapping any OpWrapper if present.

class interpretune.config.AnalysisRunnerCfg(*, it_session=None, module=None, datamodule=None, limit_train_batches=-1, limit_val_batches=-1, limit_test_batches=-1, max_steps=-1, max_epochs=-1, analysis_cfgs=None, limit_analysis_batches=-1, cache_dir=None, op_output_dataset_path=None, latent_analysis_targets=None, artifact_cfg=<factory>, ignore_manual=False)[source]#
Parameters:
class interpretune.config.AutoCompConf(*, auto_comp_cfg=None)[source]#
class interpretune.config.AutoCompConfig(*, module_cfg_name, module_cfg_mixin, target_adapters=None, _orig_cfg_cls=None)[source]#
Parameters:
class interpretune.config.BaseGenerationConfig(*, generate_kwargs=<factory>)[source]#
Parameters:

generate_kwargs (dict)

class interpretune.config.CircuitTracerConfig(*, backend='transformerlens', model_name=None, transcoder_set='gemma', dtype=torch.bfloat16, max_n_logits=10, desired_logit_prob=0.95, batch_size=256, max_feature_nodes=None, offload=None, 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)[source]#

Configuration for Circuit Tracer functionality.

This configuration extends ITLensConfig with circuit tracing specific parameters for generating attribution graphs using ReplacementModel and transcoders.

Parameters:

  • backend (str)

  • model_name (str | None)

  • transcoder_set (str)

  • dtype (dtype)

  • max_n_logits (int)

  • desired_logit_prob (float)

  • batch_size (int)

  • max_feature_nodes (int | None)

  • offload (str | None)

  • lazy_encoder (bool | None)

  • lazy_decoder (bool)

  • verbose (bool)

  • default_node_threshold (float)

  • default_edge_threshold (float)

  • save_graphs (bool)

  • graph_output_dir (str | None)

  • analysis_target_tokens (list[str] | None)

  • target_token_ids (list[int] | Tensor | str | None)

  • use_neuronpedia (bool)

  • intervention_scale_factor (float)

  • intervention_max_influence_norm_scale (bool)

  • intervention_sign_aware_scale (bool)

  • intervention_value (float | None)

  • intervention_value_source (str)

  • intervention_constrained_layers (list[int] | None)

  • intervention_freeze_attention (bool | None)

  • intervention_apply_activation_function (bool | None)

  • intervention_sparse (bool)

  • intervention_return_activations (bool)

  • nnsight_remote (bool)

  • ndif_api_key (str | None)

analysis_target_tokens: list[str] | None = None#

A tensor of pre-tokenized target token IDs for analysis or a module attribute to be used as a source for them.

backend: str = 'transformerlens'#

Model name to use for attribution. If None, uses the base model name.

batch_size: int = 256#

Maximum number of feature nodes to include in attribution.

default_edge_threshold: float = 0.98#

Whether to automatically save generated graphs.

default_node_threshold: float = 0.8#

Default threshold for edge pruning in visualization.

desired_logit_prob: float = 0.95#

Batch size for backward passes during attribution.

dtype: dtype = torch.bfloat16#

Maximum number of logit nodes to attribute from.

graph_output_dir: str | None = None#

Specific tokens to analyze, will use tokens associated with top max_n_logits if None.

intervention_apply_activation_function: bool | None = None#

Whether to request sparse intervention activations from circuit-tracer.

intervention_constrained_layers: list[int] | None = None#

Optional passthrough for circuit-tracer attention freezing during intervention.

intervention_freeze_attention: bool | None = None#

Optional passthrough controlling activation-function application during intervention.

intervention_max_influence_norm_scale: bool = False#

Use top_feature_scores sign to choose intervention direction when scores are available.

intervention_return_activations: bool = False#

Whether to use remote execution for NNsight backend.

Only applicable when backend=’nnsight’. Enables running models on NNsight’s remote infrastructure.

intervention_scale_factor: float = 1.0#

Scale each feature by abs(score) / max(abs(score)) before applying intervention_scale_factor.

intervention_sign_aware_scale: bool = True#

Optional constant value for all interventions.

When None, the op uses per-feature values from top_feature_scores.

intervention_sparse: bool = False#

Whether to request intervention activations alongside logits from circuit-tracer.

intervention_value: float | None = None#

Source for per-feature intervention values when intervention_value is unset.

intervention_value_source: str = 'top_feature_scores'#

Optional explicit constrained layer list passed to circuit-tracer intervention APIs.

lazy_decoder: bool = True#

Whether to display detailed progress information.

lazy_encoder: bool | None = None#

Lazily load transcoder decoder weights from disk on access (default True).

max_feature_nodes: int | None = None#

Memory optimization option (‘cpu’, ‘disk’, or None).

max_n_logits: int = 10#

Cumulative probability threshold for top logits.

model_name: str | None = None#

Transcoder set to use.

Can be ‘gemma’, ‘llama’, or path to custom config.

nnsight_remote: bool = False#

API key for NNsight remote execution.

Only applicable when backend=’nnsight’ and nnsight_remote=True.

offload: str | None = None#

Lazily load transcoder encoder weights from disk on access.

When None (default), auto-enabled when offload='cpu' to avoid OOM during transcoder loading for large transcoder widths (e.g. 262k with 4B+ models).

save_graphs: bool = True#

Directory to save attribution graphs.If None, uses analysis output directory.

target_token_ids: list[int] | Tensor | str | None = None#

Whether to prepare graphs for Neuronpedia graph storage and analysis.

transcoder_set: str = 'gemma'#

Data type for model and transcoders.

use_neuronpedia: bool = False#

Base scale factor applied to each constructed intervention value.

verbose: bool = True#

Default threshold for node pruning in visualization.

class interpretune.config.CoreGenerationConfig(*, generate_kwargs=<factory>, max_new_tokens=5, do_sample=True, top_p=1.0, top_k=50, temperature=1.0, return_dict_in_generate=None, output_logits=None)[source]#
Parameters:

  • generate_kwargs (dict)

  • max_new_tokens (int)

  • do_sample (bool)

  • top_p (float)

  • top_k (int)

  • temperature (float)

  • return_dict_in_generate (bool | None)

  • output_logits (bool | None)

interpretune.config.ExtensionConf#

alias of ExtensionsConf

class interpretune.config.GenerativeClassificationConfig(*, enabled=False, lm_generation_cfg=<factory>, input_inspection_enabled=True)[source]#
Parameters:
class interpretune.config.HFFromPretrainedConfig(*, pretrained_kwargs=<factory>, dynamic_module_cfg=<factory>, use_model_cache=False, model_head='', lora_cfg=<factory>, bitsandbytesconfig=<factory>, activation_checkpointing=False, enable_input_require_grads=True, default_head='transformers.AutoModelForCausalLM')[source]#

HFFromPretrainedConfig: Configuration for loading a pretrained model from Huggingface along with configuration options contingent on the HF pretrained model type.

Parameters:
class interpretune.config.HFGenerationConfig(*, generate_kwargs=<factory>, model_config=<factory>, default_overrides=<factory>)[source]#
Parameters:

  • generate_kwargs (dict)

  • model_config (dict)

  • default_overrides (dict)

class interpretune.config.ITConfig(*, debug_lm_cfg=<factory>, memprofiler_cfg=<factory>, auto_comp_cfg=None, compatibility_attrs=<factory>, experiment_tag='default', log_env_details=True, core_log_dir=None, logging_level='INFO', log_to_stream=True, log_to_file=True, analysis_cfg=None, generative_step_cfg=<factory>, hf_from_pretrained_cfg=None, classification_mapping=None, classification_mapping_indices=None, optimizer_init=<factory>, lr_scheduler_init=<factory>, pl_lrs_cfg=<factory>, model_class=None, model_cfg=<factory>, cust_fwd_kwargs=<factory>, model_name_or_path='', task_name='', tokenizer_name=None, tokenizer=None, os_env_model_auth_key=None, tokenizer_id_overrides=<factory>, tokenizer_kwargs=<factory>, defer_model_init=False)[source]#
Parameters:
class interpretune.config.ITDataModuleConfig(*, remove_unused_columns=True, text_fields=None, dataset_path=None, enable_datasets_cache=False, data_collator_cfg=<factory>, signature_columns=<factory>, prepare_data_map_cfg=<factory>, tokenizers_parallelism=True, local_fast_tokenizer_path=None, cust_tokenization_pattern=None, special_tokens_dict=<factory>, max_seq_length=2048, model_name_or_path='', task_name='', tokenizer_name=None, tokenizer=None, os_env_model_auth_key=None, tokenizer_id_overrides=<factory>, tokenizer_kwargs=<factory>, defer_model_init=False, train_batch_size=32, eval_batch_size=32, dataloader_kwargs=<factory>, prompt_cfg=<factory>)[source]#
Parameters:

  • remove_unused_columns (bool)

  • text_fields (Tuple | None)

  • dataset_path (str | Path | None)

  • enable_datasets_cache (bool | None)

  • data_collator_cfg (dict[str, Any])

  • signature_columns (List | None)

  • prepare_data_map_cfg (dict[str, Any])

  • tokenizers_parallelism (bool)

  • local_fast_tokenizer_path (str | None)

  • cust_tokenization_pattern (str | None)

  • special_tokens_dict (dict[str, Any])

  • max_seq_length (int)

  • model_name_or_path (str)

  • task_name (str)

  • tokenizer_name (str | None)

  • tokenizer (PreTrainedTokenizerBase | None)

  • os_env_model_auth_key (str | None)

  • tokenizer_id_overrides (dict | None)

  • tokenizer_kwargs (dict[str, Any])

  • defer_model_init (bool | None)

  • train_batch_size (int)

  • eval_batch_size (int)

  • dataloader_kwargs (dict[str, Any])

  • prompt_cfg (PromptConfig)

class interpretune.config.ITExtension(ext_attr, ext_cls_fqn, ext_cfg_fqn)[source]#

Create new instance of ITExtension(ext_attr, ext_cls_fqn, ext_cfg_fqn)

Parameters:

  • ext_attr (str)

  • ext_cls_fqn (str)

  • ext_cfg_fqn (str)

ext_attr: str#

Alias for field number 0

ext_cfg_fqn: str#

Alias for field number 2

ext_cls_fqn: str#

Alias for field number 1

class interpretune.config.ITLensBridgeConfig(*, move_to_device=True, default_padding_side='right', use_bridge=True, model_name='gpt2-small', transformer_bridge_config_overrides=None, enable_compatibility_mode=False, enable_compatibility_mode_kwargs=None, device=None, dtype='float32')[source]#

TransformerBridge-specific configuration for TransformerLens v3 integration.

This config provides explicit control over TransformerBridge initialization and compatibility mode settings. Use this config when you need fine-grained control over how the bridge is configured (e.g., enabling/disabling weight processing).

Parameters:

  • model_name (str) – The model name for TransformerBridge (passed to TransformerBridgeConfig).

  • transformer_bridge_config_overrides (dict[str, Any] | None) – Optional dict of kwargs to pass to/override in the TransformerBridgeConfig constructor. Use this to set device, dtype, or any other TransformerBridgeConfig fields.

  • enable_compatibility_mode (bool) – Whether to call enable_compatibility_mode() on the TransformerBridge after instantiation. Default: False.

  • enable_compatibility_mode_kwargs (dict[str, Any] | None) –

    Optional dict of kwargs for enable_compatibility_mode(). Supported kwargs:

    • disable_warnings: bool (default False)

    • no_processing: bool (default False) - disables ALL weight processing

    • fold_ln: bool (default True) - fold LayerNorm weights

    • center_writing_weights: bool (default True)

    • center_unembed: bool (default True)

    • fold_value_biases: bool (default True)

    • refactor_factored_attn_matrices: bool (default False)

  • move_to_device (bool | None)

  • default_padding_side (Literal['left', 'right'] | None)

  • use_bridge (bool | None)

  • device (str | None)

  • dtype (str)

Example:

# Basic bridge config (no processing, current default behavior)
config = ITLensBridgeConfig(model_name="gpt2-small")

# Enable weight processing via compatibility mode
config = ITLensBridgeConfig(
    model_name="gpt2-small",
    enable_compatibility_mode=True,
    enable_compatibility_mode_kwargs={"fold_ln": True, "fold_value_biases": True}
)

# Use no_processing (analogous to ITLensFromPretrainedNoProcessingConfig behavior)
config = ITLensBridgeConfig(
    model_name="gpt2-small",
    enable_compatibility_mode=True,
    enable_compatibility_mode_kwargs={"no_processing": True}
)
class interpretune.config.ITLensConfig(*, debug_lm_cfg=<factory>, memprofiler_cfg=<factory>, auto_comp_cfg=None, compatibility_attrs=<factory>, experiment_tag='default', log_env_details=True, core_log_dir=None, logging_level='INFO', log_to_stream=True, log_to_file=True, analysis_cfg=None, generative_step_cfg=<factory>, hf_from_pretrained_cfg=None, classification_mapping=None, classification_mapping_indices=None, optimizer_init=<factory>, lr_scheduler_init=<factory>, pl_lrs_cfg=<factory>, model_class=None, model_cfg=<factory>, cust_fwd_kwargs=<factory>, model_name_or_path='', task_name='', tokenizer_name=None, tokenizer=None, os_env_model_auth_key=None, tokenizer_id_overrides=<factory>, tokenizer_kwargs=<factory>, defer_model_init=False, tl_cfg)[source]#

Dataclass to encapsulate the ITModule internal state.

Parameters:
class interpretune.config.ITLensCustomConfig(*, move_to_device=True, default_padding_side='right', use_bridge=False, cfg)[source]#

Custom TL config for creating a HookedTransformer from a TL config.

NOTE: TransformerBridge is not supported with config-only initialization. Set use_bridge=False (default) or interpretune will force the value to False and warn.

Parameters:

  • move_to_device (bool | None)

  • default_padding_side (Literal['left', 'right'] | None)

  • use_bridge (bool | None)

  • cfg (Any)

class interpretune.config.ITLensFromPretrainedConfig(*, move_to_device=True, default_padding_side='right', use_bridge=True, model_name='gpt2-small', fold_ln=True, center_writing_weights=True, center_unembed=True, refactor_factored_attn_matrices=False, checkpoint_index=None, checkpoint_value=None, hf_model=None, device=None, n_devices=1, tokenizer=None, fold_value_biases=True, default_prepend_bos=True, dtype='float32')[source]#
Parameters:

  • move_to_device (bool | None)

  • default_padding_side (Literal['left', 'right'] | None)

  • use_bridge (bool | None)

  • model_name (str)

  • fold_ln (bool | None)

  • center_writing_weights (bool | None)

  • center_unembed (bool | None)

  • refactor_factored_attn_matrices (bool | None)

  • checkpoint_index (int | None)

  • checkpoint_value (int | None)

  • hf_model (AutoModelForCausalLM | str | None)

  • device (str | None)

  • n_devices (int | None)

  • tokenizer (PreTrainedTokenizerBase | None)

  • fold_value_biases (bool | None)

  • default_prepend_bos (bool | None)

  • dtype (str)

class interpretune.config.ITLensFromPretrainedNoProcessingConfig(*, move_to_device=True, default_padding_side='right', use_bridge=True, model_name='gpt2-small', fold_ln=False, center_writing_weights=False, center_unembed=False, refactor_factored_attn_matrices=False, checkpoint_index=None, checkpoint_value=None, hf_model=None, device=None, n_devices=1, tokenizer=None, fold_value_biases=False, default_prepend_bos=True, dtype='float32')[source]#
Parameters:

  • move_to_device (bool | None)

  • default_padding_side (Literal['left', 'right'] | None)

  • use_bridge (bool | None)

  • model_name (str)

  • fold_ln (bool | None)

  • center_writing_weights (bool | None)

  • center_unembed (bool | None)

  • refactor_factored_attn_matrices (bool | None)

  • checkpoint_index (int | None)

  • checkpoint_value (int | None)

  • hf_model (AutoModelForCausalLM | str | None)

  • device (str | None)

  • n_devices (int | None)

  • tokenizer (PreTrainedTokenizerBase | None)

  • fold_value_biases (bool | None)

  • default_prepend_bos (bool | None)

  • dtype (str)

class interpretune.config.ITNNsightConfig(*, debug_lm_cfg=<factory>, memprofiler_cfg=<factory>, auto_comp_cfg=None, compatibility_attrs=<factory>, experiment_tag='default', log_env_details=True, core_log_dir=None, logging_level='INFO', log_to_stream=True, log_to_file=True, analysis_cfg=None, generative_step_cfg=<factory>, hf_from_pretrained_cfg=None, classification_mapping=None, classification_mapping_indices=None, optimizer_init=<factory>, lr_scheduler_init=<factory>, pl_lrs_cfg=<factory>, model_class=None, model_cfg=<factory>, cust_fwd_kwargs=<factory>, model_name_or_path='', task_name='', tokenizer_name=None, tokenizer=None, os_env_model_auth_key=None, tokenizer_id_overrides=<factory>, tokenizer_kwargs=<factory>, defer_model_init=False, nnsight_cfg)[source]#

Dataclass to encapsulate ITModule configuration for NNsight models.

This class extends ITConfig to provide NNsight-specific model initialization and configuration handling. Similar to ITLensConfig for TransformerLens, ITNNsightConfig manages the integration between Interpretune’s module system and NNsight’s LanguageModel wrapper.

NNsight wraps HuggingFace models directly without weight conversion, providing tracing-based activation access for analysis and intervention.

Parameters:
nnsight_cfg#

Configuration for NNsight model initialization.

Example

Basic GPT-2 configuration:

>>> config = ITNNsightConfig(
...     model_name_or_path="openai-community/gpt2",
...     nnsight_cfg=NNsightConfig(model_name="openai-community/gpt2"),
... )

Using registry-based configuration:

from interpretune import ITSession
from interpretune.protocol import Adapter
session = ITSession(model_src_key="gpt2", adapter_ctx=(Adapter.core, Adapter.nnsight))
class interpretune.config.ITSerializableCfg[source]#

Base class for serializable Interpretune configs.

Automatically registers subclasses and Path types as safe globals for PyTorch checkpoint loading.

class interpretune.config.ITSharedConfig(*, model_name_or_path='', task_name='', tokenizer_name=None, tokenizer=None, os_env_model_auth_key=None, tokenizer_id_overrides=<factory>, tokenizer_kwargs=<factory>, defer_model_init=False)[source]#
Parameters:

  • model_name_or_path (str)

  • task_name (str)

  • tokenizer_name (str | None)

  • tokenizer (PreTrainedTokenizerBase | None)

  • os_env_model_auth_key (str | None)

  • tokenizer_id_overrides (dict | None)

  • tokenizer_kwargs (dict[str, Any])

  • defer_model_init (bool | None)

class interpretune.config.ITState(_it_lr_scheduler_configs=<factory>, _it_optimizers=<factory>, _log_dir=None, _datamodule=None, _device=None, _extensions=<factory>, _session_complete=False, _init_hparams=<factory>, _current_epoch=0, _global_step=0)[source]#

Dataclass to encapsulate the ITModule internal state and keep top-level namespace as clean as possible.

Parameters:
to_dict()[source]#

Return a JSON-serializable summary dict of the ITState.

Return type:

dict

interpretune.config.NNsightCfg#

alias of NNsightConfig

class interpretune.config.NNsightConfig(*, model_name='openai-community/gpt2', device_map=None, torch_dtype='float32', dispatch=True, tokenizer_kwargs=None, trust_remote_code=False, attn_implementation=None, default_padding_side='left', remote=False, api_key=None)[source]#

Configuration for NNsight model integration.

NNsight wraps HuggingFace models directly without weight conversion, providing tracing-based activation access and intervention capabilities.

This config provides IT-specific settings that control how NNsight’s LanguageModel is initialized and configured. Most kwargs are passed through to the underlying HuggingFace from_pretrained call.

Parameters:

  • model_name (str | None)

  • device_map (str | dict[str, Any] | None)

  • torch_dtype (str | dtype | None)

  • dispatch (bool)

  • tokenizer_kwargs (dict[str, Any] | None)

  • trust_remote_code (bool)

  • attn_implementation (str | None)

  • default_padding_side (Literal['left', 'right'] | None)

  • remote (bool)

  • api_key (str | None)

model_name#

The model name or path for loading (e.g., “openai-community/gpt2”). If None, uses the base model name from ITConfig.model_name_or_path.

device_map#

Device mapping strategy. Options: “auto”, “cpu”, “cuda:0”, or dict mapping.

torch_dtype#

Data type for model weights (e.g., “float32”, “float16”, “bfloat16”).

dispatch#

Whether to load model immediately (True) or defer loading (False). Default True for typical usage.

tokenizer_kwargs#

Additional kwargs passed to tokenizer initialization.

trust_remote_code#

Whether to trust remote code for custom model architectures.

attn_implementation#

Attention implementation to use (e.g., “flash_attention_2”).

Example

Basic GPT-2 configuration:

>>> config = NNsightConfig(model_name="openai-community/gpt2")

Larger model with precision settings:

>>> config = NNsightConfig(
...     model_name="meta-llama/Llama-3.1-8B",
...     device_map="auto",
...     torch_dtype="bfloat16",
...     trust_remote_code=True,
... )
get_nnsight_kwargs()[source]#

Generate kwargs dict for NNsight LanguageModel initialization.

Return type:

dict[str, Any]

Returns:

Dictionary of kwargs to pass to LanguageModel constructor.

api_key: str | None = None#

API key for NDIF remote execution.

If not provided, will be read from NDIF_API_KEY environment variable. Only used when remote=True.

attn_implementation: str | None = None#

Attention implementation to use.

Options include:

  • None: Use default implementation

  • “flash_attention_2”: Use Flash Attention 2 (requires compatible hardware)

  • “sdpa”: Use PyTorch’s scaled dot product attention

default_padding_side: Optional[Literal['left', 'right']] = 'left'#

Default padding side for tokenization.

NNsight defaults to left padding.

device_map: str | dict[str, Any] | None = None#

Device mapping strategy for HuggingFace Accelerate.

Common options:

  • “auto”: Automatically distribute across available devices

  • “cpu”: Force CPU execution

  • “cuda:0”: Specific CUDA device

  • dict: Custom layer-to-device mapping

dispatch: bool = True#

Whether to load model immediately (True) or defer (False).

When True (default), the model is loaded into memory immediately upon initialization. When False, model loading is deferred until first use (useful for lazy initialization).

model_name: str | None = 'openai-community/gpt2'#

Model name or path for loading.

If None, uses ITConfig.model_name_or_path.

remote: bool = False#

Whether to use NNsight’s remote execution via NDIF.

When True, model execution happens on NDIF servers rather than locally. Requires NDIF_API_KEY environment variable to be set.

property resolved_dtype: dtype | str | None#

Get the resolved torch dtype.

tokenizer_kwargs: dict[str, Any] | None = None#

Additional kwargs passed to AutoTokenizer.from_pretrained().

Example

>>> config = NNsightConfig(
...     model_name="gpt2",
...     tokenizer_kwargs={"padding_side": "left", "add_bos_token": True}
... )
torch_dtype: str | dtype | None = 'float32'#

Data type for model weights.

Accepts string (“float32”, “bfloat16”) or torch.dtype.

trust_remote_code: bool = False#

Whether to trust remote code for custom model architectures.

Required for some models that use custom code not in the transformers library.

class interpretune.config.PromptConfig(*, cust_task_prompt=<factory>)[source]#
Parameters:

cust_task_prompt (dict[str, Any])

class interpretune.config.SAELensConfig(*, debug_lm_cfg=<factory>, memprofiler_cfg=<factory>, auto_comp_cfg=None, compatibility_attrs=<factory>, experiment_tag='default', log_env_details=True, core_log_dir=None, logging_level='INFO', log_to_stream=True, log_to_file=True, analysis_cfg=None, generative_step_cfg=<factory>, hf_from_pretrained_cfg=None, classification_mapping=None, classification_mapping_indices=None, optimizer_init=<factory>, lr_scheduler_init=<factory>, pl_lrs_cfg=<factory>, model_class=None, model_cfg=<factory>, cust_fwd_kwargs=<factory>, model_name_or_path='', task_name='', tokenizer_name=None, tokenizer=None, os_env_model_auth_key=None, tokenizer_id_overrides=<factory>, tokenizer_kwargs=<factory>, defer_model_init=False, backend='transformerlens', use_bridge=True, tl_cfg=None, nnsight_cfg=None, sae_cfgs, add_saes_on_init=False)[source]#

Configuration for SAE Lens adapter.

Supports both TransformerLens and NNsight backends via the backend field. When backend="transformerlens" (default), tl_cfg must be provided. When backend="nnsight", nnsight_cfg must be provided instead.

Inherits from ITConfig (not ITLensConfig) so that it is backend-agnostic at the type level. TL-specific initialization logic is provided by TLConfigInitMixin, which is shared with ITLensConfig.

The use_bridge field is only meaningful when backend="transformerlens" and controls whether a SAETransformerBridge (True) or HookedSAETransformer (False) is used.

Parameters:
class interpretune.config.SAELensCustomConfig(*, cfg)[source]#
Parameters:

cfg (Any)

class interpretune.config.SAELensFromPretrainedConfig(*, release, sae_id, device=None, dtype=None)[source]#
Parameters:

  • release (str)

  • sae_id (str)

  • device (str | None)

  • dtype (str | None)

class interpretune.config.SessionRunnerCfg(*, it_session=None, module=None, datamodule=None, limit_train_batches=-1, limit_val_batches=-1, limit_test_batches=-1, max_steps=-1, max_epochs=-1)[source]#
Parameters:

  • it_session (ITSession | None)

  • module (BuiltTrainSteppable | BuiltValidationSteppable | BuiltTestSteppable | BuiltPredictSteppable | None)

  • datamodule (BuiltTrainLoadable | BuiltValLoadable | BuiltTestLoadable | BuiltPredictLoadable | None)

  • limit_train_batches (int)

  • limit_val_batches (int)

  • limit_test_batches (int)

  • max_steps (int)

  • max_epochs (int)

class interpretune.config.TLConfigInitMixin[source]#

Mixin providing TransformerLens config initialization logic.

Shared by ITLensConfig (pure TL config) and SAELensConfig (multi-backend config that delegates to TL when backend="transformerlens"). Separating these helpers avoids code duplication without requiring SAELensConfig to inherit from ITLensConfig.

Note: This mixin always co-inherits with ITConfig, which provides hf_from_pretrained_cfg, model_name_or_path, tokenizer_kwargs, etc. Only attributes unique to TL initialization are declared here.

class interpretune.config.TLensGenerationConfig(*, generate_kwargs=<factory>, max_new_tokens=5, do_sample=True, top_p=1.0, top_k=50, temperature=1.0, return_dict_in_generate=None, output_logits=None, stop_at_eos=True, eos_token_id=None, freq_penalty=0.0, use_past_kv_cache=True, prepend_bos=None, padding_side=None, return_type='input', verbose=True)[source]#
Parameters:

  • generate_kwargs (dict)

  • max_new_tokens (int)

  • do_sample (bool)

  • top_p (float)

  • top_k (int)

  • temperature (float)

  • return_dict_in_generate (bool | None)

  • output_logits (bool | None)

  • stop_at_eos (bool)

  • eos_token_id (int | None)

  • freq_penalty (float)

  • use_past_kv_cache (bool)

  • prepend_bos (bool | None)

  • padding_side (Literal['left', 'right'] | None)

  • return_type (str | None)

  • verbose (bool)

interpretune.config.init_analysis_cfgs(module, analysis_cfgs, cache_dir=None, op_output_dataset_path=None, latent_analysis_targets=None, ignore_manual=False)[source]#

Initialize analysis configurations for the given module.

Parameters:
Return type:

None

interpretune.config.init_analysis_dirs(module, cache_dir=None, op_output_dataset_path=None, analysis_cfgs=None)[source]#

Initialize the analysis directories for the given module and analysis configurations.

Parameters:

  • module (LatentAnalysisModuleProtocol) – The module to set up analysis directories for

  • cache_dir (str | Path | None) – Optional path to cache directory, will be created if it doesn’t exist

  • op_output_dataset_path (str | Path | None) – Optional path for analysis outputs, will be created if it doesn’t exist

  • analysis_cfgs (list[AnalysisCfg] | None) – Optional list of analysis configurations to check for op directories

Return type:

tuple[Path, Path]

Returns:

Tuple of (cache_dir, op_output_dataset_path) as Path objects

On this page
Show Source