wip

Author: mivanit

spd/clustering/pipeline/clustering_pipeline.py

@@ -5,26 +5,33 @@
 Each batch loads its own model and WandB run to match original design.
 """
 
+from collections.abc import Iterator
+from typing import Any
+
 from spd.clustering.merge_run_config import RunConfig
 
 
 def main(config: RunConfig) -> None:
-    from spd.clustering.consts import DistancesArray, MergesArray
+    from spd.clustering.consts import DistancesArray, DistancesMethod, MergesArray
     from spd.clustering.math.merge_distances import (
         compute_distances,
     )
-    from spd.clustering.pipeline.s1_split_dataset import split_dataset
+    from spd.clustering.pipeline.s1_split_dataset import BatchTensor, split_dataset
     from spd.clustering.pipeline.s2_clustering import ClusteringResult, process_batches_parallel
     from spd.clustering.pipeline.s3_normalize_histories import normalize_and_save
     from spd.clustering.pipeline.s4_compute_distances import create_clustering_report
     from spd.clustering.pipeline.storage import ClusteringStorage
 
-    storage = ClusteringStorage(base_path=config.base_path, run_identifier=config.config_identifier)
+    storage: ClusteringStorage = ClusteringStorage(
+        base_path=config.base_path, run_identifier=config.config_identifier
+    )
 
     print(f"Run record saved to {storage.run_config_file}")
     storage.save_run_config(config)
 
     print(f"Splitting dataset into {config.n_batches} batches...")
+    batches: Iterator[BatchTensor]
+    dataset_config: dict[str, Any]
     batches, dataset_config = split_dataset(config=config)
     storage.save_batches(batches=batches, config=dataset_config)
 
@@ -40,7 +47,8 @@ def main(config: RunConfig) -> None:
 
     normalized_merge_array: MergesArray = normalize_and_save(storage=storage)
 
-    method = "perm_invariant_hamming"
+    # TODO: read method from config
+    method: DistancesMethod = "perm_invariant_hamming"
     distances: DistancesArray = compute_distances(
         normalized_merge_array=normalized_merge_array,
         method=method,

tests/clustering/scripts/cluster_resid_mlp.py

@@ -70,10 +70,14 @@
 # %%
 # Get component activations
 # ============================================================
+# Get a single batch from the dataloader
+BATCH_DATA: tuple[Tensor, Tensor] = next(iter(DATALOADER))
+BATCH: Tensor = BATCH_DATA[0]
+
 COMPONENT_ACTS: dict[str, Tensor] = component_activations(
     model=MODEL,
     device=DEVICE,
-    dataloader=DATALOADER,
+    batch=BATCH,
     sigmoid_type="hard",
 )
 
@@ -88,13 +92,13 @@
 PROCESSED_ACTIVATIONS: ProcessedActivations = process_activations(
     COMPONENT_ACTS,
     filter_dead_threshold=FILTER_DEAD_THRESHOLD,
-    sort_components=False,  # Test the new sorting functionality
 )
 
 
 plot_activations(
     processed_activations=PROCESSED_ACTIVATIONS,
-    save_pdf=False,
+    save_dir=None,
+    wandb_run=None,
 )
 
 # %%
@@ -113,27 +117,24 @@
 
 
 def _plot_func(
-    costs: torch.Tensor,
-    # merge_history: MergeHistory,
-    current_merge: Any,
     current_coact: torch.Tensor,
-    # current_act_mask: torch.Tensor,
-    i: int,
-    # k_groups: int,
-    # activation_mask_orig: torch.Tensor,
     component_labels: list[str],
-    # sweep_params: dict[str, Any],
-    **kwargs: Any,
+    current_merge: Any,
+    costs: torch.Tensor,
+    merge_history: MergeHistory,
+    iter_idx: int,
+    k_groups: int,
+    merge_pair_cost: float,
+    mdl_loss: float,
+    mdl_loss_norm: float,
+    diag_acts: torch.Tensor,
 ) -> None:
-    assert kwargs
-    if (i % 50 == 0 and i > 0) or i == 1:
-        # latest = merge_history.latest()
-        # latest['merges'].plot()
+    if (iter_idx % 50 == 0 and iter_idx > 0) or iter_idx == 1:
         plot_merge_iteration(
             current_merge=current_merge,
             current_coact=current_coact,
             costs=costs,
-            iteration=i,
+            iteration=iter_idx,
             component_labels=component_labels,
             show=True,  # Show the plot interactively
         )
@@ -144,7 +145,7 @@ def _plot_func(
     batch_id="batch_0",
     activations=PROCESSED_ACTIVATIONS.activations,
     component_labels=PROCESSED_ACTIVATIONS.labels,
-    log_callback=None,
+    log_callback=_plot_func,
 )
 
 # %%
@@ -162,18 +163,18 @@ def _plot_func(
 
 # Modern approach: run merge_iteration multiple times to create ensemble
 ENSEMBLE_SIZE: int = 4
-histories: list[MergeHistory] = []
+HISTORIES: list[MergeHistory] = []
 for i in range(ENSEMBLE_SIZE):
-    history: MergeHistory = merge_iteration(
+    HISTORY: MergeHistory = merge_iteration(
         merge_config=MERGE_CFG,
         batch_id=f"batch_{i}",
         activations=PROCESSED_ACTIVATIONS.activations,
         component_labels=PROCESSED_ACTIVATIONS.labels,
         log_callback=None,
     )
-    histories.append(history)
+    HISTORIES.append(HISTORY)
 
-ENSEMBLE: MergeHistoryEnsemble = MergeHistoryEnsemble(data=histories)
+ENSEMBLE: MergeHistoryEnsemble = MergeHistoryEnsemble(data=HISTORIES)
 
 DISTANCES = ENSEMBLE.get_distances(method="perm_invariant_hamming")
 

tests/clustering/scripts/cluster_ss.py

@@ -1,7 +1,6 @@
 # %%
 
 import matplotlib.pyplot as plt
-import numpy as np
 import torch
 from jaxtyping import Int
 from muutils.dbg import dbg_auto
@@ -15,7 +14,8 @@
 from spd.clustering.merge import merge_iteration
 from spd.clustering.merge_config import MergeConfig
 from spd.clustering.merge_history import MergeHistory, MergeHistoryEnsemble
-from spd.clustering.pipeline.s1_split_dataset import split_dataset_lm
+from spd.clustering.merge_run_config import RunConfig
+from spd.clustering.pipeline.s1_split_dataset import split_dataset
 from spd.clustering.plotting.activations import plot_activations
 from spd.clustering.plotting.merge import plot_dists_distribution
 from spd.models.component_model import ComponentModel, SPDRunInfo
@@ -31,23 +31,25 @@
 # ============================================================
 MODEL_PATH: str = "wandb:goodfire/spd/runs/ioprgffh"
 
-_, DATA_CFG = split_dataset_lm(
-    model_path=MODEL_PATH,
-    n_batches=1,
-    batch_size=2,
-)
-DATASET_PATH: str = DATA_CFG["output_files"][0]
-
 SPD_RUN: SPDRunInfo = SPDRunInfo.from_path(MODEL_PATH)
 MODEL: ComponentModel = ComponentModel.from_pretrained(SPD_RUN.checkpoint_path)
 MODEL.to(DEVICE)
 SPD_CONFIG = SPD_RUN.config
 
+# Use split_dataset with RunConfig to get real data
+CONFIG: RunConfig = RunConfig(
+    merge_config=MergeConfig(),
+    model_path=MODEL_PATH,
+    task_name="lm",
+    n_batches=1,
+    batch_size=2,
+)
+BATCHES, _ = split_dataset(config=CONFIG)
 
 # %%
 # Load data batch
 # ============================================================
-DATA_BATCH: Int[Tensor, "batch_size n_ctx"] = torch.tensor(np.load(DATASET_PATH)["input_ids"])
+DATA_BATCH: Int[Tensor, "batch_size n_ctx"] = next(BATCHES)
 
 # %%
 # Get component activations
@@ -75,7 +77,8 @@
 
 plot_activations(
     processed_activations=PROCESSED_ACTIVATIONS,
-    save_pdf=False,
+    save_dir=None,
+    wandb_run=None,
 )
 
 # %%
@@ -94,18 +97,18 @@
 
 # Modern approach: run merge_iteration multiple times to create ensemble
 ENSEMBLE_SIZE: int = 2
-histories: list[MergeHistory] = []
+HISTORIES: list[MergeHistory] = []
 for i in range(ENSEMBLE_SIZE):
-    history: MergeHistory = merge_iteration(
+    HISTORY: MergeHistory = merge_iteration(
         merge_config=MERGE_CFG,
         batch_id=f"batch_{i}",
         activations=PROCESSED_ACTIVATIONS.activations,
         component_labels=PROCESSED_ACTIVATIONS.labels,
         log_callback=None,
     )
-    histories.append(history)
+    HISTORIES.append(HISTORY)
 
-ENSEMBLE: MergeHistoryEnsemble = MergeHistoryEnsemble(data=histories)
+ENSEMBLE: MergeHistoryEnsemble = MergeHistoryEnsemble(data=HISTORIES)
 
 
 # %%