From cd3f990d093027cafaf4e527a5c34bb0df5f8d27 Mon Sep 17 00:00:00 2001 From: HoagyC Date: Mon, 11 Dec 2023 11:11:25 +0000 Subject: [PATCH 1/5] Save checkpoints in own folder with config json. --- sparse_autoencoder/train/pipeline.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/sparse_autoencoder/train/pipeline.py b/sparse_autoencoder/train/pipeline.py index e9a63f5d..f12ed785 100644 --- a/sparse_autoencoder/train/pipeline.py +++ b/sparse_autoencoder/train/pipeline.py @@ -1,6 +1,9 @@ """Default pipeline.""" from collections.abc import Iterable +from datetime import datetime, timezone from functools import partial +from json import dumps +import os from pathlib import Path from typing import final from urllib.parse import quote_plus @@ -121,7 +124,7 @@ def __init__( # noqa: PLR0913 self.loss = loss self.metrics = metrics self.optimizer = optimizer - self.run_name = run_name + self.run_name = run_name + datetime.now(tz=timezone.utc).strftime("-%Y-%m-%d-%H-%M-%S") self.source_data_batch_size = source_data_batch_size self.source_dataset = source_dataset self.source_model = source_model @@ -343,9 +346,17 @@ 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) + run_directory = self.checkpoint_directory / run_name_file_system_safe + if not run_directory.exists(): + run_directory.mkdir(parents=True) + if "config.json" not in os.listdir(run_directory): + with Path.open(run_directory / "config.json", "w") as config_file: + config_file.write(dumps(dict(wandb.config))) + file_path: Path = ( self.checkpoint_directory - / f"{run_name_file_system_safe}-{self.total_activations_trained_on}.pt" + / run_name_file_system_safe + / f"checkpoint-{self.total_activations_trained_on}activations.pt" ) torch.save(self.autoencoder.state_dict(), file_path) From 5d38e1b64cc5d63602818d5404c613534c4c4552 Mon Sep 17 00:00:00 2001 From: HoagyC Date: Mon, 11 Dec 2023 11:27:10 +0000 Subject: [PATCH 2/5] Add the hash of the git commit to help with debugging models. --- sparse_autoencoder/train/pipeline.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/sparse_autoencoder/train/pipeline.py b/sparse_autoencoder/train/pipeline.py index f12ed785..968e0daf 100644 --- a/sparse_autoencoder/train/pipeline.py +++ b/sparse_autoencoder/train/pipeline.py @@ -5,8 +5,10 @@ from json import dumps import os from pathlib import Path +import subprocess from typing import final from urllib.parse import quote_plus +import warnings from jaxtyping import Int, Int64 import torch @@ -341,6 +343,22 @@ def validate_sae(self, validation_number_activations: int) -> None: if wandb.run is not None: wandb.log(data=calculated, commit=False) + @staticmethod + def get_git_commit_hash() -> None | str: + """Get the Git commit hash of the current directory.""" + try: + return ( + subprocess.check_output(["/usr/bin/git", "rev-parse", "HEAD"]) # noqa: S603 + .decode("ascii") + .strip() + ) + except subprocess.CalledProcessError: + # Handle the case where the directory is not a Git repository + warnings.warn( + "Directory is not a Git repository, not logging commit hash", stacklevel=1 + ) + return None + @final def save_checkpoint(self) -> None: """Save the model as a checkpoint.""" @@ -350,8 +368,12 @@ def save_checkpoint(self) -> None: if not run_directory.exists(): run_directory.mkdir(parents=True) if "config.json" not in os.listdir(run_directory): + config_dict = dict(wandb.config) + git_hash = self.get_git_commit_hash() + if git_hash is not None: + config_dict["git_hash"] = git_hash with Path.open(run_directory / "config.json", "w") as config_file: - config_file.write(dumps(dict(wandb.config))) + config_file.write(dumps(config_dict, indent=4)) file_path: Path = ( self.checkpoint_directory From c57117cd1df5492d9d0cef5c667cdf5426ae7485 Mon Sep 17 00:00:00 2001 From: HoagyC Date: Mon, 11 Dec 2023 13:11:36 +0000 Subject: [PATCH 3/5] Use importlib.metadata to get version. --- sparse_autoencoder/train/pipeline.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/sparse_autoencoder/train/pipeline.py b/sparse_autoencoder/train/pipeline.py index 968e0daf..658451e9 100644 --- a/sparse_autoencoder/train/pipeline.py +++ b/sparse_autoencoder/train/pipeline.py @@ -2,6 +2,7 @@ from collections.abc import Iterable from datetime import datetime, timezone from functools import partial +import importlib.metadata from json import dumps import os from pathlib import Path @@ -355,10 +356,20 @@ def get_git_commit_hash() -> None | str: except subprocess.CalledProcessError: # Handle the case where the directory is not a Git repository warnings.warn( - "Directory is not a Git repository, not logging commit hash", stacklevel=1 + "Directory is not a Git repository, not logging commit hash", stacklevel=2 ) return None + @staticmethod + def get_package_version() -> str | None: + """Get the version of the package.""" + try: + return importlib.metadata.version("sparse-autoencoder") + except importlib.metadata.PackageNotFoundError: + # Handle the case where the directory is not a Git repository + warnings.warn("Package not found, not logging version", stacklevel=2) + return None + @final def save_checkpoint(self) -> None: """Save the model as a checkpoint.""" @@ -369,9 +380,13 @@ def save_checkpoint(self) -> None: run_directory.mkdir(parents=True) if "config.json" not in os.listdir(run_directory): config_dict = dict(wandb.config) + git_hash = self.get_git_commit_hash() - if git_hash is not None: - config_dict["git_hash"] = git_hash + config_dict["git_hash"] = git_hash + + package_version = self.get_package_version() + config_dict["package_version"] = package_version + with Path.open(run_directory / "config.json", "w") as config_file: config_file.write(dumps(config_dict, indent=4)) From 67a641ca9f9c6d8cdba9b9dd335f143173864547 Mon Sep 17 00:00:00 2001 From: HoagyC Date: Mon, 11 Dec 2023 13:15:30 +0000 Subject: [PATCH 4/5] Remove unnecessary and/or wrong comments. --- sparse_autoencoder/train/pipeline.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/sparse_autoencoder/train/pipeline.py b/sparse_autoencoder/train/pipeline.py index 658451e9..9a73fd41 100644 --- a/sparse_autoencoder/train/pipeline.py +++ b/sparse_autoencoder/train/pipeline.py @@ -354,7 +354,6 @@ def get_git_commit_hash() -> None | str: .strip() ) except subprocess.CalledProcessError: - # Handle the case where the directory is not a Git repository warnings.warn( "Directory is not a Git repository, not logging commit hash", stacklevel=2 ) @@ -366,7 +365,6 @@ def get_package_version() -> str | None: try: return importlib.metadata.version("sparse-autoencoder") except importlib.metadata.PackageNotFoundError: - # Handle the case where the directory is not a Git repository warnings.warn("Package not found, not logging version", stacklevel=2) return None From 31c2687ba787bad38043445a8a0943f591092cbc Mon Sep 17 00:00:00 2001 From: HoagyC Date: Mon, 11 Dec 2023 15:14:07 +0000 Subject: [PATCH 5/5] Add mean to see if less noisy. --- .../metrics/validate/model_reconstruction_score.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sparse_autoencoder/metrics/validate/model_reconstruction_score.py b/sparse_autoencoder/metrics/validate/model_reconstruction_score.py index c3eedfc0..5f940ada 100644 --- a/sparse_autoencoder/metrics/validate/model_reconstruction_score.py +++ b/sparse_autoencoder/metrics/validate/model_reconstruction_score.py @@ -9,6 +9,7 @@ if TYPE_CHECKING: from sparse_autoencoder.tensor_types import Axis + from jaxtyping import Float from torch import Tensor @@ -67,6 +68,10 @@ def calculate(self, data: ValidationMetricData) -> dict[str, Any]: zero_ablate_loss_minus_reconstruction_loss / zero_ablate_loss_minus_default_loss ) model_reconstruction_score: float = model_reconstruction_score_itemwise.mean().item() + model_reconstruction_score_average: float = ( + zero_ablate_loss_minus_reconstruction_loss.mean().item() + / zero_ablate_loss_minus_default_loss.mean().item() + ) # Get the other metrics validation_baseline_loss: float = data.source_model_loss.mean().item() @@ -82,4 +87,5 @@ def calculate(self, data: ValidationMetricData) -> dict[str, Any]: "validate/loss_with_reconstruction": validation_loss_with_reconstruction, "validate/loss_with_zero_ablation": validation_loss_with_zero_ablation, "validate/model_reconstruction_score": model_reconstruction_score, + "validate/model_reconstruction_score_average": model_reconstruction_score_average, }