feat: add Qwen3.5 MoE hybrid layer support (#187)

* feat: add Qwen3.5 MoE hybrid layer support

Qwen3.5 MoE uses GatedDeltaNet (linear attention) on some layers instead
of standard self-attention, causing abliteration to fail because
self_attn.o_proj doesn't exist on those layers.

Changes:
- Wrap self_attn.o_proj in suppress(Exception) and add linear_attn.out_proj
  as alternative attention out-projection for GatedDeltaNet layers
- Scan all layers in get_abliterable_components() instead of only layer 0,
  since hybrid models have different components on different layers
- Derive LoRA target_modules from actual named_modules() instead of
  splitting component keys, which fails when module names differ across
  layers (e.g. "o_proj" vs "out_proj")

Tested with Qwen3.5-397B-A17B (7/100 refusals, KL 0.2676).

Relates to #43

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Apply suggestion from @gemini-code-assist[bot]

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Philipp Emanuel Weidmann <pew@worldwidemann.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Matthias Stegner
2026-03-06 08:33:57 +01:00
committed by GitHub
parent 303ba9d978
commit 5e3c04c802
+27 -12
View File
@@ -161,14 +161,19 @@ class Model:
assert isinstance(self.model, PreTrainedModel) assert isinstance(self.model, PreTrainedModel)
# Always use LoRA adapters for abliteration (faster reload, no weight modification). # Always use LoRA adapters for abliteration (faster reload, no weight modification).
# We use the leaf names (e.g. "o_proj") as target modules. # Collect actual leaf module names from the model for LoRA targeting.
# This may cause LoRA adapters to be attached to unrelated modules (e.g. "conv.o_proj"), # This is more robust than splitting component keys (e.g. "attn.o_proj" -> "o_proj")
# but this is harmless as we only abliterate the modules we target in `abliterate()`, # because hybrid models like Qwen3.5 MoE have modules with different names
# leaving the others at their default (identity) state. # across layers (e.g. "o_proj" on attention layers, "out_proj" on linear attention layers).
# NOTE: This will need to be updated when hybrid layer support (#43) is merged. target_modules_set: set[str] = set()
target_modules = [ layers = self.get_layers()
comp.split(".")[-1] for comp in self.get_abliterable_components() for layer_index, layer in enumerate(layers):
] module_id_to_leaf_name = {id(m): name.split(".")[-1] for name, m in layer.named_modules()}
for modules_list in self.get_layer_modules(layer_index).values():
for mod in modules_list:
if id(mod) in module_id_to_leaf_name:
target_modules_set.add(module_id_to_leaf_name[id(mod)])
target_modules = list(target_modules_set)
if self.settings.row_normalization != RowNormalization.FULL: if self.settings.row_normalization != RowNormalization.FULL:
# Rank 1 is sufficient for directional ablation without renormalization. # Rank 1 is sufficient for directional ablation without renormalization.
@@ -340,9 +345,14 @@ class Model:
f"Unexpected Tensor in {component} - expected nn.Module" f"Unexpected Tensor in {component} - expected nn.Module"
) )
# Exceptions aren't suppressed here, because there is currently # Standard self-attention out-projection (most models).
# no alternative location for the attention out-projection. 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 standard self-attention, so self_attn.o_proj doesn't exist on those layers.
with suppress(Exception):
try_add("attn.o_proj", layer.linear_attn.out_proj) # ty:ignore[possibly-missing-attribute]
# Most dense models. # Most dense models.
with suppress(Exception): with suppress(Exception):
@@ -374,7 +384,12 @@ class Model:
return modules return modules
def get_abliterable_components(self) -> list[str]: def get_abliterable_components(self) -> list[str]:
return list(self.get_layer_modules(0).keys()) # 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: set[str] = set()
for layer_index in range(len(self.get_layers())):
components.update(self.get_layer_modules(layer_index).keys())
return sorted(components)
def abliterate( def abliterate(
self, self,