circuit-tracer integration for CrossLayerTranscoders

.github/workflows/pr-pytest.yml

@@ -18,10 +18,10 @@ jobs:
           python-version: "3.11"
 
       - name: Install uv
-        run: pip install uv
+        uses: astral-sh/setup-uv@v7
 
-      - name: Install dependencies
-        run: uv install
+      - name: Install the project
+        run: uv sync --locked --all-extras --dev
 
-      - name: Run pytest
-        run: uv run pytest
+      - name: Run tests
+        run: uv run pytest --ignore=tests/test_text_dataset.py

config/circuit-tracer.yaml

@@ -0,0 +1,172 @@
+# Default configuration for CrossLayer Transcoder training
+# This file uses Lightning CLI's automatic class construction
+
+seed_everything: 42
+
+trainer:
+  # max_steps is number of gradient updates. If using gradient accumulation, this is not the number of batches.
+  max_steps: 100_000 
+  val_check_interval: 1_000
+  limit_val_batches: 1
+  enable_checkpointing: false  # We use custom end-of-training checkpoint
+  num_sanity_val_steps: 0  # Can't run replacement model before standardizers are initialized
+  precision: "16-mixed"
+  accelerator: "gpu"
+  devices: [0]  # [0] means cuda:0
+  accumulate_grad_batches: 1
+  logger:  # WandB logger is recommended but other loggers are supported as well
+    class_path: lightning.pytorch.loggers.WandbLogger
+    init_args:
+      project: "clt"
+      name: "circuit-tracer"
+      save_dir: "./wandb"
+  callbacks:
+    - class_path: crosslayer_transcoder.utils.callbacks.EndOfTrainingCheckpointCallback
+      init_args:
+        checkpoint_dir: "checkpoints"
+    - class_path: crosslayer_transcoder.utils.callbacks.ModelConversionCallback
+      init_args:
+        converter:
+          class_path: crosslayer_transcoder.utils.model_converters.circuit_tracer.CircuitTracerConverter
+          init_args:
+            save_dir: "circuit-tracer"
+            feature_input_hook: "hook_resid_mid"
+            feature_output_hook: "hook_mlp_out"
+        on_events: ["on_train_batch_end"]
+
+model:
+  class_path: crosslayer_transcoder.model.clt_lightning.JumpReLUCrossLayerTranscoderModule
+  init_args:
+    model:
+      class_path: crosslayer_transcoder.model.clt.CrossLayerTranscoder
+      init_args:
+        encoder:
+          class_path: crosslayer_transcoder.model.clt.Encoder
+          init_args:
+            d_acts: 768
+            d_features: 10_000
+            n_layers: 12
+
+        decoder:
+          class_path: crosslayer_transcoder.model.clt.CrosslayerDecoder
+          init_args:
+            d_acts: 768
+            d_features: 10_000
+            n_layers: 12
+
+        nonlinearity:
+          class_path: crosslayer_transcoder.model.jumprelu.JumpReLU
+          init_args:
+            theta: 0.03
+            bandwidth: 0.01
+            n_layers: 12
+            d_features: 10_000
+
+        input_standardizer:
+          class_path: crosslayer_transcoder.model.standardize.DimensionwiseInputStandardizer
+          init_args:
+            n_layers: 12
+            activation_dim: 768
+
+        output_standardizer:
+          class_path: crosslayer_transcoder.model.standardize.DimensionwiseOutputStandardizer
+          init_args:
+            n_layers: 12
+            activation_dim: 768
+
+    # Pre-constructed replacement model
+    replacement_model:
+      class_path: crosslayer_transcoder.metrics.replacement_model_accuracy.ReplacementModelAccuracy
+      init_args:
+        model_name: "openai-community/gpt2"
+        device_map: "cuda:0"  # should match trainer.devices
+        loader_batch_size: 2
+
+    # Pre-constructed dead features metric
+    dead_features:
+      class_path: crosslayer_transcoder.metrics.dead_features.DeadFeatures
+      init_args:
+        n_features: 10_000
+        n_layers: 12
+        return_per_layer: true
+        return_log_freqs: true
+        return_neuron_indices: true
+
+
+    # Training parameters
+    learning_rate: 3e-4
+    compile: true  # if using torch.compile
+    lr_decay_step: 16_000  # lr is scaled by lr_decay_factor after this many steps
+    lr_decay_factor: 0.1
+
+    lambda_sparsity: 0.0007  # sparsity loss weight
+    c_sparsity: 1  # sparsity loss coefficient
+    use_tanh: true  # use tanh nonlinearity in the JumpReLU
+    pre_actv_loss: 1e-6  # pre-activation loss weight
+
+    # Dead features computation settings
+    compute_dead_features: true
+    compute_dead_features_every: 500
+
+data:
+  class_path: crosslayer_transcoder.data.datamodule.ActivationDataModule
+  init_args:
+    # Buffer settings
+    buffer_size: 500_000 # number of activations to store in the buffer
+    n_in_out: 2  # number of input and output layers
+    n_layers: 12  # number of layers in the model
+    activation_dim: 768  # dimension of the activations
+    dtype: "float16"  # dtype of the activations
+    max_batch_size: 8000  # maximum batch size for the data loader
+
+    # Model settings for activation generation
+    model_name: "openai-community/gpt2"
+    model_dtype: "float32"
+
+    # Dataset settings
+    dataset_name: "Skylion007/openwebtext"
+    dataset_split: "train"
+    max_sequence_length: 1024
+
+    # Generation settings
+    generation_batch_size: 10
+    refresh_interval: 0.1  # time (s) between shell logs updates
+
+    # Memory settings
+    shared_memory_name: "activation_buffer"
+    timeout_seconds: 30
+
+    # File paths
+    init_file: null  # path to file with shuffled activations to initialize the buffer fast
+    # if null, activations are generated and training starts when the buffer is at least minimum_fill_threshold full
+
+    # DataLoader settings
+    batch_size: 1000
+    num_workers: 10
+    prefetch_factor: 2
+    shuffle: true
+    persistent_workers: true
+    pin_memory: true
+
+    minimum_fill_threshold: 0.2  # Only provide activations when buffer is at least 20% full
+    # to maintain sufficient shuffling
+
+    use_shared_memory: true
+
+    # Device configuration
+    device_map: "cuda:0"  # "cpu", "auto", "cuda:0", "cuda:0,1,2,3"
+    deployment_policy: "gpu_only"  # "cpu_only", "gpu_only", or "dynamic"
+    # dynamic will use CPU and only GPU if the buffer is almost empty to refill fast. Use this if you use a single GPU and have a beefy CPU.
+
+    # WandB logging configuration for data generation
+    wandb_logging:
+      enabled: true                           # Enable WandB logging for data generation
+      project: "clt"        # WandB project (should match trainer logger)
+      group: null                             # Group name (null = auto-generated from training run)
+      run_name: "data-generator-jumprelu"              # Run name suffix
+      tags: ["data-generation"]               # Tags for the data generation run
+      save_dir: "./wandb"                     # Directory for WandB files
+      log_interval: 5.0                       # Logging interval in seconds
+      offline: true                           # Offline mode for WandB logging
+
+ckpt_path: null

crosslayer_transcoder/model/clt.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Tuple, Union
 
+import einops
 import torch
 import torch.nn as nn
 import yaml
@@ -194,6 +195,17 @@ def to_config(self) -> Dict[str, Any]:
             },
         }
 
+    def to_circuit_tracer(self):
+        W = einops.rearrange(
+            self.get_parameter("W"),
+            "n_layers d_acts d_features -> n_layers d_features d_acts",
+        ).contiguous()
+        b = self.get_parameter("b")
+        return {
+            "W": W,
+            "b": b,
+        }
+
 
 class Decoder(SerializableModule):
     def __init__(self, d_acts: int, d_features: int, n_layers: int):
@@ -202,16 +214,16 @@ def __init__(self, d_acts: int, d_features: int, n_layers: int):
         self.d_features = d_features
         self.n_layers = n_layers
         self.register_parameter(
-            f"W", nn.Parameter(torch.empty((n_layers, d_features, d_acts)))
+            "W", nn.Parameter(torch.empty((n_layers, d_features, d_acts)))
         )
-        self.register_parameter(f"b", nn.Parameter(torch.empty((n_layers, d_acts))))
+        self.register_parameter("b", nn.Parameter(torch.empty((n_layers, d_acts))))
         self._is_folded = False
         self.reset_parameters()
 
     def reset_parameters(self):
         dec_uniform_thresh = 1 / ((self.d_acts * self.n_layers) ** 0.5)
-        self.get_parameter(f"W").data.uniform_(-dec_uniform_thresh, dec_uniform_thresh)
-        self.get_parameter(f"b").data.zero_()
+        self.get_parameter("W").data.uniform_(-dec_uniform_thresh, dec_uniform_thresh)
+        self.get_parameter("b").data.zero_()
 
     @torch.no_grad()
     def forward_layer(
@@ -224,7 +236,7 @@ def forward_layer(
         return (
             einsum(
                 features,
-                self.get_parameter(f"W")[layer],
+                self.get_parameter("W")[layer],
                 "batch_size seq d_features, d_features d_acts -> batch_size seq d_acts",
             )
             + self.b[layer]
@@ -275,6 +287,12 @@ def to_config(self) -> Dict[str, Any]:
             },
         }
 
+    def to_circuit_tracer(self):
+        return {
+            "W": self.W,
+            "b": self.b,
+        }
+
 
 class CrosslayerDecoder(SerializableModule):
     def __init__(self, d_acts: int, d_features: int, n_layers: int):
@@ -287,7 +305,7 @@ def __init__(self, d_acts: int, d_features: int, n_layers: int):
                 f"W_{i}", nn.Parameter(torch.empty((i + 1, d_features, d_acts)))
             )
         self._is_folded = False
-        self.register_parameter(f"b", nn.Parameter(torch.empty((n_layers, d_acts))))
+        self.register_parameter("b", nn.Parameter(torch.empty((n_layers, d_acts))))
         self.reset_parameters()
 
     def reset_parameters(self):
@@ -331,15 +349,15 @@ def forward(
             device=features.device,
             dtype=features.dtype,
         )
-        for l in range(self.n_layers):
-            W = self.get_parameter(f"W_{l}")
-            selected_features = features[:, : l + 1]
+        for layer_idx in range(self.n_layers):
+            W = self.get_parameter(f"W_{layer_idx}")
+            selected_features = features[:, : layer_idx + 1]
             l_recons = einsum(
                 selected_features,
                 W,
                 "batch_size n_layers d_features, n_layers d_features d_acts -> batch_size d_acts",
             )
-            recons[:, l, :] = l_recons
+            recons[:, layer_idx, :] = l_recons
         recons = recons + self.b.to(features.dtype)
         return recons
 
@@ -369,6 +387,32 @@ def to_config(self) -> Dict[str, Any]:
             },
         }
 
+    def to_circuit_tracer(self):
+        output_decs = []
+        for source_layer in range(self.n_layers):
+            output_dec_i = torch.zeros(
+                [self.d_features, self.n_layers - source_layer, self.d_acts],
+            )
+
+            for k in range(source_layer, self.n_layers):
+                # get decoder mat for layer i --> k
+                decoder_w_k = self.get_parameter(f"W_{k}")
+
+                dec_i_k = decoder_w_k[source_layer, ...]
+                assert dec_i_k.shape == (
+                    self.d_features,
+                    self.d_acts,
+                )
+
+                output_dec_i[:, k - source_layer, ...] = dec_i_k
+
+            output_decs.append(output_dec_i)
+
+        return {
+            "W": output_decs,
+            "b": self.b,
+        }
+
 
 class CrossLayerTranscoder(SerializableModule):
     def __init__(
@@ -465,3 +509,20 @@ def save_pretrained(self, directory: Path, fold_standardizers: bool = True):
             yaml.dump({"model": config}, f)
 
         save_file(self.state_dict(), directory / "checkpoint.safetensors")
+
+    def to_circuit_tracer(self):
+        # NOTE: this mutates the model in-place. Potentially bad, but a tradeoff for copying a huge model.
+        self.fold()
+
+        encoder = self.encoder.to_circuit_tracer()
+        decoder = self.decoder.to_circuit_tracer()
+
+        is_per_layer_decoder = isinstance(self.decoder, Decoder)
+
+        config = {
+            "is_per_layer_decoder": is_per_layer_decoder,
+            "encoder": encoder,
+            "decoder": decoder,
+        }
+
+        return config

crosslayer_transcoder/model/jumprelu.py

@@ -44,7 +44,9 @@ def backward(ctx, grad_output):
         grad_input = grad_output.clone()
         grad_input[input < 0] = 0
 
-        theta_grad = -(theta / bandwidth) * rectangle((input - theta) / bandwidth) * grad_output
+        theta_grad = (
+            -(theta / bandwidth) * rectangle((input - theta) / bandwidth) * grad_output
+        )
         return grad_input, theta_grad, None
 
 
@@ -78,7 +80,9 @@ class HeavysideStep(torch.autograd.Function):
     def forward(ctx, input, theta, bandwidth):
         ctx.save_for_backward(input, theta)
         ctx.bandwidth = bandwidth
-        return torch.where(input - theta > 0, torch.ones_like(input), torch.zeros_like(input))
+        return torch.where(
+            input - theta > 0, torch.ones_like(input), torch.zeros_like(input)
+        )
 
     @staticmethod
     def backward(ctx, grad_output):
@@ -87,5 +91,7 @@ def backward(ctx, grad_output):
         grad_input = grad_output.clone()
         grad_input = grad_output * 0.0
 
-        theta_grad = -(1.0 / bandwidth) * rectangle((input - theta) / bandwidth) * grad_output
+        theta_grad = (
+            -(1.0 / bandwidth) * rectangle((input - theta) / bandwidth) * grad_output
+        )
         return grad_input, theta_grad, None