diff --git a/src/clt_forge/clt.py b/src/clt_forge/clt.py index e15a0cb..b44f63c 100644 --- a/src/clt_forge/clt.py +++ b/src/clt_forge/clt.py @@ -13,7 +13,7 @@ from clt_forge.config import CLTConfig from clt_forge.utils import DTYPE_MAP, CLT_WEIGHTS_FILENAME, CLT_CFG_FILENAME -from clt_forge.training.optim import JumpReLU +from clt_forge.training.optim import JumpReLU, fused_encoder from clt_forge import logger C_l0_COEF = 4 @@ -97,6 +97,8 @@ def __init__(self, cfg: CLTConfig, rank: int = 0, world_size: int = 1) -> None: self.log_threshold = nn.Parameter( torch.full((self.N_layers, self.local_d_latent), math.log(cfg.jumprelu_init_threshold), dtype=self.dtype, device=init_device) ) + if cfg.activation_fn != "jumprelu": + self.log_threshold.requires_grad = False self.bandwidth = cfg.jumprelu_bandwidth self.register_buffer('feature_count', @@ -153,7 +155,7 @@ def _initialize_b_enc(self, hidden_pre: Float[torch.Tensor, "..."]) -> None: for feature in range(self.local_d_latent): feature_pre_acts = hidden_pre[:, layer, feature] # [B] sorted_acts, _ = torch.sort(feature_pre_acts, descending=True) - target_idx = int(target_activation_rate * B) + 1 + target_idx = min(int(target_activation_rate * B) + 1, B - 1) threshold_value = sorted_acts[target_idx] required_bias = thresh[layer, feature] - threshold_value @@ -182,24 +184,40 @@ def encode( output: tuple([B, N_layers, local_d_latent], [B, N_layers, local_d_latent]) if layer is None, else [B, local_d_latent] """ - if layer is None: - hidden_pre = (torch.einsum( # double check einsum and autocast - "bnd,ndk->bnk", - x, - self.W_enc, - ) + self.b_enc) - - thresh = torch.exp(self.log_threshold) #shape [N_layers, d_latent] - else: - assert 0 <= layer < self.N_layers, f"Layer {layer} out of range" - hidden_pre = F.linear( - x, - self.W_enc[layer].T, - self.b_enc[layer] - ) - thresh = torch.exp(self.log_threshold[layer]) - - feat_act = JumpReLU.apply(hidden_pre, thresh, self.bandwidth) + if self.cfg.activation_fn in ["topk", "groupmax"]: + assert self.cfg.k is not None, f"k must be specified in config for {self.cfg.activation_fn} activation" + if layer is None: + weight = self.W_enc + bias = self.b_enc + else: + assert 0 <= layer < self.N_layers, f"Layer {layer} out of range" + weight = self.W_enc[layer] + bias = self.b_enc[layer] + + values, indices, preacts = fused_encoder(x, weight, bias, self.cfg.k, self.cfg.activation_fn) + feat_act = torch.zeros_like(preacts) + feat_act.scatter_(-1, indices, values) + hidden_pre = preacts + elif self.cfg.activation_fn == "jumprelu": + if layer is None: + hidden_pre = (torch.einsum( + "bnd,ndk->bnk", + x, + self.W_enc, + ) + self.b_enc) + thresh = torch.exp(self.log_threshold) + else: + assert 0 <= layer < self.N_layers, f"Layer {layer} out of range" + hidden_pre = F.linear( + x, + self.W_enc[layer].T, + self.b_enc[layer] + ) + thresh = torch.exp(self.log_threshold[layer]) + + feat_act = JumpReLU.apply(hidden_pre, thresh, self.bandwidth) + else: + raise ValueError(f"Unsupported activation_fn: {self.cfg.activation_fn}") return feat_act, hidden_pre def decode( @@ -305,36 +323,41 @@ def loss(self, act_in: torch.Tensor, act_out: torch.Tensor, l0_coef: float, df_c mse_loss_accross_layers = mse_loss_tensor.sum(dim=-1).mean(dim=0) mse_loss = mse_loss_accross_layers.sum() - if self.cfg.cross_layer_decoders: - squared_norms = (self.W_dec.float()**2).sum(dim=2) - feature_norms_local = torch.sqrt(torch.matmul(self.layer_mask.float(), squared_norms)) - else: - feature_norms_local = self.W_dec.float().norm(dim=2) - - # Compute L0 loss local - weighted_activations = feat_act.float() * feature_norms_local - tanh_weighted_activations = torch.tanh(C_l0_COEF * weighted_activations) - l0_loss_accross_layers = l0_coef * tanh_weighted_activations.sum(dim=-1).mean(dim=0) - l0_loss = l0_loss_accross_layers.sum().float() - - # SUM losses across ranks using autograd-aware all_reduce - if self.cfg.is_sharded: - l0_loss = all_reduce(l0_loss, op=dist.ReduceOp.SUM) - # l0_loss /= self.world_size - l0_loss_accross_layers = all_reduce(l0_loss_accross_layers, op=dist.ReduceOp.SUM) - # l0_loss_accross_layers /= self.world_size - - if self.cfg.debug: - self.log_loss_debug(feat_act, feature_norms_local, l0_loss) - - ### Dead feature penalty - dead_feature_loss = df_coef * torch.relu(torch.exp(self.log_threshold.float()) - hidden_pre.float()) * feature_norms_local - dead_feature_loss = dead_feature_loss.sum(dim=-1).mean(dim=0).sum() - - # SUM losses across ranks using autograd-aware all_reduce - if self.cfg.is_sharded: - dead_feature_loss = all_reduce(dead_feature_loss, op=dist.ReduceOp.SUM) - # dead_feature_loss /= self.world_size + if self.cfg.activation_fn in ["topk", "groupmax"]: + l0_loss = torch.tensor(0.0, device=act_in.device, dtype=torch.float32) + l0_loss_accross_layers = torch.zeros(self.N_layers, device=act_in.device, dtype=torch.float32) + dead_feature_loss = torch.tensor(0.0, device=act_in.device, dtype=torch.float32) + else: + if self.cfg.cross_layer_decoders: + squared_norms = (self.W_dec.float()**2).sum(dim=2) + feature_norms_local = torch.sqrt(torch.matmul(self.layer_mask.float(), squared_norms)) + else: + feature_norms_local = self.W_dec.float().norm(dim=2) + + # Compute L0 loss local + weighted_activations = feat_act.float() * feature_norms_local + tanh_weighted_activations = torch.tanh(C_l0_COEF * weighted_activations) + l0_loss_accross_layers = l0_coef * tanh_weighted_activations.sum(dim=-1).mean(dim=0) + l0_loss = l0_loss_accross_layers.sum().float() + + # SUM losses across ranks using autograd-aware all_reduce + if self.cfg.is_sharded: + l0_loss = all_reduce(l0_loss, op=dist.ReduceOp.SUM) + # l0_loss /= self.world_size + l0_loss_accross_layers = all_reduce(l0_loss_accross_layers, op=dist.ReduceOp.SUM) + # l0_loss_accross_layers /= self.world_size + + if self.cfg.debug: + self.log_loss_debug(feat_act, feature_norms_local, l0_loss) + + ### Dead feature penalty + dead_feature_loss = df_coef * torch.relu(torch.exp(self.log_threshold.float()) - hidden_pre.float()) * feature_norms_local + dead_feature_loss = dead_feature_loss.sum(dim=-1).mean(dim=0).sum() + + # SUM losses across ranks using autograd-aware all_reduce + if self.cfg.is_sharded: + dead_feature_loss = all_reduce(dead_feature_loss, op=dist.ReduceOp.SUM) + # dead_feature_loss /= self.world_size ### Dead feature count local with torch.no_grad(): diff --git a/src/clt_forge/clt_training_runner.py b/src/clt_forge/clt_training_runner.py index 9f34624..fbd4ab0 100644 --- a/src/clt_forge/clt_training_runner.py +++ b/src/clt_forge/clt_training_runner.py @@ -164,6 +164,7 @@ def run(self): trainer = CLTTrainer( clt=self.clt, activations_store=self.activations_store, + val_activations_store=getattr(self, "val_activations_store", None), save_checkpoint_fn=self.save_checkpoint, cfg=self.cfg, rank=self.rank, diff --git a/src/clt_forge/config/clt_config.py b/src/clt_forge/config/clt_config.py index 0fa1262..a922a8d 100644 --- a/src/clt_forge/config/clt_config.py +++ b/src/clt_forge/config/clt_config.py @@ -24,6 +24,8 @@ class CLTConfig(BaseModel): cross_layer_decoders: bool context_size: int functional_loss: Optional[str] = None + activation_fn: str = "jumprelu" + k: Optional[int] = None # -----Sparsity--------------------------- l0_coefficient: float diff --git a/src/clt_forge/config/clt_training_runner_config.py b/src/clt_forge/config/clt_training_runner_config.py index c80b5c0..f497fdd 100644 --- a/src/clt_forge/config/clt_training_runner_config.py +++ b/src/clt_forge/config/clt_training_runner_config.py @@ -38,6 +38,8 @@ class CLTTrainingRunnerConfig(BaseModel): jumprelu_init_threshold: float = 0.03 jumprelu_bandwidth: float = 1. normalize_decoder: bool = False + activation_fn: str = "jumprelu" + k: Optional[int] = None # -----ActivationStore Parameters--------- context_size: int = 32 @@ -81,6 +83,9 @@ class CLTTrainingRunnerConfig(BaseModel): # -----Metrics---------------------------- dead_feature_window: int = 250 # n_eval_batches: int = 10 + eval_step_size: Optional[int] = None + val_step_size: Optional[int] = None + tau_sweep_steps: int = 10 # -----WANDB------------------------------ log_to_wandb: bool = True diff --git a/src/clt_forge/training/clt_trainer.py b/src/clt_forge/training/clt_trainer.py index 7be97fb..4d8248d 100644 --- a/src/clt_forge/training/clt_trainer.py +++ b/src/clt_forge/training/clt_trainer.py @@ -30,6 +30,7 @@ def __init__( activations_store: ActivationsStore, cfg: CLTTrainingRunnerConfig, save_checkpoint_fn: Callable[["CLTTrainer", str], None], + val_activations_store: Optional[ActivationsStore] = None, rank: int = 0, world_size: int = 1 ) -> None: @@ -38,9 +39,13 @@ def __init__( self.is_main_process = rank == 0 self.clt = clt self.activations_store = activations_store + self.val_activations_store = val_activations_store self.cfg = cfg self.save_checkpoint_fn = save_checkpoint_fn self.n_training_steps: int = 0 + + # Evaluator instantiated dynamically to save memory if not evaluating + self.evaluator = None self.checkpoint_thresholds = [] if self.cfg.n_checkpoints > 0: @@ -74,7 +79,7 @@ def __init__( print(f"warm up {cfg.l0_warm_up_steps}") self.optimizer = Adam( - self.clt.parameters(), + [p for p in self.clt.parameters() if p.requires_grad], lr=cfg.lr, betas=( cfg.adam_beta1, @@ -101,7 +106,6 @@ def __init__( self.accumulation_step: int = 0 def _initialize_b_enc(self, n_batches: int = 5): - model = self._get_clt() def get_hidden_pre(acts_in): @@ -195,6 +199,14 @@ def fit(self): if self.accumulation_step == 0: self._log_train_step(loss_metrics) self._checkpoint_if_needed() + + # Check for evaluation + if self.is_main_process and self.cfg.eval_step_size is not None and self.n_training_steps > 0 and self.n_training_steps % self.cfg.eval_step_size == 0: + self._evaluate_and_log_metrics(prefix="train") + + # Check for validation + if self.is_main_process and self.cfg.val_step_size is not None and self.n_training_steps > 0 and self.n_training_steps % self.cfg.val_step_size == 0: + self._evaluate_and_log_metrics(prefix="val") # if self.cfg.functional_loss is not None and self.fc_scheduler.get_lr() > 0 and start_func_finetuning: # self._enable_functional_training() @@ -222,6 +234,9 @@ def fit(self): if self.accumulation_step == 0: self.n_training_steps += 1 + if self.is_main_process and self.cfg.eval_step_size is not None: + self._evaluate_and_log_metrics(prefix="train") + self.save_checkpoint_fn( trainer=self, checkpoint_name=f"final_{self.n_tokens}", @@ -604,6 +619,81 @@ def _update_layer_metrics_history( return log_dict + def _evaluate_and_log_metrics(self, prefix: str = "train"): + """Evaluates metrics via CLTEvaluator and logs to W&B.""" + if not self.cfg.log_to_wandb: + return + + import numpy as np + + if self.evaluator is None: + from clt_forge.training.eval_metrics import CLTEvaluator + self.evaluator = CLTEvaluator(self._get_clt(), self.cfg) + + store = self.val_activations_store if prefix == "val" and self.val_activations_store else self.activations_store + + # Get a batch of tokens + # The activations store __iter__ yields: `*tokens, acts_in, acts_out` + try: + batch = next(store.__iter__()) + tokens = batch[0] if len(batch) > 2 else None + except StopIteration: + logger.warning(f"Failed to get evaluation batch for {prefix}.") + return + + if tokens is None: + logger.warning(f"No tokens available in ActivationsStore for evaluation.") + return + + logger.info(f"Running evaluation ({prefix}) at step {self.n_training_steps}...") + + # Evaluate Replacement Score and KL Divergence + metrics_kl = self.evaluator.evaluate_replacement_and_kl(tokens) + + # Evaluate Pruning-time Sparsity Sweeps + tau_values = np.linspace(0, 1.0, self.cfg.tau_sweep_steps) + pruning_table = wandb.Table(columns=["tau", "pruned_l0_norm"]) + + base_sparsity_metrics = None + for tau in tau_values: + sparsity_metrics = self.evaluator.evaluate_sparsity_and_l0(tokens, tau=tau) + pruning_table.add_data(tau, sparsity_metrics["pruned_l0_norm"]) + if tau == 0.0: + base_sparsity_metrics = sparsity_metrics + + # Evaluate Activation Density + density_metrics = self.evaluator.evaluate_activation_density(tokens) + + # Log basic metrics + log_dict = { + f"{prefix}_eval/replacement_score": metrics_kl["replacement_score"], + f"{prefix}_eval/kl_divergence": metrics_kl["kl_divergence"], + f"{prefix}_eval/l0_norm": base_sparsity_metrics["l0_norm"], + f"{prefix}_eval/mean_activation_density": density_metrics["mean_activation_density"], + f"{prefix}_eval/dead_features_ratio": density_metrics["dead_features_ratio"], + } + + # 1. Pareto Frontier + pareto_table = wandb.Table(columns=["l0_norm", "replacement_score", "step"]) + pareto_table.add_data(base_sparsity_metrics["l0_norm"], metrics_kl["replacement_score"], self.n_training_steps) + log_dict[f"{prefix}_eval/Pareto_Frontier"] = wandb.plot.scatter(pareto_table, "l0_norm", "replacement_score", title=f"Pareto Frontier ({prefix})") + + # 3. Pruning Sweep + log_dict[f"{prefix}_eval/Pruning_Sweep"] = wandb.plot.line(pruning_table, "tau", "pruned_l0_norm", title=f"Pruning-time Sparsity Sweep ({prefix})") + + # 4. Dead Feature Histogram + density_vector = density_metrics["activation_density_per_feature"].flatten().cpu().numpy().tolist() + log_dict[f"{prefix}_eval/Activation_Density_Hist"] = wandb.Histogram(density_vector) + + # 5. Per-Layer Density Table for W&B Native Custom Charts + layer_density_table = wandb.Table(columns=["Layer", "Step", "Mean_Density"]) + density_per_layer = density_metrics["activation_density_per_feature"].mean(dim=1).cpu().numpy() + for l_idx, density in enumerate(density_per_layer): + layer_density_table.add_data(l_idx, self.n_training_steps, float(density)) + log_dict[f"{prefix}_eval/Layer_Density_Table"] = layer_density_table + + wandb.log(log_dict) + # def _enable_functional_training(self): # """Enable functional training by configuring activations store for token return.""" diff --git a/src/clt_forge/training/eval_metrics.py b/src/clt_forge/training/eval_metrics.py new file mode 100644 index 0000000..89c6278 --- /dev/null +++ b/src/clt_forge/training/eval_metrics.py @@ -0,0 +1,206 @@ +import torch +import torch.nn.functional as F +import numpy as np +from typing import List, Callable, Optional, Dict, Any +from transformer_lens import HookedTransformer + +class CLTEvaluator: + def __init__(self, clt, cfg): + """ + clt: The CLT model. + cfg: The CLTTrainingRunnerConfig instance. + """ + self.clt = clt + self.cfg = cfg + + from clt_forge import logger + logger.info(f"Instantiating base model {cfg.model_name} in CLTEvaluator") + self.base_model = HookedTransformer.from_pretrained( + cfg.model_name, + device=clt.device, + **(cfg.model_from_pretrained_kwargs if cfg.model_from_pretrained_kwargs else {}) + ) + self.base_model.eval() + + # Assuming residual stream hook names match layers directly + self.hook_names = [f"blocks.{i}.hook_resid_post" for i in range(clt.N_layers)] + + @torch.no_grad() + def evaluate_replacement_and_kl( + self, + tokens: torch.Tensor, + layers: Optional[List[int]] = None, + metric_fn: Optional[Callable] = None + ) -> Dict[str, float]: + """ + Calculates the Replacement Score and KL Divergence for the specified layers. + If layers is None, evaluates all layers. + """ + if layers is None: + layers = list(range(self.clt.N_layers)) + + if metric_fn is None: + def default_metric(logits, tokens): + shift_logits = logits[:, :-1, :].contiguous() + shift_labels = tokens[:, 1:].contiguous() + return F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)).item() + metric_fn = default_metric + + # 1. Clean Run + clean_logits = self.base_model(tokens, return_type="logits") + M_clean = metric_fn(clean_logits, tokens) + P_clean = F.log_softmax(clean_logits, dim=-1) + + # 2. Zero (Ablated) Run + def zero_hook_fn(act, hook): + return torch.zeros_like(act) + + zero_hooks = [(self.hook_names[i], zero_hook_fn) for i in layers] + zero_logits = self.base_model.run_with_hooks(tokens, return_type="logits", fwd_hooks=zero_hooks) + M_zero = metric_fn(zero_logits, tokens) + + # 3. Transcoder (Replaced) Run + replace_hooks = [] + + if self.clt.cfg.cross_layer_decoders: + # We need to hook all layers up to max(layers) to accumulate correctly + max_layer = max(layers) + layers_to_hook = list(range(max_layer + 1)) + else: + layers_to_hook = layers + + B = tokens.size(0) + seq_len = tokens.size(1) + recon_acc = torch.zeros(B * seq_len, self.clt.N_layers, self.clt.d_in, device=self.clt.device, dtype=self.clt.dtype) + + for layer_idx in layers_to_hook: + def get_replace_hook(l_idx, replace): + def replace_hook_fn(act, hook): + act_flat = act.view(-1, self.clt.d_in).to(self.clt.dtype) + + z_i, _ = self.clt.encode(act_flat, layer=l_idx) + out_i = self.clt.decode(z_i, layer=l_idx) + + if self.clt.cfg.cross_layer_decoders: + indices = (self.clt.l_idx == l_idx).nonzero(as_tuple=True)[0] + target_layers = self.clt.k_idx[indices] + + for idx, target_layer in enumerate(target_layers): + recon_acc[:, target_layer, :] += out_i[:, idx, :] + + recon = recon_acc[:, l_idx, :] + else: + recon = out_i + + if replace: + return recon.view(*act.shape).to(act.dtype) + else: + return act + return replace_hook_fn + + should_replace = layer_idx in layers + replace_hooks.append((self.hook_names[layer_idx], get_replace_hook(layer_idx, should_replace))) + + transcoder_logits = self.base_model.run_with_hooks(tokens, return_type="logits", fwd_hooks=replace_hooks) + M_transcoder = metric_fn(transcoder_logits, tokens) + P_transcoder = F.log_softmax(transcoder_logits, dim=-1) + + denom = (M_clean - M_zero) + if denom == 0: + replacement_score = float('nan') + else: + replacement_score = (M_transcoder - M_zero) / denom + + prob_clean = torch.exp(P_clean) + kl_div = F.kl_div(P_transcoder, prob_clean, reduction='batchmean', log_target=False).item() + + return { + "M_clean": M_clean, + "M_zero": M_zero, + "M_transcoder": M_transcoder, + "replacement_score": replacement_score, + "kl_divergence": kl_div + } + + @torch.no_grad() + def evaluate_sparsity_and_l0(self, tokens: torch.Tensor, layers: Optional[List[int]] = None, tau: float = 0.0) -> Dict[str, float]: + """ + Calculates the L0 Norm and Sparsity Comparisons. + """ + if layers is None: + layers = list(range(self.clt.N_layers)) + + _, cache = self.base_model.run_with_cache(tokens, names_filter=self.hook_names) + + B, seq_len = tokens.shape + d_in = self.clt.d_in + + act_in_list = [] + for name in self.hook_names: + if name in cache: + act_in_list.append(cache[name].view(-1, d_in)) + else: + raise ValueError(f"Hook name {name} not found in model cache.") + + act_in = torch.stack(act_in_list, dim=1).to(self.clt.dtype).to(self.clt.device) + + feat_act, _ = self.clt.encode(act_in) + + feat_act_selected = feat_act[:, layers, :] + + active_counts = (feat_act_selected > 0).float().sum(dim=-1) + l0_norm = active_counts.mean().item() + + active_counts_pruned = (feat_act_selected > tau).float().sum(dim=-1) + pruned_l0_norm = active_counts_pruned.mean().item() + + total_features = self.clt.local_d_latent + training_sparsity = 1.0 - (l0_norm / total_features) + pruning_sparsity = 1.0 - (pruned_l0_norm / total_features) + + return { + "l0_norm": l0_norm, + "pruned_l0_norm": pruned_l0_norm, + "training_sparsity": training_sparsity, + "pruning_sparsity": pruning_sparsity + } + + @torch.no_grad() + def evaluate_activation_density(self, tokens: torch.Tensor, layers: Optional[List[int]] = None) -> Dict[str, Any]: + """ + Calculates the activation density (proportion of tokens each feature is active for). + """ + if layers is None: + layers = list(range(self.clt.N_layers)) + + _, cache = self.base_model.run_with_cache(tokens, names_filter=self.hook_names) + + B, seq_len = tokens.shape + d_in = self.clt.d_in + + act_in_list = [] + for name in self.hook_names: + if name in cache: + act_in_list.append(cache[name].view(-1, d_in)) + else: + raise ValueError(f"Hook name {name} not found in model cache.") + + act_in = torch.stack(act_in_list, dim=1).to(self.clt.dtype).to(self.clt.device) + + feat_act, _ = self.clt.encode(act_in) + + feat_act_selected = feat_act[:, layers, :] + + density = (feat_act_selected > 0).float().mean(dim=0) + + mean_density = density.mean().item() + + dead_features = (density == 0).sum().item() + total_features = len(layers) * self.clt.local_d_latent + + return { + "activation_density_per_feature": density, + "mean_activation_density": mean_density, + "dead_features_count": dead_features, + "dead_features_ratio": dead_features / total_features + } diff --git a/src/clt_forge/training/optim.py b/src/clt_forge/training/optim.py index 5ada06e..4b7a212 100644 --- a/src/clt_forge/training/optim.py +++ b/src/clt_forge/training/optim.py @@ -1,6 +1,7 @@ import math import torch -from typing import Any +import torch.nn.functional as F +from typing import Any, Literal class LearningRateScheduler: def __init__( @@ -135,3 +136,112 @@ def backward( # type: ignore[override] dim=0, ) return x_grad, threshold_grad, None + + +class FusedEncoder(torch.autograd.Function): + @staticmethod + def forward( + ctx, input, weight, bias, k: int, activation: Literal["groupmax", "topk"] + ): + """ + input: (B, L, D) + weight: (L, D, M) + bias: (L, M) + k: int (number of top elements to select along dim=-1) + """ + preacts = torch.einsum("bld,ldm->blm", input, weight) + if bias is not None: + preacts = preacts + bias.unsqueeze(0) + preacts = F.relu(preacts) + + # Get top-k values and indices for each row + if activation == "topk": + values, indices = torch.topk(preacts, k, dim=-1, sorted=False) + elif activation == "groupmax": + values, indices = preacts.unflatten(-1, (k, -1)).max(dim=-1) + + # torch.max gives us indices into each group, but we want indices into the + # flattened tensor. Add the offsets to get the correct indices. + num_latents = preacts.shape[-1] + offsets = torch.arange( + 0, num_latents, num_latents // k, device=preacts.device + ) + indices = offsets + indices + else: + raise ValueError(f"Unknown activation: {activation}") + + ctx.save_for_backward(input, weight, bias, indices) + ctx.k = k + ctx.activation = activation + return values, indices, preacts + + @staticmethod + def backward(ctx, grad_values, grad_indices, grad_preacts): + input, weight, bias, indices = ctx.saved_tensors + + B, L, D = input.shape + _, _, M = weight.shape + + grad_input = grad_weight = grad_bias = None + + # --- Grad w.r.t. input --- + if ctx.needs_input_grad[0]: + grad_input = torch.zeros_like(input) + for l in range(L): + embedding_weight = weight[l].T # Shape: (M, D) + grad_input[:, l, :] = F.embedding_bag( + indices[:, l, :], + embedding_weight, + mode="sum", + per_sample_weights=grad_values[:, l, :].type_as(embedding_weight), + ) + + # --- Grad w.r.t. weight --- + if ctx.needs_input_grad[1]: + grad_weight = torch.zeros_like(weight) + for l in range(L): + # Compute contributions from each top-k element for layer l + contributions_l = grad_values[:, l, :].unsqueeze(2) * input[:, l, :].unsqueeze(1) # Shape: (B, k, D) + contributions_l = contributions_l.reshape(-1, D) # Shape: (B*k, D) + + grad_weight_l_T = torch.zeros(M, D, device=weight.device, dtype=weight.dtype) + grad_weight_l_T.index_add_(0, indices[:, l, :].flatten(), contributions_l.type_as(weight)) + grad_weight[l] = grad_weight_l_T.T + + # --- Grad w.r.t. bias --- + if bias is not None and ctx.needs_input_grad[2]: + grad_bias = torch.zeros_like(bias) + for l in range(L): + grad_bias[l].index_add_( + 0, indices[:, l, :].flatten(), grad_values[:, l, :].flatten().type_as(bias) + ) + + return grad_input, grad_weight, grad_bias, None, None + + +def fused_encoder( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + k: int, + activation: Literal["groupmax", "topk"], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Convenience wrapper that performs a multi-layer projection followed by `activation` + with a backward pass optimized using index_add and embedding_bag. + """ + is_2d = (input.dim() == 2) + if is_2d: + input = input.unsqueeze(1) # (B, 1, D) + weight = weight.unsqueeze(0) # (1, D, M) + if bias is not None: + bias = bias.unsqueeze(0) # (1, M) + + values, indices, preacts = FusedEncoder.apply(input, weight, bias, k, activation) + + if is_2d: + values = values.squeeze(1) + indices = indices.squeeze(1) + preacts = preacts.squeeze(1) + + return values, indices, preacts diff --git a/tests/test_eval_metrics.py b/tests/test_eval_metrics.py new file mode 100644 index 0000000..c174ec4 --- /dev/null +++ b/tests/test_eval_metrics.py @@ -0,0 +1,121 @@ +import pytest +import torch +import torch.nn as nn +from unittest.mock import Mock, patch + +from clt_forge.training.eval_metrics import CLTEvaluator + +class MockCLTConfig: + def __init__(self): + self.cross_layer_decoders = False + self.model_name = "mock_model" + self.model_from_pretrained_kwargs = {} + +class MockCLT: + def __init__(self, n_layers, d_in, d_latent): + self.N_layers = n_layers + self.d_in = d_in + self.d_latent = d_latent + self.local_d_latent = d_latent + self.device = torch.device('cpu') + self.dtype = torch.float32 + self.cfg = MockCLTConfig() + + def encode(self, act_in, layer=None): + # mock encode: just return act_in expanded to d_latent, and act_in + shape = list(act_in.shape) + shape[-1] = self.d_latent + feat_act = torch.ones(*shape, device=self.device, dtype=self.dtype) + return feat_act, act_in + + def decode(self, z, layer=None): + # mock decode: return z reduced to d_in + shape = list(z.shape) + shape[-1] = self.d_in + return torch.ones(*shape, device=self.device, dtype=self.dtype) * 0.5 + +class MockHookedTransformer(nn.Module): + def __init__(self, hook_names): + super().__init__() + self.hook_names = hook_names + + def forward(self, tokens, return_type="logits"): + # Returns fake logits + batch_size, seq_len = tokens.shape + # d_vocab = 100 + return torch.ones((batch_size, seq_len, 100)) + + def run_with_hooks(self, tokens, return_type="logits", fwd_hooks=[]): + # Execute hooks manually on dummy activations + batch_size, seq_len = tokens.shape + dummy_act = torch.ones((batch_size, seq_len, 12)) # d_in = 12 + for name, hook in fwd_hooks: + # Just call the hook + dummy_act = hook(dummy_act, Mock()) + return self.forward(tokens, return_type=return_type) + + def run_with_cache(self, tokens, names_filter=None): + batch_size, seq_len = tokens.shape + logits = self.forward(tokens) + cache = {name: torch.ones((batch_size, seq_len, 12)) for name in self.hook_names} + return logits, cache + + def eval(self): + pass + +@patch("clt_forge.training.eval_metrics.HookedTransformer.from_pretrained") +def test_evaluate_replacement_and_kl(mock_from_pretrained): + clt = MockCLT(n_layers=2, d_in=12, d_latent=24) + hook_names = ["blocks.0.hook_resid_post", "blocks.1.hook_resid_post"] + + # Setup mock + mock_model = MockHookedTransformer(hook_names) + mock_from_pretrained.return_value = mock_model + + evaluator = CLTEvaluator(clt, clt.cfg) + + tokens = torch.randint(0, 100, (2, 5)) # B=2, seq_len=5 + + metrics = evaluator.evaluate_replacement_and_kl(tokens) + + assert "replacement_score" in metrics + assert "kl_divergence" in metrics + assert "M_clean" in metrics + assert "M_zero" in metrics + assert "M_transcoder" in metrics + +@patch("clt_forge.training.eval_metrics.HookedTransformer.from_pretrained") +def test_evaluate_sparsity_and_l0(mock_from_pretrained): + clt = MockCLT(n_layers=2, d_in=12, d_latent=24) + hook_names = ["blocks.0.hook_resid_post", "blocks.1.hook_resid_post"] + + mock_model = MockHookedTransformer(hook_names) + mock_from_pretrained.return_value = mock_model + + evaluator = CLTEvaluator(clt, clt.cfg) + + tokens = torch.randint(0, 100, (2, 5)) + metrics = evaluator.evaluate_sparsity_and_l0(tokens, tau=0.5) + + assert "l0_norm" in metrics + assert "pruned_l0_norm" in metrics + assert "training_sparsity" in metrics + assert "pruning_sparsity" in metrics + +@patch("clt_forge.training.eval_metrics.HookedTransformer.from_pretrained") +def test_evaluate_activation_density(mock_from_pretrained): + clt = MockCLT(n_layers=2, d_in=12, d_latent=24) + hook_names = ["blocks.0.hook_resid_post", "blocks.1.hook_resid_post"] + + mock_model = MockHookedTransformer(hook_names) + mock_from_pretrained.return_value = mock_model + + evaluator = CLTEvaluator(clt, clt.cfg) + + tokens = torch.randint(0, 100, (2, 5)) + metrics = evaluator.evaluate_activation_density(tokens) + + assert "activation_density_per_feature" in metrics + assert "mean_activation_density" in metrics + assert "dead_features_count" in metrics + assert "dead_features_ratio" in metrics diff --git a/tests/test_topk.py b/tests/test_topk.py new file mode 100644 index 0000000..62ee121 --- /dev/null +++ b/tests/test_topk.py @@ -0,0 +1,121 @@ +import pytest +import torch +from clt_forge.config import CLTConfig +from clt_forge.clt import CLT + +def get_base_config(activation_fn="topk", k=4, cross_layer_decoders=True): + return CLTConfig( + device="cpu", + dtype="float32", + seed=42, + model_name="dummy", + d_in=16, + d_latent=32, + n_layers=2, + jumprelu_bandwidth=1.0, + jumprelu_init_threshold=0.03, + normalize_decoder=False, + dead_feature_window=250, + cross_layer_decoders=cross_layer_decoders, + context_size=32, + l0_coefficient=1e-3, + activation_fn=activation_fn, + k=k + ) + +@pytest.mark.parametrize("activation_fn", ["topk", "groupmax"]) +@pytest.mark.parametrize("cross_layer_decoders", [True, False]) +def test_topk_groupmax_sparsity(activation_fn, cross_layer_decoders): + k = 4 + cfg = get_base_config(activation_fn=activation_fn, k=k, cross_layer_decoders=cross_layer_decoders) + clt = CLT(cfg) + + # Check requires_grad is False for log_threshold + assert clt.log_threshold.requires_grad is False + + # Input tensor shape: [B, N_layers, d_in] + B = 5 + acts_in = torch.randn(B, cfg.n_layers, cfg.d_in) + + feat_act, hidden_pre = clt.encode(acts_in) + + assert feat_act.shape == (B, cfg.n_layers, cfg.d_latent) + + # Assert exactly k active features per token/layer + for b in range(B): + for l in range(cfg.n_layers): + active_count = (feat_act[b, l] > 0).sum().item() + assert active_count == k, f"Expected exactly {k} active features, got {active_count}" + + # For groupmax, assert exactly 1 active feature per group + if activation_fn == "groupmax": + group_size = cfg.d_latent // k + for b in range(B): + for l in range(cfg.n_layers): + for g in range(k): + group_acts = feat_act[b, l, g*group_size:(g+1)*group_size] + group_active_count = (group_acts > 0).sum().item() + assert group_active_count == 1, f"Expected exactly 1 active feature in group {g}, got {group_active_count}" + +@pytest.mark.parametrize("activation_fn", ["topk", "groupmax"]) +def test_topk_groupmax_loss_and_gradients(activation_fn): + cfg = get_base_config(activation_fn=activation_fn, k=4) + clt = CLT(cfg) + + B = 5 + acts_in = torch.randn(B, cfg.n_layers, cfg.d_in) + acts_out = torch.randn(B, cfg.n_layers, cfg.d_in) + + # Compute loss + loss_metrics = clt.loss(acts_in, acts_out, l0_coef=1e-3, df_coef=1e-5) + + # Assert L0 loss and dead feature loss are exactly 0 + assert loss_metrics.l0_loss.item() == 0.0 + assert loss_metrics.dead_feature_loss.item() == 0.0 + assert loss_metrics.mse_loss.item() > 0.0 + + total_loss = loss_metrics.mse_loss + loss_metrics.l0_loss + loss_metrics.dead_feature_loss + total_loss.backward() + + # Verify gradients flow to weights + assert clt.W_enc.grad is not None + assert clt.b_enc.grad is not None + assert clt.W_dec.grad is not None + + # Ensure no gradients flow to log_threshold + assert clt.log_threshold.grad is None or (clt.log_threshold.grad == 0).all() + +@pytest.mark.parametrize("activation_fn", ["topk", "groupmax"]) +def test_topk_groupmax_b_enc_initialization(activation_fn): + # Use larger d_latent to test in-bounds target_idx + cfg = CLTConfig( + device="cpu", + dtype="float32", + seed=42, + model_name="dummy", + d_in=16, + d_latent=16384, + n_layers=2, + jumprelu_bandwidth=1.0, + jumprelu_init_threshold=0.03, + normalize_decoder=False, + dead_feature_window=250, + cross_layer_decoders=True, + context_size=32, + l0_coefficient=1e-3, + activation_fn=activation_fn, + k=4 + ) + clt = CLT(cfg) + + # Assert initial b_enc is all zeros + assert (clt.b_enc == 0).all() + + B = 300 + # hidden_pre needs to be shape [B, N_layers, d_latent] + hidden_pre = torch.randn(B, cfg.n_layers, cfg.d_latent) + + clt._initialize_b_enc(hidden_pre) + + # Assert b_enc is no longer all zeros + assert not (clt.b_enc == 0).all()