Merge pull request #71 from OpenMOSS/sae/persistant_dataset_norm

feat(sae): support saving/loading dataset_average_activation_norm to/from SAE state dict

Author: Frankstein73

src/lm_saes/sae.py

@@ -196,6 +196,11 @@ def _get_full_state_dict(self):  # should be overridden by subclasses
         if self.device_mesh and self.device_mesh["model"].size(0) > 1:
             state_dict = {k: v.full_tensor() if isinstance(v, DTensor) else v for k, v in state_dict.items()}
 
+        # Add dataset_average_activation_norm to state dict
+        if self.dataset_average_activation_norm is not None:
+            for hook_point, value in self.dataset_average_activation_norm.items():
+                state_dict[f"dataset_average_activation_norm.{hook_point}"] = torch.tensor(value)
+
         # If sparsity_include_decoder_norm is False, we need to normalize the decoder weight before saving
         # We use a deepcopy to avoid modifying the original weight to avoid affecting the training progress
         if not self.cfg.sparsity_include_decoder_norm:
@@ -489,6 +494,15 @@ def compute_loss(
             return loss, (loss_dict, aux_data)
         return loss
 
+    def _load_full_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
+        # Extract and set dataset_average_activation_norm if present
+        norm_keys = [k for k in state_dict.keys() if k.startswith("dataset_average_activation_norm.")]
+        if norm_keys:
+            dataset_norm = {key.split(".", 1)[1]: state_dict[key].item() for key in norm_keys}
+            self.set_dataset_average_activation_norm(dataset_norm)
+            state_dict = {k: v for k, v in state_dict.items() if not k.startswith("dataset_average_activation_norm.")}
+        self.load_state_dict(state_dict, strict=self.cfg.strict_loading)
+
     @classmethod
     def from_config(cls, cfg: SAEConfig) -> "SparseAutoEncoder":
         if cfg.sae_pretrained_name_or_path is None:
@@ -512,12 +526,16 @@ def from_config(cls, cfg: SAEConfig) -> "SparseAutoEncoder":
                 raise FileNotFoundError(f"Pretrained model not found at {cfg.sae_pretrained_name_or_path}")
 
         if ckpt_path.endswith(".safetensors"):
-            state_dict = safe.load_file(ckpt_path, device=cfg.device)
+            state_dict: dict[str, torch.Tensor] = safe.load_file(ckpt_path, device=cfg.device)
         else:
-            state_dict = torch.load(ckpt_path, map_location=cfg.device)["sae"]
+            state_dict: dict[str, torch.Tensor] = torch.load(
+                ckpt_path,
+                map_location=cfg.device,
+                weights_only=True,
+            )["sae"]
 
         model = cls(cfg)
-        model.load_state_dict(state_dict, strict=cfg.strict_loading)
+        model._load_full_state_dict(state_dict)
         return model
 
     @classmethod

tests/unit/test_sae.py

@@ -3,9 +3,6 @@
 import pytest
 import torch
 
-if not torch.cuda.is_available():
-    pytest.skip("CUDA device not available", allow_module_level=True)
-
 from lm_saes.config import SAEConfig
 from lm_saes.sae import SparseAutoEncoder
 
@@ -21,6 +18,7 @@ def sae_config() -> SAEConfig:
         dtype=torch.float32,
         act_fn="topk",
         jump_relu_threshold=2.0,
+        top_k=2,
     )
 
 
@@ -42,7 +40,7 @@ def sae(sae_config: SAEConfig, generator: torch.Generator) -> SparseAutoEncoder:
     )
     if sae_config.use_decoder_bias:
         sae.decoder.bias.data = torch.randn(
-            sae_config.d_sae, generator=generator, device=sae_config.device, dtype=sae_config.dtype
+            sae_config.d_model, generator=generator, device=sae_config.device, dtype=sae_config.dtype
         )
     if sae_config.use_glu_encoder:
         sae.encoder_glu.weight.data = torch.randn(
@@ -156,6 +154,20 @@ def test_compute_norm_factor(sae_config: SAEConfig, sae: SparseAutoEncoder):
     )
 
 
+def test_persistent_dataset_average_activation_norm(sae_config: SAEConfig, sae: SparseAutoEncoder):
+    sae.set_dataset_average_activation_norm({"in": 3.0, "out": 2.0})
+    assert sae.dataset_average_activation_norm == {"in": 3.0, "out": 2.0}
+    state_dict = sae._get_full_state_dict()
+    assert state_dict["dataset_average_activation_norm.in"] == 3.0
+    assert state_dict["dataset_average_activation_norm.out"] == 2.0
+
+    new_sae = SparseAutoEncoder(sae_config)
+    new_sae._load_full_state_dict(state_dict)
+    assert new_sae.cfg == sae.cfg
+    assert all(torch.allclose(p, q, atol=1e-4, rtol=1e-5) for p, q in zip(new_sae.parameters(), sae.parameters()))
+    assert new_sae.dataset_average_activation_norm == {"in": 3.0, "out": 2.0}
+
+
 def test_get_full_state_dict(sae_config: SAEConfig, sae: SparseAutoEncoder):
     sae_config.sparsity_include_decoder_norm = False
     state_dict = sae._get_full_state_dict()
@@ -205,3 +217,9 @@ def test_standardize_parameters_of_dataset_norm(sae_config: SAEConfig, sae: Spar
         assert torch.allclose(
             sae.decoder.bias.data, decoder_bias_data / math.sqrt(sae_config.d_model) * 2.0, atol=1e-4, rtol=1e-5
         )
+
+
+def test_forward(sae_config: SAEConfig, sae: SparseAutoEncoder):
+    sae.set_dataset_average_activation_norm({"in": 3.0, "out": 2.0})
+    output = sae.forward(torch.tensor([[1.0, 2.0]], device=sae_config.device, dtype=sae_config.dtype))
+    assert output.shape == (1, 2)