Merge pull request #91 from OpenMOSS/implement_analyze_dispatch

implement dynamic dispatch for sae instantiation in analyze_sae(runner)

Author: Hzfinfdu

src/lm_saes/runner.py

@@ -26,8 +26,10 @@
     TrainerConfig,
     WandbConfig,
 )
+from lm_saes.crosscoder import CrossCoder
 from lm_saes.database import MongoClient
 from lm_saes.initializer import Initializer
+from lm_saes.mixcoder import MixCoder
 from lm_saes.resource_loaders import load_dataset, load_model
 from lm_saes.sae import SparseAutoEncoder
 from lm_saes.trainer import Trainer
@@ -406,7 +408,15 @@ def analyze_sae(settings: AnalyzeSAESettings) -> None:
     mongo_client = MongoClient(settings.mongo)
     activation_factory = ActivationFactory(settings.activation_factory)
 
-    sae = SparseAutoEncoder.from_config(settings.sae)
+    if settings.sae.sae_type == "sae":
+        sae = SparseAutoEncoder.from_config(settings.sae)
+    elif settings.sae.sae_type == "crosscoder":
+        sae = CrossCoder.from_config(settings.sae)
+    elif settings.sae.sae_type == "mixcoder":
+        sae = MixCoder.from_config(settings.sae)
+    else:
+        # TODO: add support for different SAE config types, e.g. MixCoderConfig, CrossCoderConfig, etc.
+        raise ValueError(f"SAE type {settings.sae.sae_type} not supported.")
 
     analyzer = FeatureAnalyzer(settings.analyzer)
 

src/lm_saes/trainer.py

@@ -137,14 +137,16 @@ def _log(self, sae: SparseAutoEncoder, log_info: dict, batch: dict[str, Tensor])
                 "sparsity/below_1e-5": (feature_sparsity < 1e-5).sum().item(),
                 "sparsity/below_1e-6": (feature_sparsity < 1e-6).sum().item(),
             }
-            if sae.cfg.sae_type == 'crosscoder':
-                overall_act_freq_scores = all_reduce_tensor(feature_sparsity, aggregate='max')
-                wandb_log_dict.update({
-                    "sparsity/overall_above_1e-1": (overall_act_freq_scores > 1e-1).sum().item(),
-                    "sparsity/overall_above_1e-2": (overall_act_freq_scores > 1e-2).sum().item(),
-                    "sparsity/overall_below_1e-5": (overall_act_freq_scores < 1e-5).sum().item(),
-                    "sparsity/overall_below_1e-6": (overall_act_freq_scores < 1e-6).sum().item(),
-                })
+            if sae.cfg.sae_type == "crosscoder":
+                overall_act_freq_scores = all_reduce_tensor(feature_sparsity, aggregate="max")
+                wandb_log_dict.update(
+                    {
+                        "sparsity/overall_above_1e-1": (overall_act_freq_scores > 1e-1).sum().item(),
+                        "sparsity/overall_above_1e-2": (overall_act_freq_scores > 1e-2).sum().item(),
+                        "sparsity/overall_below_1e-5": (overall_act_freq_scores < 1e-5).sum().item(),
+                        "sparsity/overall_below_1e-6": (overall_act_freq_scores < 1e-6).sum().item(),
+                    }
+                )
 
             self.wandb_logger.log(wandb_log_dict, step=self.cur_step + 1)
             log_info["act_freq_scores"] = torch.zeros_like(log_info["act_freq_scores"])
@@ -161,7 +163,11 @@ def _log(self, sae: SparseAutoEncoder, log_info: dict, batch: dict[str, Tensor])
             wandb_log_dict = {
                 # losses
                 "losses/mse_loss": l_rec.item(),
-                **({"losses/sparsity_loss": log_info["l_s"].mean().item()} if log_info.get("l_s", None) is not None else {}),
+                **(
+                    {"losses/sparsity_loss": log_info["l_s"].mean().item()}
+                    if log_info.get("l_s", None) is not None
+                    else {}
+                ),
                 "losses/overall_loss": log_info["loss"].item(),
                 # variance explained
                 "metrics/explained_variance": explained_variance.mean().item(),
@@ -179,10 +185,16 @@ def _log(self, sae: SparseAutoEncoder, log_info: dict, batch: dict[str, Tensor])
                 "details/n_training_tokens": self.cur_tokens,
             }
             wandb_log_dict.update(sae.log_statistics())
-            if sae.cfg.sae_type == 'crosscoder':
-                wandb_log_dict.update({
-                    "metrics/overall_l0": all_reduce_tensor(log_info["feature_acts"], aggregate='max').gt(0).float().sum(-1).mean()
-                })
+            if sae.cfg.sae_type == "crosscoder":
+                wandb_log_dict.update(
+                    {
+                        "metrics/overall_l0": all_reduce_tensor(log_info["feature_acts"], aggregate="max")
+                        .gt(0)
+                        .float()
+                        .sum(-1)
+                        .mean()
+                    }
+                )
             elif sae.cfg.sae_type == "mixcoder":
                 assert isinstance(sae, MixCoder)
                 for modality, (start, end) in sae.modality_index.items():