3 changed files with 13 additions and 21 deletions
+3 -8
View File
@@ -17,15 +17,11 @@ def _is_help_invocation() -> bool:
if _is_help_invocation(): if _is_help_invocation():
Settings() # ty:ignore[missing-argument] Settings() # ty:ignore[missing-argument]
# FIXME: Rich progress bars are currently disabled because of rendering issues
# when used from multiple threads in parallel (e.g. by huggingface_hub).
"""
from .progress import patch_tqdm from .progress import patch_tqdm
# This patches tqdm class definitions, which must happen # This patches tqdm class definitions, which must happen
# before any other module imports tqdm. # before any other module imports tqdm.
patch_tqdm() patch_tqdm()
"""
import logging import logging
import math import math
@@ -429,6 +425,9 @@ def run():
needs_full_residuals = settings.print_residual_geometry or settings.plot_residuals needs_full_residuals = settings.print_residual_geometry or settings.plot_residuals
good_residuals = None
bad_residuals = None
if needs_full_residuals: if needs_full_residuals:
print("* Obtaining residuals for good prompts...") print("* Obtaining residuals for good prompts...")
good_residuals = model.get_residuals_batched(good_prompts) good_residuals = model.get_residuals_batched(good_prompts)
@@ -466,12 +465,8 @@ def run():
refusal_directions - projection_vector.unsqueeze(1) * good_directions refusal_directions - projection_vector.unsqueeze(1) * good_directions
) )
refusal_directions = F.normalize(refusal_directions, p=2, dim=1) refusal_directions = F.normalize(refusal_directions, p=2, dim=1)
del good_directions, projection_vector
del good_means, bad_means
# Clear cache before starting the optimization study. # Clear cache before starting the optimization study.
# This should free up memory from the objects released with the del statements above.
empty_cache() empty_cache()
trial_index = 0 trial_index = 0
+6 -9
View File
@@ -154,15 +154,13 @@ class Model:
# so we don't need to do anything manually. # so we don't need to do anything manually.
print(f"* Transformer model with [bold]{len(self.get_layers())}[/] layers") print(f"* Transformer model with [bold]{len(self.get_layers())}[/] layers")
print("* Abliterable components:")
all_components = {} all_components = {}
for layer_index in range(len(self.get_layers())): for layer_index in range(len(self.get_layers())):
for component, modules in self.get_layer_modules(layer_index).items(): for component, modules in self.get_layer_modules(layer_index).items():
if component not in all_components: if component not in all_components:
all_components[component] = 0 all_components[component] = 0
all_components[component] += len(modules) all_components[component] += len(modules)
print("* Abliterable components:")
for component, count in all_components.items(): for component, count in all_components.items():
print(f" * [bold]{component}[/]: [bold]{count}[/] modules total") print(f" * [bold]{component}[/]: [bold]{count}[/] modules total")
@@ -370,8 +368,8 @@ class Model:
with suppress(Exception): with suppress(Exception):
try_add("attn.o_proj", layer.self_attn.o_proj) # ty:ignore[possibly-missing-attribute] try_add("attn.o_proj", layer.self_attn.o_proj) # ty:ignore[possibly-missing-attribute]
# Qwen3.5 MoE hybrid layers use GatedDeltaNet (linear attention) instead of # Qwen3.5 MoE hybrid layers use GatedDeltaNet (linear attention) instead
# standard self-attention, so self_attn.o_proj doesn't exist on those layers. # of standard self-attention, so self_attn.o_proj doesn't exist on those layers.
with suppress(Exception): with suppress(Exception):
try_add("attn.o_proj", layer.linear_attn.out_proj) # ty:ignore[possibly-missing-attribute] try_add("attn.o_proj", layer.linear_attn.out_proj) # ty:ignore[possibly-missing-attribute]
@@ -405,13 +403,11 @@ class Model:
return modules return modules
def get_abliterable_components(self) -> list[str]: def get_abliterable_components(self) -> list[str]:
components: set[str] = set()
# Scan all layers because hybrid models (e.g. Qwen3.5 MoE) have different # Scan all layers because hybrid models (e.g. Qwen3.5 MoE) have different
# components on different layers (some have self_attn, others linear_attn). # components on different layers (some have self_attn, others linear_attn).
components: set[str] = set()
for layer_index in range(len(self.get_layers())): for layer_index in range(len(self.get_layers())):
components.update(self.get_layer_modules(layer_index).keys()) components.update(self.get_layer_modules(layer_index).keys())
return sorted(components) return sorted(components)
def abliterate( def abliterate(
@@ -748,8 +744,9 @@ class Model:
# The returned tensor has shape (prompt, token). # The returned tensor has shape (prompt, token).
logprobs = F.log_softmax(logits, dim=-1) logprobs = F.log_softmax(logits, dim=-1)
del outputs
if self.settings.offload_outputs_to_cpu: if self.settings.offload_outputs_to_cpu:
del outputs, logits
logprobs = logprobs.cpu() logprobs = logprobs.cpu()
empty_cache() empty_cache()
Generated
+4 -4
View File
@@ -8,7 +8,7 @@ resolution-markers = [
] ]
[options] [options]
exclude-newer = "2026-04-16T07:30:09.771407348Z" exclude-newer = "2026-04-14T22:48:57.86057843Z"
exclude-newer-span = "P7D" exclude-newer-span = "P7D"
[[package]] [[package]]
@@ -1508,14 +1508,14 @@ wheels = [
[[package]] [[package]]
name = "mako" name = "mako"
version = "1.3.11" version = "1.3.10"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
dependencies = [ dependencies = [
{ name = "markupsafe" }, { name = "markupsafe" },
] ]
sdist = { url = "https://files.pythonhosted.org/packages/59/8a/805404d0c0b9f3d7a326475ca008db57aea9c5c9f2e1e39ed0faa335571c/mako-1.3.11.tar.gz", hash = "sha256:071eb4ab4c5010443152255d77db7faa6ce5916f35226eb02dc34479b6858069", size = 399811, upload-time = "2026-04-14T20:19:51.493Z" } sdist = { url = "https://files.pythonhosted.org/packages/9e/38/bd5b78a920a64d708fe6bc8e0a2c075e1389d53bef8413725c63ba041535/mako-1.3.10.tar.gz", hash = "sha256:99579a6f39583fa7e5630a28c3c1f440e4e97a414b80372649c0ce338da2ea28", size = 392474, upload-time = "2025-04-10T12:44:31.16Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/68/a5/19d7aaa7e433713ffe881df33705925a196afb9532efc8475d26593921a6/mako-1.3.11-py3-none-any.whl", hash = "sha256:e372c6e333cf004aa736a15f425087ec977e1fcbd2966aae7f17c8dc1da27a77", size = 78503, upload-time = "2026-04-14T20:19:53.233Z" }, { url = "https://files.pythonhosted.org/packages/87/fb/99f81ac72ae23375f22b7afdb7642aba97c00a713c217124420147681a2f/mako-1.3.10-py3-none-any.whl", hash = "sha256:baef24a52fc4fc514a0887ac600f9f1cff3d82c61d4d700a1fa84d597b88db59", size = 78509, upload-time = "2025-04-10T12:50:53.297Z" },
] ]
[[package]] [[package]]