diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index dedeb55..8caffb6 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -1033,6 +1033,10 @@ class AbliterationPipeline: self._routing_harmful: dict[int, list[torch.Tensor]] = {} self._routing_harmless: dict[int, list[torch.Tensor]] = {} self._routing_is_harmful: bool = True # flag for routing hooks + # Attention mask for the synchronous forward pass currently observed + # by router hooks. Flattened MoE routers need this shape information + # to recover the original (batch, sequence) prompt association. + self._routing_attention_mask: torch.Tensor | None = None def log(self, msg: str): self._on_log(msg) @@ -1126,17 +1130,39 @@ class AbliterationPipeline: def make_hook(layer_idx: int): def hook_fn(module, input, output): - logits = output if isinstance(output, torch.Tensor) else output[0] - # Extract router logits — use mean across positions for - # CoT-aware models so we capture expert routing at reasoning - # tokens, not just the final output token. - if logits.dim() == 3: - if getattr(self, "cot_aware", False) and logits.shape[1] > 4: - logits = logits.mean(dim=1) # (batch, num_experts) - else: - logits = logits[:, -1, :] # (batch, num_experts) - elif logits.dim() == 2 and logits.shape[0] > 1: - logits = logits[-1:, :] + if isinstance(output, torch.Tensor): + logits = output + elif ( + isinstance(output, (tuple, list)) + and output + and isinstance(output[0], torch.Tensor) + ): + logits = output[0] + else: + warnings.warn( + f"Skipping router profiling for layer {layer_idx}: " + "router output does not contain a tensor", + RuntimeWarning, + stacklevel=2, + ) + return + router_input = ( + input[0] if input and isinstance(input[0], torch.Tensor) else None + ) + try: + logits = self._router_logits_per_prompt( + logits, + router_input=router_input, + attention_mask=self._routing_attention_mask, + cot_aware=getattr(self, "cot_aware", False), + ) + except ValueError as exc: + warnings.warn( + f"Skipping router profiling for layer {layer_idx}: {exc}", + RuntimeWarning, + stacklevel=2, + ) + return target = (self._routing_harmful if self._routing_is_harmful else self._routing_harmless) @@ -1156,6 +1182,86 @@ class AbliterationPipeline: self.log(f" Router profiling hooks installed on {len(hooks)} MoE layers") return hooks + @staticmethod + def _router_logits_per_prompt( + logits: torch.Tensor, + *, + router_input: torch.Tensor | None, + attention_mask: torch.Tensor | None, + cot_aware: bool, + ) -> torch.Tensor: + """Return one router-logit vector per prompt without mixing batches. + + Some MoE implementations expose router logits as ``(batch, sequence, + experts)`` while others flatten the first two dimensions. A flattened + tensor is reshaped only when the active attention mask or a 3-D router + input proves the original dimensions. Ambiguous layouts are rejected + instead of silently treating a token as an entire prompt. + """ + if not isinstance(logits, torch.Tensor): + raise ValueError("router output does not contain a tensor") + if logits.dim() == 1: + return logits.unsqueeze(0) + if logits.dim() not in (2, 3): + raise ValueError( + f"unsupported router-logit shape {tuple(logits.shape)}; " + "expected 1-D, 2-D, or 3-D" + ) + + mask = attention_mask + if mask is not None: + if mask.dim() != 2: + raise ValueError( + f"attention mask has shape {tuple(mask.shape)}; " + "expected (batch, sequence)" + ) + mask = mask.to(device=logits.device, dtype=torch.bool) + + if logits.dim() == 2: + rows, experts = logits.shape + batch_seq: tuple[int, int] | None = None + if mask is not None: + batch_seq = (mask.shape[0], mask.shape[1]) + elif router_input is not None and router_input.dim() == 3: + batch_seq = (router_input.shape[0], router_input.shape[1]) + elif rows == 1: + return logits + else: + raise ValueError( + "multi-row 2-D router logits are ambiguous without a 2-D " + "attention mask or 3-D router input" + ) + + batch, sequence = batch_seq + if rows != batch * sequence: + raise ValueError( + f"router-logit rows ({rows}) do not match batch × sequence " + f"({batch} × {sequence})" + ) + logits = logits.reshape(batch, sequence, experts) + + batch, sequence, _ = logits.shape + if mask is None: + mask = torch.ones((batch, sequence), dtype=torch.bool, device=logits.device) + elif tuple(mask.shape) != (batch, sequence): + raise ValueError( + f"attention mask shape {tuple(mask.shape)} does not match router logits " + f"{(batch, sequence)}" + ) + if (~mask.any(dim=1)).any(): + raise ValueError("attention mask contains a prompt with no valid tokens") + + positions = torch.arange(sequence, device=logits.device).expand(batch, -1) + last_positions = positions.masked_fill(~mask, -1).max(dim=1).values + batch_positions = torch.arange(batch, device=logits.device) + last_logits = logits[batch_positions, last_positions] + if not cot_aware: + return last_logits + + valid_counts = mask.sum(dim=1) + means = (logits * mask.unsqueeze(-1)).sum(dim=1) / valid_counts.unsqueeze(-1) + return torch.where((valid_counts > 4).unsqueeze(-1), means, last_logits) + def run(self) -> Path: """Execute the full abliteration pipeline. Returns path to saved model.""" # Remove any steering hooks left from a previous run() call @@ -1736,13 +1842,18 @@ class AbliterationPipeline: max_length=max_length, ) inputs = {k: v.to(device) for k, v in inputs.items()} - with torch.no_grad(): - model(**inputs) + self._routing_attention_mask = inputs.get("attention_mask") + try: + with torch.no_grad(): + model(**inputs) + finally: + self._routing_attention_mask = None del inputs # Free GPU memory every few batches, not every prompt if (batch_end % (batch_size * 4) == 0) or batch_end == len(prompts): self._free_gpu_memory() finally: + self._routing_attention_mask = None tokenizer.padding_side = orig_padding_side for h in hooks: h.remove() diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py index 28f91c8..ef75506 100644 --- a/tests/test_abliterate.py +++ b/tests/test_abliterate.py @@ -2186,21 +2186,152 @@ class TestRouterProfilingHooks: # Simulate harmful forward pass pipeline._routing_is_harmful = True x = torch.randn(1, 5, 16) + harmful_expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[0, -1] layer.mlp.gate(x) # triggers hook assert len(pipeline._routing_harmful[0]) == 1 assert pipeline._routing_harmful[0][0].shape[0] == 4 # n_experts + assert torch.equal(pipeline._routing_harmful[0][0], harmful_expected) # Simulate harmless forward pass pipeline._routing_is_harmful = False - layer.mlp.gate(x) + harmless_x = x + 1 + harmless_expected = torch.nn.functional.linear( + harmless_x, layer.mlp.gate.weight, + )[0, -1] + layer.mlp.gate(harmless_x) assert len(pipeline._routing_harmless[0]) == 1 + assert torch.equal(pipeline._routing_harmless[0][0], harmless_expected) + assert not torch.equal( + pipeline._routing_harmful[0][0], pipeline._routing_harmless[0][0], + ) finally: for h in hooks: h.remove() abl_module.get_ffn_module = orig_get_ffn + def test_hooks_record_single_prompt_vector_output(self): + """A one-dimensional single-prompt router output should remain intact.""" + pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers() + hooks = [] + try: + hooks = pipeline._install_router_profiling_hooks(layers) + x = torch.arange(16, dtype=torch.float32) / 100 + expected = torch.nn.functional.linear(x, layer.mlp.gate.weight) + + layer.mlp.gate(x) + + assert len(pipeline._routing_harmful[0]) == 1 + assert torch.equal(pipeline._routing_harmful[0][0], expected) + finally: + for h in hooks: + h.remove() + abl_module.get_ffn_module = orig_get_ffn + + def test_hooks_record_one_vector_for_each_3d_batch_prompt(self): + """Native 3-D router outputs should retain every batch member.""" + pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers() + hooks = [] + try: + hooks = pipeline._install_router_profiling_hooks(layers) + x = torch.arange(2 * 5 * 16, dtype=torch.float32).reshape(2, 5, 16) / 100 + expected = torch.nn.functional.linear(x, layer.mlp.gate.weight) + + layer.mlp.gate(x) + + recorded = pipeline._routing_harmful[0] + assert len(recorded) == 2 + assert torch.equal(recorded[0], expected[0, -1]) + assert torch.equal(recorded[1], expected[1, -1]) + finally: + for h in hooks: + h.remove() + abl_module.get_ffn_module = orig_get_ffn + + def test_hooks_preserve_flattened_prompt_association_with_padding(self): + """Flattened token rows should become one vector per padded prompt.""" + pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers() + hooks = [] + try: + hooks = pipeline._install_router_profiling_hooks(layers) + pipeline._routing_attention_mask = torch.tensor( + [[1, 1, 0, 0], [1, 1, 1, 1]], dtype=torch.long, + ) + x = torch.arange(8 * 16, dtype=torch.float32).reshape(8, 16) / 100 + expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 4, 4) + + layer.mlp.gate(x) + + recorded = pipeline._routing_harmful[0] + assert len(recorded) == 2 + assert torch.equal(recorded[0], expected[0, 1]) + assert torch.equal(recorded[1], expected[1, 3]) + finally: + pipeline._routing_attention_mask = None + for h in hooks: + h.remove() + abl_module.get_ffn_module = orig_get_ffn + + def test_hooks_average_only_valid_cot_tokens_per_prompt(self): + """CoT aggregation should exclude padding without mixing prompts.""" + pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers() + hooks = [] + try: + hooks = pipeline._install_router_profiling_hooks(layers) + pipeline.cot_aware = True + pipeline._routing_attention_mask = torch.tensor( + [[1, 1, 1, 1, 1, 0], [0, 1, 1, 1, 1, 1]], dtype=torch.long, + ) + x = torch.arange(12 * 16, dtype=torch.float32).reshape(12, 16) / 100 + expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 6, 4) + + layer.mlp.gate(x) + + recorded = pipeline._routing_harmful[0] + assert len(recorded) == 2 + assert torch.allclose(recorded[0], expected[0, :5].mean(dim=0)) + assert torch.allclose(recorded[1], expected[1, 1:].mean(dim=0)) + finally: + pipeline._routing_attention_mask = None + for h in hooks: + h.remove() + abl_module.get_ffn_module = orig_get_ffn + + def test_hooks_skip_ambiguous_flattened_logits(self): + """Ambiguous flattened layouts should fail visibly instead of collapsing.""" + pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers() + hooks = [] + try: + hooks = pipeline._install_router_profiling_hooks(layers) + + with pytest.warns(RuntimeWarning, match="multi-row 2-D router logits are ambiguous"): + layer.mlp.gate(torch.randn(5, 16)) + + assert pipeline._routing_harmful[0] == [] + finally: + for h in hooks: + h.remove() + abl_module.get_ffn_module = orig_get_ffn + + def test_hooks_skip_mismatched_flattened_logits(self): + """A stale or incompatible mask should produce an actionable warning.""" + pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers() + hooks = [] + try: + hooks = pipeline._install_router_profiling_hooks(layers) + pipeline._routing_attention_mask = torch.ones(2, 4, dtype=torch.long) + + with pytest.warns(RuntimeWarning, match="do not match batch × sequence"): + layer.mlp.gate(torch.randn(7, 16)) + + assert pipeline._routing_harmful[0] == [] + finally: + pipeline._routing_attention_mask = None + for h in hooks: + h.remove() + abl_module.get_ffn_module = orig_get_ffn + def test_no_handle_returns_empty(self): """Should return empty list when handle is None.""" pipeline = AbliterationPipeline(model_name="test", method="surgical")