Add save to wandb (#153)

Author: alan-cooney

sparse_autoencoder/train/pipeline.py

@@ -2,6 +2,7 @@
 from collections.abc import Iterable
 from functools import partial
 from pathlib import Path
+import tempfile
 from typing import final
 from urllib.parse import quote_plus
 
@@ -32,6 +33,9 @@
 from sparse_autoencoder.train.utils import get_model_device
 
 
+DEFAULT_CHECKPOINT_DIRECTORY: Path = Path(tempfile.gettempdir()) / "sparse_autoencoder"
+
+
 class Pipeline:
     """Pipeline for training a Sparse Autoencoder on TransformerLens activations.
 
@@ -90,7 +94,7 @@ def __init__(
         source_dataset: SourceDataset,
         source_model: HookedTransformer,
         run_name: str = "sparse_autoencoder",
-        checkpoint_directory: Path | None = None,
+        checkpoint_directory: Path = DEFAULT_CHECKPOINT_DIRECTORY,
         log_frequency: int = 100,
         metrics: MetricsContainer = default_metrics,
         source_data_batch_size: int = 12,
@@ -339,15 +343,34 @@ def validate_sae(self, validation_number_activations: int) -> None:
                 wandb.log(data=calculated, commit=False)
 
     @final
-    def save_checkpoint(self) -> None:
-        """Save the model as a checkpoint."""
-        if self.checkpoint_directory:
-            run_name_file_system_safe = quote_plus(self.run_name)
-            file_path: Path = (
-                self.checkpoint_directory
-                / f"{run_name_file_system_safe}-{self.total_activations_trained_on}.pt"
-            )
-            torch.save(self.autoencoder.state_dict(), file_path)
+    def save_checkpoint(self, *, is_final: bool = False) -> Path:
+        """Save the model as a checkpoint.
+
+        Args:
+            is_final: Whether this is the final checkpoint.
+
+        Returns:
+            Path to the saved checkpoint.
+        """
+        # Create the name
+        name: str = f"{self.run_name}_{'final' if is_final else self.total_activations_trained_on}"
+        safe_name = quote_plus(name, safe="_")
+
+        # Save locally
+        self.checkpoint_directory.mkdir(parents=True, exist_ok=True)
+        file_path: Path = self.checkpoint_directory / f"{safe_name}.pt"
+        torch.save(
+            self.autoencoder.state_dict(),
+            file_path,
+        )
+
+        # Upload to wandb
+        if wandb.run is not None:
+            artifact = wandb.Artifact(safe_name, type="model")
+            artifact.add_file(str(file_path))
+            wandb.log_artifact(artifact)
+
+        return file_path
 
     def run_pipeline(
         self,
@@ -440,6 +463,9 @@ def run_pipeline(
                 # Update the progress bar
                 progress_bar.update(store_size)
 
+        # Save the final checkpoint
+        self.save_checkpoint(is_final=True)
+
     @staticmethod
     def stateful_dataloader_iterable(
         dataloader: DataLoader[TorchTokenizedPrompts],

sparse_autoencoder/train/sweep.py

@@ -217,6 +217,7 @@ def run_training_pipeline(
     hook_point = get_act_name(
         hyperparameters["source_model"]["hook_site"], hyperparameters["source_model"]["hook_layer"]
     )
+
     pipeline = Pipeline(
         activation_resampler=activation_resampler,
         autoencoder=autoencoder,

sparse_autoencoder/train/tests/test_pipeline.py

@@ -1,5 +1,5 @@
 """Test the pipeline module."""
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from unittest.mock import MagicMock
 
 import pytest
@@ -26,6 +26,10 @@
 from sparse_autoencoder.source_data.mock_dataset import MockDataset
 
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
 @pytest.fixture()
 def pipeline_fixture() -> Pipeline:
     """Fixture to create a Pipeline instance for testing."""
@@ -289,6 +293,22 @@ def calculate(self, data: ValidationMetricData) -> dict[str, Any]:
         ), "Zero ablation loss should be more than base loss."
 
 
+class TestSaveCheckpoint:
+    """Test the save_checkpoint method."""
+
+    def test_saves_locally(self, pipeline_fixture: Pipeline) -> None:
+        """Test that the save_checkpoint method saves the checkpoint locally."""
+        saved_checkpoint: Path = pipeline_fixture.save_checkpoint()
+        assert saved_checkpoint.exists(), "Checkpoint file should exist."
+
+    def test_saves_final(self, pipeline_fixture: Pipeline) -> None:
+        """Test that the save_checkpoint method saves the final checkpoint."""
+        saved_checkpoint: Path = pipeline_fixture.save_checkpoint(is_final=True)
+        assert (
+            "final.pt" in saved_checkpoint.name
+        ), "Checkpoint file should be named '<run_name>_final.pt'."
+
+
 class TestRunPipeline:
     """Test the run_pipeline method."""
 
@@ -306,15 +326,15 @@ def test_run_pipeline_calls_all_methods(self, pipeline_fixture: Pipeline) -> Non
 
         total_loops = 5
         validate_expected_calls = 2
-        checkpoint_expected_calls = 5
+        checkpoint_expected_calls = 6  # Includes final
 
         pipeline_fixture.run_pipeline(
             train_batch_size=train_batch_size,
             max_store_size=store_size,
             max_activations=store_size * 5,
             validation_number_activations=store_size,
             validate_frequency=store_size * (total_loops // validate_expected_calls),
-            checkpoint_frequency=store_size * (total_loops // checkpoint_expected_calls),
+            checkpoint_frequency=store_size * (total_loops // checkpoint_expected_calls - 1),
         )
 
         # Check the number of calls