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/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..b581cb0 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 diff --git a/src/clt_forge/training/clt_trainer.py b/src/clt_forge/training/clt_trainer.py index 7be97fb..0ff471a 100644 --- a/src/clt_forge/training/clt_trainer.py +++ b/src/clt_forge/training/clt_trainer.py @@ -74,7 +74,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 +101,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): 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_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()