mi-crow 0.1.1.post12__py3-none-any.whl

amber/mechanistic/sae/training/wandb_logger.py

@@ -0,0 +1,222 @@
+"""Wandb logging utilities for SAE training."""
+
+from typing import Any, Optional
+
+from amber.mechanistic.sae.sae_trainer import SaeTrainingConfig
+from amber.utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class WandbLogger:
+    """
+    Handles wandb logging for SAE training.
+    
+    Encapsulates all wandb-related operations including initialization,
+    metric logging, and summary updates.
+    """
+
+    def __init__(self, config: SaeTrainingConfig, run_id: str):
+        """
+        Initialize WandbLogger.
+        
+        Args:
+            config: Training configuration
+            run_id: Training run identifier
+        """
+        self.config = config
+        self.run_id = run_id
+        self.wandb_run: Optional[Any] = None
+        self._initialized = False
+
+    def initialize(self) -> bool:
+        """
+        Initialize wandb run if enabled in config.
+        
+        Returns:
+            True if wandb was successfully initialized, False otherwise
+        """
+        if not self.config.use_wandb:
+            return False
+
+        try:
+            import wandb
+        except ImportError:
+            logger.warning("[WandbLogger] wandb not installed, skipping wandb logging")
+            logger.warning("[WandbLogger] Install with: pip install wandb")
+            return False
+
+        try:
+            wandb_project = self.config.wandb_project or "sae-training"
+            wandb_name = self.config.wandb_name or self.run_id
+            wandb_mode = self.config.wandb_mode.lower() if self.config.wandb_mode else "online"
+
+            self.wandb_run = wandb.init(
+                project=wandb_project,
+                entity=self.config.wandb_entity,
+                name=wandb_name,
+                mode=wandb_mode,
+                config=self._build_wandb_config(),
+                tags=self.config.wandb_tags or [],
+            )
+            self._initialized = True
+            return True
+        except Exception as e:
+            logger.warning(f"[WandbLogger] Unexpected error initializing wandb: {e}")
+            logger.warning("[WandbLogger] Continuing training without wandb logging")
+            return False
+
+    def _build_wandb_config(self) -> dict[str, Any]:
+        """Build wandb configuration dictionary."""
+        return {
+            "run_id": self.run_id,
+            "epochs": self.config.epochs,
+            "batch_size": self.config.batch_size,
+            "lr": self.config.lr,
+            "l1_lambda": self.config.l1_lambda,
+            "device": str(self.config.device),
+            "dtype": str(self.config.dtype) if self.config.dtype else None,
+            "use_amp": self.config.use_amp,
+            "clip_grad": self.config.clip_grad,
+            "max_batches_per_epoch": self.config.max_batches_per_epoch,
+            **(self.config.wandb_config or {}),
+        }
+
+    def log_metrics(
+            self,
+            history: dict[str, list[float | None]],
+            verbose: bool = False
+    ) -> None:
+        """
+        Log training metrics to wandb.
+        
+        Args:
+            history: Dictionary with training history (loss, r2, l1, l0, etc.)
+            verbose: Whether to log verbose information
+        """
+        if not self._initialized or self.wandb_run is None:
+            return
+
+        try:
+            num_epochs = len(history.get("loss", []))
+            slow_metrics_freq = self.config.wandb_slow_metrics_frequency
+
+            # Helper to get last known value for slow metrics
+            def get_last_known_value(values: list[float | None], idx: int) -> float:
+                """Get the last non-None value up to idx, or 0.0 if none found."""
+                for i in range(idx, -1, -1):
+                    if i < len(values) and values[i] is not None:
+                        return values[i]
+                return 0.0
+
+            # Log metrics for each epoch
+            for epoch in range(1, num_epochs + 1):
+                epoch_idx = epoch - 1
+                should_log_slow = (epoch % slow_metrics_freq == 0) or (epoch == num_epochs)
+
+                metrics = self._build_epoch_metrics(history, epoch_idx, should_log_slow, get_last_known_value)
+                self.wandb_run.log(metrics)
+
+            # Log final summary metrics
+            self._log_summary_metrics(history, get_last_known_value)
+
+            if verbose:
+                self._log_wandb_url()
+
+        except Exception as e:
+            logger.warning(f"[WandbLogger] Failed to log metrics to wandb: {e}")
+
+    def _build_epoch_metrics(
+            self,
+            history: dict[str, list[float | None]],
+            epoch_idx: int,
+            should_log_slow: bool,
+            get_last_known_value: Any
+    ) -> dict[str, Any]:
+        """Build metrics dictionary for a single epoch."""
+        metrics = {
+            "epoch": epoch_idx + 1,
+            "train/loss": self._safe_get(history["loss"], epoch_idx, 0.0),
+            "train/reconstruction_mse": self._safe_get(history["recon_mse"], epoch_idx, 0.0),
+            "train/r2_score": self._safe_get(history["r2"], epoch_idx, 0.0),
+            "train/l1_penalty": self._safe_get(history["l1"], epoch_idx, 0.0),
+            "train/learning_rate": self.config.lr,
+        }
+
+        # Add slow metrics if computed this epoch
+        if should_log_slow:
+            l0_val = self._get_slow_metric(history["l0"], epoch_idx, get_last_known_value)
+            dead_pct = self._get_slow_metric(history["dead_features_pct"], epoch_idx, get_last_known_value)
+            metrics["train/l0_sparsity"] = l0_val
+            metrics["train/dead_features_pct"] = dead_pct
+
+        return metrics
+
+    def _get_slow_metric(
+            self,
+            values: list[float | None],
+            epoch_idx: int,
+            get_last_known_value: Any
+    ) -> float:
+        """Get slow metric value, using last known value if current is None."""
+        if epoch_idx < len(values) and values[epoch_idx] is not None:
+            return values[epoch_idx]
+        return get_last_known_value(values, epoch_idx)
+
+    def _safe_get(self, values: list[float | None], idx: int, default: float) -> float:
+        """Safely get value from list, returning default if out of bounds."""
+        if idx < len(values) and values[idx] is not None:
+            return values[idx]
+        return default
+
+    def _log_summary_metrics(
+            self,
+            history: dict[str, list[float | None]],
+            get_last_known_value: Any
+    ) -> None:
+        """Log final summary metrics to wandb."""
+        if self.wandb_run is None:
+            return
+
+        # Get last computed values for slow metrics
+        final_l0 = get_last_known_value(history["l0"], len(history["l0"]) - 1) if history.get("l0") else 0.0
+        final_dead_pct = get_last_known_value(
+            history["dead_features_pct"],
+            len(history["dead_features_pct"]) - 1
+        ) if history.get("dead_features_pct") else 0.0
+
+        final_metrics = {
+            "final/loss": history["loss"][-1] if history.get("loss") else 0.0,
+            "final/reconstruction_mse": history["recon_mse"][-1] if history.get("recon_mse") else 0.0,
+            "final/r2_score": history["r2"][-1] if history.get("r2") else 0.0,
+            "final/l1_penalty": history["l1"][-1] if history.get("l1") else 0.0,
+            "final/l0_sparsity": final_l0,
+            "final/dead_features_pct": final_dead_pct,
+            "training/num_epochs": len(history.get("loss", [])),
+        }
+
+        # Add best metrics
+        if history.get("loss"):
+            best_loss_idx = min(range(len(history["loss"])), key=lambda i: history["loss"][i] or float('inf'))
+            final_metrics["best/loss"] = history["loss"][best_loss_idx] or 0.0
+            final_metrics["best/loss_epoch"] = best_loss_idx + 1
+
+        if history.get("r2"):
+            best_r2_idx = max(range(len(history["r2"])), key=lambda i: history["r2"][i] or -float('inf'))
+            final_metrics["best/r2_score"] = history["r2"][best_r2_idx] or 0.0
+            final_metrics["best/r2_epoch"] = best_r2_idx + 1
+
+        self.wandb_run.summary.update(final_metrics)
+
+    def _log_wandb_url(self) -> None:
+        """Log wandb run URL if available."""
+        if self.wandb_run is None:
+            return
+
+        try:
+            url = self.wandb_run.url
+            logger.info(f"[WandbLogger] Metrics logged to wandb: {url}")
+        except (AttributeError, RuntimeError):
+            # Offline mode or URL not available
+            logger.info("[WandbLogger] Metrics logged to wandb (offline mode)")
+

amber/store/__init__.py

@@ -0,0 +1,5 @@
+from amber.store.store import Store
+from amber.store.local_store import LocalStore
+
+__all__ = ["Store", "LocalStore"]
+

amber/store/local_store.py

@@ -0,0 +1,437 @@
+import json
+from pathlib import Path
+from typing import Dict, Any, List
+import shutil
+
+import torch
+
+from amber.store.store import Store, TensorMetadata
+import safetensors.torch as storch
+
+
+class LocalStore(Store):
+    """Local filesystem implementation of Store interface."""
+
+    def __init__(
+            self,
+            base_path: Path | str = '',
+            runs_prefix: str = "runs",
+            dataset_prefix: str = "datasets",
+            model_prefix: str = "models",
+    ):
+        """Initialize LocalStore.
+        
+        Args:
+            base_path: Base directory path for the store
+            runs_prefix: Prefix for runs directory
+            dataset_prefix: Prefix for datasets directory
+            model_prefix: Prefix for models directory
+        """
+        super().__init__(base_path, runs_prefix, dataset_prefix, model_prefix)
+
+    def _full(self, key: str) -> Path:
+        p = self.base_path / key
+        p.parent.mkdir(parents=True, exist_ok=True)
+        return p
+
+    def put_tensor(self, key: str, tensor: torch.Tensor) -> None:
+        path = self._full(key)
+        storch.save_file({"tensor": tensor}, str(path))
+
+    def get_tensor(self, key: str) -> torch.Tensor:
+        loaded = storch.load_file(str(self._full(key)))
+        return loaded["tensor"]
+
+    def _validate_run_id(self, run_id: str) -> None:
+        if not run_id or not isinstance(run_id, str) or not run_id.strip():
+            raise ValueError(f"run_id must be a non-empty string, got: {run_id!r}")
+
+    def _validate_batch_index(self, batch_index: int) -> None:
+        if batch_index < 0:
+            raise ValueError(f"batch_index must be non-negative, got: {batch_index}")
+
+    def _validate_layer_key(self, layer: str, key: str) -> None:
+        if not layer or not isinstance(layer, str) or not layer.strip():
+            raise ValueError(f"layer must be a non-empty string, got: {layer!r}")
+        if not key or not isinstance(key, str) or not key.strip():
+            raise ValueError(f"key must be a non-empty string, got: {key!r}")
+
+    def _ensure_directory(self, path: Path) -> None:
+        path.mkdir(parents=True, exist_ok=True)
+
+    def put_run_batch(self, run_id: str, batch_index: int,
+                      tensors: List[torch.Tensor] | Dict[str, torch.Tensor]) -> str:
+        if isinstance(tensors, dict):
+            to_save = tensors
+        elif isinstance(tensors, list):
+            if len(tensors) == 0:
+                to_save = {"_empty_list": torch.tensor([])}
+            else:
+                to_save = {f"item_{i}": t for i, t in enumerate(tensors)}
+        else:
+            to_save = {}
+        batch_path = self.base_path / f"{self.runs_prefix}/{run_id}/batch_{batch_index:06d}.safetensors"
+        self._ensure_directory(batch_path.parent)
+        storch.save_file(to_save, str(batch_path))
+        return f"{self.runs_prefix}/{run_id}/batch_{batch_index:06d}.safetensors"
+
+    def get_run_batch(self, run_id: str, batch_index: int) -> List[torch.Tensor] | Dict[
+        str, torch.Tensor]:
+
+        batch_path = self.base_path / f"{self.runs_prefix}/{run_id}/batch_{batch_index:06d}.safetensors"
+        if batch_path.exists():
+            loaded = storch.load_file(str(batch_path))
+            keys = list(loaded.keys())
+            if keys == ["_empty_list"]:
+                return []
+            if keys and all(k.startswith("item_") for k in keys):
+                try:
+                    items = sorted(((int(k.split("_", 1)[1]), v) for k, v in loaded.items()), key=lambda x: x[0])
+                    if [i for i, _ in items] == list(range(len(items))):
+                        return [v for _, v in items]
+                except Exception:
+                    pass
+            return loaded
+
+        detector_base = self.base_path / self.runs_prefix / run_id / f"batch_{batch_index}"
+        if detector_base.exists():
+            result: Dict[str, torch.Tensor] = {}
+
+            layer_dirs = [d for d in detector_base.iterdir() if d.is_dir()]
+            for layer_dir in layer_dirs:
+                activations_path = layer_dir / "activations.safetensors"
+                if activations_path.exists():
+                    try:
+                        loaded_tensor = storch.load_file(str(activations_path))["tensor"]
+                        # Use layer_signature as key, or "activations" if only one layer
+                        layer_sig = layer_dir.name
+                        if len(layer_dirs) == 1:
+                            # Only one layer, use simple "activations" key for compatibility
+                            result["activations"] = loaded_tensor
+                        else:
+                            # Multiple layers, use layer-specific key
+                            result[f"activations_{layer_sig}"] = loaded_tensor
+                    except Exception:
+                        pass
+
+            if result:
+                return result
+
+        # If neither exists, raise FileNotFoundError
+        raise FileNotFoundError(f"Batch {batch_index} not found for run {run_id}")
+
+    def list_run_batches(self, run_id: str) -> List[int]:
+        base = self.base_path / self.runs_prefix / run_id
+        if not base.exists():
+            return []
+        out: set[int] = set()
+        
+        for p in sorted(base.glob("batch_*.safetensors")):
+            name = p.name
+            try:
+                idx = int(name[len("batch_"): len("batch_") + 6])
+                out.add(idx)
+            except Exception:
+                continue
+        
+        for p in sorted(base.glob("batch_*")):
+            if p.is_dir():
+                name = p.name
+                try:
+                    idx = int(name[len("batch_"):])
+                    out.add(idx)
+                except Exception:
+                    continue
+        
+        return sorted(list(out))
+
+    def delete_run(self, run_id: str) -> None:
+        base = self.base_path / self.runs_prefix / run_id
+        if not base.exists():
+            return
+        for p in base.glob("batch_*.safetensors"):
+            if p.is_file():
+                p.unlink()
+        for p in base.glob("batch_*"):
+            if p.is_dir():
+                shutil.rmtree(p, ignore_errors=True)
+        metadata_path = self._run_metadata_key(run_id)
+        if metadata_path.exists():
+            metadata_path.unlink()
+
+    def put_run_metadata(self, run_id: str, meta: Dict[str, Any]) -> str:
+        self._validate_run_id(run_id)
+        
+        metadata_path = self._run_metadata_key(run_id)
+        self._ensure_directory(metadata_path.parent)
+        
+        try:
+            with metadata_path.open("w", encoding="utf-8") as f:
+                json.dump(meta, f, ensure_ascii=False, indent=2)
+        except (TypeError, ValueError) as e:
+            raise ValueError(
+                f"Metadata is not JSON-serializable for run_id={run_id!r}. "
+                f"Error: {e}"
+            ) from e
+        except OSError as e:
+            raise OSError(
+                f"Failed to write metadata file at {metadata_path} for run_id={run_id!r}. "
+                f"Error: {e}"
+            ) from e
+        
+        return f"{self.runs_prefix}/{run_id}/meta.json"
+
+    def get_run_metadata(self, run_id: str) -> Dict[str, Any]:
+        self._validate_run_id(run_id)
+        
+        metadata_path = self._run_metadata_key(run_id)
+        if not metadata_path.exists():
+            return {}
+        
+        try:
+            with metadata_path.open("r", encoding="utf-8") as f:
+                return json.load(f)
+        except json.JSONDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Invalid JSON in metadata file at {metadata_path} for run_id={run_id!r}",
+                e.doc,
+                e.pos
+            ) from e
+
+    def put_detector_metadata(
+            self,
+            run_id: str,
+            batch_index: int,
+            metadata: Dict[str, Any],
+            tensor_metadata: TensorMetadata
+    ) -> str:
+        self._validate_run_id(run_id)
+        self._validate_batch_index(batch_index)
+        
+        batch_dir = self._run_batch_key(run_id, batch_index)
+        self._ensure_directory(batch_dir)
+
+        tensor_metadata_names = {
+            str(layer_signature): list(detector_tensors.keys())
+            for layer_signature, detector_tensors in tensor_metadata.items()
+            if detector_tensors
+        }
+        metadata_with_tensor_names = {
+            **metadata,
+            "_tensor_metadata_names": tensor_metadata_names
+        }
+
+        detector_metadata_path = batch_dir / "metadata.json"
+        try:
+            with detector_metadata_path.open("w", encoding="utf-8") as f:
+                json.dump(metadata_with_tensor_names, f, ensure_ascii=False, indent=2)
+        except (TypeError, ValueError) as e:
+            raise ValueError(
+                f"Metadata is not JSON-serializable for run_id={run_id!r}, "
+                f"batch_index={batch_index}. Error: {e}"
+            ) from e
+        except OSError as e:
+            raise OSError(
+                f"Failed to write metadata file at {detector_metadata_path} for "
+                f"run_id={run_id!r}, batch_index={batch_index}. Error: {e}"
+            ) from e
+
+        for layer_signature, detector_tensors in tensor_metadata.items():
+            if not detector_tensors:
+                continue
+
+            layer_dir = batch_dir / layer_signature
+            self._ensure_directory(layer_dir)
+
+            # Save each tensor key (e.g., "activations") as a separate safetensors file
+            for tensor_key, tensor in detector_tensors.items():
+                tensor_filename = f"{tensor_key}.safetensors"
+                tensor_path = layer_dir / tensor_filename
+                try:
+                    storch.save_file({"tensor": tensor}, str(tensor_path))
+                except Exception as e:
+                    raise OSError(
+                        f"Failed to save tensor at {tensor_path} for run_id={run_id!r}, "
+                        f"batch_index={batch_index}, layer={layer_signature!r}, "
+                        f"key={tensor_key!r}. Error: {e}"
+                    ) from e
+
+        return f"{self.runs_prefix}/{run_id}/batch_{batch_index}"
+
+    def get_detector_metadata(
+            self,
+            run_id: str,
+            batch_index: int
+    ) -> tuple[Dict[str, Any], TensorMetadata]:
+        self._validate_run_id(run_id)
+        self._validate_batch_index(batch_index)
+        
+        batch_dir = self._run_batch_key(run_id, batch_index)
+        metadata_path = batch_dir / "metadata.json"
+        
+        if not metadata_path.exists():
+            return {}, {}
+        
+        try:
+            with metadata_path.open("r", encoding="utf-8") as f:
+                metadata = json.load(f)
+        except json.JSONDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Invalid JSON in metadata file at {metadata_path} for "
+                f"run_id={run_id!r}, batch_index={batch_index}",
+                e.doc,
+                e.pos
+            ) from e
+
+        tensor_metadata: Dict[str, Dict[str, torch.Tensor]] = {}
+        tensor_metadata_names = metadata.pop("_tensor_metadata_names", None)
+
+        if tensor_metadata_names is not None:
+            for layer_signature, tensor_keys in tensor_metadata_names.items():
+                layer_dir = batch_dir / layer_signature
+                detector_tensors: Dict[str, torch.Tensor] = {}
+                for tensor_key in tensor_keys:
+                    tensor_filename = f"{tensor_key}.safetensors"
+                    tensor_path = layer_dir / tensor_filename
+                    if tensor_path.exists():
+                        try:
+                            detector_tensors[tensor_key] = storch.load_file(str(tensor_path))["tensor"]
+                        except Exception as e:
+                            raise OSError(
+                                f"Failed to load tensor at {tensor_path} for "
+                                f"run_id={run_id!r}, batch_index={batch_index}, "
+                                f"layer={layer_signature!r}, key={tensor_key!r}. Error: {e}"
+                            ) from e
+                if detector_tensors:
+                    tensor_metadata[layer_signature] = detector_tensors
+        else:
+            raise ValueError(
+                f"Field '_tensor_metadata_names' not found in detector metadata at "
+                f"{metadata_path} for run_id={run_id!r}, batch_index={batch_index}. "
+                f"Cannot retrieve tensors."
+            )
+
+        return metadata, tensor_metadata
+
+    def get_detector_metadata_by_layer_by_key(
+            self,
+            run_id: str,
+            batch_index: int,
+            layer: str,
+            key: str
+    ) -> torch.Tensor:
+        self._validate_run_id(run_id)
+        self._validate_batch_index(batch_index)
+        self._validate_layer_key(layer, key)
+        
+        batch_dir = self._run_batch_key(run_id, batch_index)
+        layer_dir = batch_dir / layer
+        tensor_path = layer_dir / f"{key}.safetensors"
+        
+        if not tensor_path.exists():
+            raise FileNotFoundError(
+                f"Tensor not found at {tensor_path} for run_id={run_id!r}, "
+                f"batch_index={batch_index}, layer={layer!r}, key={key!r}"
+            )
+        
+        try:
+            return storch.load_file(str(tensor_path))["tensor"]
+        except Exception as e:
+            raise OSError(
+                f"Failed to load tensor at {tensor_path} for run_id={run_id!r}, "
+                f"batch_index={batch_index}, layer={layer!r}, key={key!r}. Error: {e}"
+            ) from e
+
+    def put_run_detector_metadata(
+            self,
+            run_id: str,
+            metadata: Dict[str, Any],
+            tensor_metadata: TensorMetadata,
+    ) -> str:
+        self._validate_run_id(run_id)
+
+        detectors_dir = self._run_key(run_id) / "detectors"
+        self._ensure_directory(detectors_dir)
+
+        metadata_path = detectors_dir / "metadata.json"
+
+        if metadata_path.exists():
+            try:
+                with metadata_path.open("r", encoding="utf-8") as f:
+                    aggregated = json.load(f)
+            except json.JSONDecodeError as e:
+                raise json.JSONDecodeError(
+                    f"Invalid JSON in unified detector metadata at {metadata_path} for run_id={run_id!r}",
+                    e.doc,
+                    e.pos,
+                ) from e
+        else:
+            aggregated = {"batches": []}
+
+        batches = aggregated.setdefault("batches", [])
+        batch_index = len(batches)
+
+        tensor_metadata_names = {
+            layer_signature: list(detector_tensors.keys())
+            for layer_signature, detector_tensors in tensor_metadata.items()
+            if detector_tensors
+        }
+
+        batch_entry = {
+            **metadata,
+            "batch_index": batch_index,
+            "_tensor_metadata_names": tensor_metadata_names,
+        }
+
+        batches.append(batch_entry)
+
+        try:
+            with metadata_path.open("w", encoding="utf-8") as f:
+                json.dump(aggregated, f, ensure_ascii=False, indent=2)
+        except (TypeError, ValueError) as e:
+            raise ValueError(
+                f"Unified detector metadata is not JSON-serializable for run_id={run_id!r}. "
+                f"Error: {e}"
+            ) from e
+        except OSError as e:
+            raise OSError(
+                f"Failed to write unified detector metadata at {metadata_path} for run_id={run_id!r}. "
+                f"Error: {e}"
+            ) from e
+
+        for layer_signature, detector_tensors in tensor_metadata.items():
+            if not detector_tensors:
+                continue
+
+            layer_dir = detectors_dir / str(layer_signature)
+            self._ensure_directory(layer_dir)
+
+            for tensor_key, tensor in detector_tensors.items():
+                tensor_filename = f"{tensor_key}.safetensors"
+                tensor_path = layer_dir / tensor_filename
+
+                if tensor_path.exists():
+                    try:
+                        existing = storch.load_file(str(tensor_path))
+                    except Exception as e:
+                        raise OSError(
+                            f"Failed to load existing unified tensor at {tensor_path} for "
+                            f"run_id={run_id!r}, layer={layer_signature!r}, key={tensor_key!r}. "
+                            f"Error: {e}"
+                        ) from e
+                else:
+                    existing = {}
+
+                batch_key = f"batch_{batch_index}"
+                existing[batch_key] = tensor
+
+                try:
+                    storch.save_file(existing, str(tensor_path))
+                except Exception as e:
+                    raise OSError(
+                        f"Failed to save unified tensor at {tensor_path} for run_id={run_id!r}, "
+                        f"layer={layer_signature!r}, key={tensor_key!r}, batch_index={batch_index}. "
+                        f"Error: {e}"
+                    ) from e
+
+        return f"{self.runs_prefix}/{run_id}/detectors"