From c0d73eb28886bf2119856a9e66637a8eb532906e Mon Sep 17 00:00:00 2001 From: Hinata Date: Sun, 11 Oct 2026 03:47:14 +0000 Subject: [PATCH 102/103] qwen4_exp: add QWEN38_ABLITERATION runtime refusal-direction removal --- vllm/models/qwen4_exp/nvidia/abliteration.py | 88 ++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model.py | 7 +- vllm/models/qwen4_exp/nvidia/mtp.py | 4 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 3 + 4 files changed, 100 insertions(+), 2 deletions(-) create mode 100644 vllm/models/qwen4_exp/nvidia/abliteration.py diff --git a/vllm/models/qwen4_exp/nvidia/abliteration.py b/vllm/models/qwen4_exp/nvidia/abliteration.py new file mode 100644 index 0000000..840168b --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/abliteration.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional refusal-direction removal ("abliteration") applied at runtime. + +``QWEN38_ABLITERATION`` names a JSON file with a ``refusal_direction`` list, or a safetensors file with a +``refusal_direction`` tensor, of length ``hidden_size``. A bare name such as ``Qwen3.8-Flash-Next-Uncensored`` resolves to +``/opt/qwen38/abliteration/.json``. +With it set, every residual-stream write loses its component along the unit direction r: + + x <- x - r (r . x) + +applied to the attention and linear-attention block outputs, the MoE block output (routed plus shared experts), +the PLE value projection and the text token embeddings. For one direction this equals the weight edit +``W' = W - r (r^T W)`` on ``self_attn.o_proj``, ``linear_attn.out_proj``, every expert's and the shared expert's +``down_proj``, ``ple.value_proj`` and the rows of ``embed_tokens`` (Arditi et al., 2024): the routed and shared +expert outputs are summed with scalar weights before the residual write, so projecting their sum is the same as +editing each down projection. The INT4 experts and every other weight stay untouched. Unset, nothing is loaded and +no work is added. +""" + +import json +import math +import os + +import torch + +ENV = "QWEN38_ABLITERATION" +BUILTIN_DIR = "/opt/qwen38/abliteration" + + +def direction_path() -> str | None: + value = os.environ.get(ENV, "").strip() + if value in ("", "0"): + return None + if os.sep not in value and not value.endswith((".json", ".safetensors")): + value = os.path.join(BUILTIN_DIR, f"{value}.json") + return value + + +def load_direction(hidden_size: int) -> torch.Tensor | None: + """The unit refusal direction as a float32 CPU tensor, or None when the switch is off.""" + path = direction_path() + if path is None: + return None + if path.endswith(".json"): + with open(path) as f: + r = torch.tensor(json.load(f)["refusal_direction"], dtype=torch.float32) + else: + from safetensors import safe_open + + with safe_open(path, framework="pt", device="cpu") as f: + r = f.get_tensor("refusal_direction") + r = r.to(torch.float32).flatten() + if r.numel() != hidden_size: + raise ValueError(f"{ENV}: {path} holds {r.numel()} values; the model's hidden size is {hidden_size}") + norm = float(r.norm()) + if not math.isfinite(norm) or norm == 0.0: + raise ValueError(f"{ENV}: {path} holds a zero or non-finite direction") + return r / norm + + +def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None: + """Attach the direction as a non-persistent buffer ``_abliteration_r`` (None when the switch is off). + + Call inside the model constructor so the buffer lands on the model's device before CUDA-graph capture. The + CPU-only PLE worker process never runs these forwards and gets no buffer. + """ + from vllm.model_executor.layers.ple_offload_layer import is_offload_process + + r = None if is_offload_process() else load_direction(hidden_size) + if r is None: + module._abliteration_r = None + return + buf = torch.empty(hidden_size, dtype=dtype) + buf.copy_(r) + module.register_buffer("_abliteration_r", buf, persistent=False) + + +def project_(x: torch.Tensor, r: torch.Tensor | None) -> torch.Tensor: + """In place: remove the component along r from x's last dimension. No-op when r is None.""" + if r is None: + return x + width = x.shape[-1] + if x.is_contiguous(): + flat = x.view(-1, width) + flat.addr_(flat @ r, r, alpha=-1.0) # rank-1 update, no full-size temporary + else: + x.sub_(torch.outer(x.reshape(-1, width) @ r, r).view(x.shape)) + return x diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index fe416d4..b4fe373 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -89,6 +89,7 @@ from vllm.v1.e2_guards import ( from vllm.v1.kv_cache_interface import MambaSpec from ..config import Qwen4ExpConfig +from . import abliteration from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .ple_layer import Qwen4ExpPLELayer @@ -384,6 +385,7 @@ class Qwen4ExpDecoderLayer(nn.Module): hc_config, prefix=maybe_prefix(prefix, "mlp_hyper_connection"), ) + abliteration.register(self, config.hidden_size, model_config.dtype) def forward( self, @@ -434,12 +436,14 @@ class Qwen4ExpDecoderLayer(nn.Module): ) else: raise ValueError("Invalid layer_type") + abliteration.project_(attn_out, self._abliteration_r) mlp_hc = self.mlp_hyper_connection hidden_states, block_input, injection = mlp_hc.combine_and_mix( hidden_states, attn_out, injection ) mlp_out = self.mlp(block_input) + abliteration.project_(mlp_out, self._abliteration_r) return hidden_states, mlp_out, injection @@ -554,6 +558,7 @@ class Qwen4ExpModel(nn.Module): self.embed_tokens = embedding_cls(self.vocab_size, config.hidden_size) else: self.embed_tokens = PPMissingLayer() + abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype) def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: layer_idx = extract_layer_index(prefix) @@ -615,7 +620,7 @@ class Qwen4ExpModel(nn.Module): self._mtp_hidden_buffer = None def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) + return abliteration.project_(self.embed_tokens(input_ids), self._abliteration_r) @staticmethod def _start_layer_ple_prefetch( diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 01fb69c..39a0ff4 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -60,6 +60,7 @@ from vllm.v1.e2_guards import ( assert_ple_mode, ) +from . import abliteration from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .model import ( @@ -220,6 +221,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): else VocabParallelEmbedding ) self.embed_tokens = embedding_cls(self.vocab_size, self.hidden_size) + abliteration.register(self, self.hidden_size, model_config.dtype) draft_vllm_config = _make_draft_vllm_config( vllm_config, self.mtp_start_layer_idx, @@ -312,7 +314,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): buffer[:num_rows].copy_(selected) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) + return abliteration.project_(self.embed_tokens(input_ids), self._abliteration_r) def forward( self, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index af8f528..443e9dd 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -28,6 +28,7 @@ from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionMetadata, ) +from . import abliteration from .ngram_embedding import Qwen4ExpNGramEmbedding from .ops.ple import ple_conv, ple_gate @@ -136,6 +137,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): prefix=f"{prefix}.kv_proj", disable_tp=True, ) + abliteration.register(self, self.hidden_size, model_config.dtype) norm_args = ( self.hc_hidden_size, config.rms_norm_eps, @@ -449,6 +451,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) kv, _ = self.kv_proj(embeddings) key, value = kv.split(self.kv_proj.output_sizes, dim=-1) + abliteration.project_(value, self._abliteration_r) gated_output, conv_input = ple_gate( key, value, -- 2.43.0