use regular eig as backup for eigh

elk/training/eigen_reporter.py

@@ -244,16 +244,20 @@ def fit_streaming(self, truncated: bool = False) -> float:
  else:
  try:
  L, Q = torch.linalg.eigh(A)
- except torch.linalg.LinAlgError as e:
- # Check if the matrix has non-finite values
- if not A.isfinite().all():
- raise ValueError(
- "Fitting the reporter failed because the VINC matrix has "
- "non-finite entries. Usually this means the hidden states "
- "themselves had non-finite values."
- ) from e
- else:
- raise e
+ except torch.linalg.LinAlgError:
+ try:
+ L, Q = torch.linalg.eig(A)
+ L, Q = L.real, Q.real
+ except torch.linalg.LinAlgError as e:
+ # Check if the matrix has non-finite values
+ if not A.isfinite().all():
+ raise ValueError(
+ "Fitting the reporter failed because the VINC matrix has "
+ "non-finite entries. Usually this means the hidden states "
+ "themselves had non-finite values."
+ ) from e
+ else:
+ raise e
 
  L, Q = L[-self.config.num_heads :], Q[:, -self.config.num_heads :]
 

elk/training/sweep.py

@@ -2,6 +2,7 @@
 from dataclasses import InitVar, dataclass
 
 import numpy as np
+import torch
 
 from ..evaluation.evaluate import Eval
 from ..extraction import Extract, PromptConfig
@@ -25,6 +26,8 @@ class Sweep:
  """The step size for hyperparameter sweeps. Performs a 2D
  sweep over a and b in (var_weight, inv_weight, neg_cov_weight) = (a, 1 - b, b)
  If negative, no hyperparameter sweeps will be performed. Only valid for Eigen."""
+ skip_transfer_eval: bool = False
+ """Whether to perform transfer eval on every pair of datasets."""
 
  name: str | None = None
 
@@ -108,25 +111,30 @@ def execute(self):
  )
 
  run.out_dir = out_dir
- run.execute()
-
- if len(eval_datasets) > 1:
- print(colorize("== Transfer eval ==", "green"))
-
- # Now evaluate the reporter on the other datasets
- for eval_dataset in eval_datasets:
- # We already evaluated on this one during training
- if eval_dataset in train_datasets:
- continue
-
- data = deepcopy(run.data)
- data.model = model_str
- data.prompts.datasets = [eval_dataset]
-
- eval = Eval(
- data=data,
- source=str(run.out_dir),
- out_dir=out_dir,
- skip_supervised=run.supervised == "none",
- )
- eval.execute(highlight_color="green")
+ try:
+ run.execute()
+ except torch._C._LinAlgError as e: # type: ignore
+ print(colorize(f"LinAlgError: {e}", "red"))
+ continue
+
+ if not self.skip_transfer_eval:
+ if len(eval_datasets) > 1:
+ print(colorize("== Transfer eval ==", "green"))
+
+ # Now evaluate the reporter on the other datasets
+ for eval_dataset in eval_datasets:
+ # We already evaluated on this one during training
+ if eval_dataset in train_datasets:
+ continue
+
+ data = deepcopy(run.data)
+ data.model = model_str
+ data.prompts.datasets = [eval_dataset]
+
+ eval = Eval(
+ data=data,
+ source=str(run.out_dir),
+ out_dir=out_dir,
+ skip_supervised=run.supervised == "none",
+ )
+ eval.execute(highlight_color="green")