From 0076adc78b27eb61ed7fe71fd0a22c32948f8494 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Thu, 15 Oct 2026 21:40:16 +0700 Subject: [PATCH] [Qwen3.5] Project the refusal direction out of the residual stream at runtime orcarouter/Qwen3.8-27B-Uncensored removes refusal by editing weights. The Swift release we serve is AutoRound W8A16, so re-editing weights means re-quantizing. Project one unit direction r out of every residual-stream writer output instead: x <- x - r (r . x), which equals the weight edit W' = W - r(r^T W) on every writer feeding that stream (Arditi et al. 2024). Sites: attention output for both the full-attention and GDN branches, the MLP output, and token embeddings, in Qwen3NextDecoderLayer.forward and Qwen3NextModel.embed_input_ids, which Qwen3_5DecoderLayer and Qwen3_5Model inherit but whose base constructors they skip, so both register there too. The MTP drafter reuses Qwen3_5DecoderLayer and calls its own embed_input_ids, so the draft path is covered at the same width, and _maybe_share_embeddings at PP=1 hands the draft the target embedding module. Sequence parallelism shards dim 0 only, so the last dim is hidden_size at every site. VLLM_ABLITERATION names a JSON file with key refusal_direction, a safetensors tensor, or a bare name resolved under /opt/qwen38/abliteration then /data/abliteration. Empty or 0 is off. register() always creates the buffer _abliteration_r, zeros when off, so one compiled graph serves ON and OFF and toggling the switch (declared in envs.py, hence in compile_factors) invalidates the compile cache. project_() is an in-place rank-1 update with no full-size temporary, and skips tensors whose width is not hidden_size. Co-authored-by: Hinata Signed-off-by: MOVZX --- tests/models/test_abliteration.py | 245 +++++++++++++++++++++ vllm/envs.py | 10 + vllm/model_executor/models/abliteration.py | 151 +++++++++++++ vllm/model_executor/models/qwen3_5.py | 3 + vllm/model_executor/models/qwen3_5_mtp.py | 4 +- vllm/model_executor/models/qwen3_next.py | 9 +- 6 files changed, 420 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_abliteration.py create mode 100644 vllm/model_executor/models/abliteration.py diff --git a/tests/models/test_abliteration.py b/tests/models/test_abliteration.py new file mode 100644 index 0000000000..5e9435ee9a --- /dev/null +++ b/tests/models/test_abliteration.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the runtime refusal-direction projection.""" + +import json +import math + +import pytest +import torch + +from vllm import envs +from vllm.model_executor.models import abliteration + +HIDDEN = 8 + + +@pytest.fixture(autouse=True) +def clear_direction_cache(): + abliteration._read.cache_clear() + yield + abliteration._read.cache_clear() + + +def write_direction(tmp_path, name="r.json", size=HIDDEN): + direction = [math.sin(i) for i in range(size)] + path = tmp_path / name + path.write_text(json.dumps({"hidden_size": size, "refusal_direction": direction})) + return str(path) + + +def registered(module, size=HIDDEN, dtype=torch.float32): + abliteration.register(module, size, dtype) + return module + + +def test_off_is_a_noop_and_still_registers(tmp_path, monkeypatch): + """With the switch off the buffer exists as zeros and x is unchanged.""" + monkeypatch.delenv("VLLM_ABLITERATION", raising=False) + monkeypatch.delenv("QWEN38_ABLITERATION", raising=False) + module = registered(torch.nn.Module()) + assert module._abliteration_on is False + assert torch.count_nonzero(module._abliteration_r) == 0 + + x = torch.randn(3, HIDDEN) + expected = x.clone() + assert abliteration.project_(module, x) is x + abliteration.project_(module, x) + assert torch.equal(x, expected) + + +def test_projection_drops_the_component(tmp_path, monkeypatch): + """After projection the tensor keeps no component along the direction.""" + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path)) + module = registered(torch.nn.Module()) + assert module._abliteration_on is True + + r = module._abliteration_r + x = torch.randn(4, HIDDEN) + abliteration.project_(module, x) + assert torch.allclose(x @ r, torch.zeros(4), atol=1e-5) + assert not torch.allclose(x, torch.zeros_like(x)) + + # Projecting an already projected tensor changes nothing. + projected = x.clone() + abliteration.project_(module, x) + assert torch.allclose(x, projected, atol=1e-6) + + +def test_registered_direction_is_unit(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path)) + module = registered(torch.nn.Module()) + assert torch.allclose(module._abliteration_r.norm(), torch.tensor(1.0), atol=1e-6) + + +def test_non_contiguous_input_is_projected(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path)) + module = registered(torch.nn.Module()) + x = torch.randn(4, HIDDEN * 2)[:, ::2] + assert not x.is_contiguous() + abliteration.project_(module, x) + assert torch.allclose(x @ module._abliteration_r, torch.zeros(4), atol=1e-5) + + +def test_width_mismatch_is_skipped(tmp_path, monkeypatch): + """A tensor without the full hidden dim must not be projected.""" + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path)) + module = registered(torch.nn.Module()) + x = torch.randn(2, HIDDEN // 2) + before = x.clone() + abliteration.project_(module, x) + assert torch.equal(x, before) + + +def test_hidden_size_mismatch_raises(tmp_path, monkeypatch): + """A direction for another model fails at boot, not silently.""" + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path, size=HIDDEN)) + with pytest.raises(ValueError, match="holds 8 values"): + registered(torch.nn.Module(), size=HIDDEN + 1) + + +def test_zero_direction_raises(tmp_path, monkeypatch): + path = tmp_path / "zero.json" + path.write_text(json.dumps({"refusal_direction": [0.0] * HIDDEN})) + monkeypatch.setenv("VLLM_ABLITERATION", str(path)) + with pytest.raises(ValueError, match="zero or non-finite"): + abliteration.load_direction(HIDDEN) + + +def test_bare_name_resolves_in_builtin_dir(tmp_path, monkeypatch): + path = write_direction(tmp_path, name="Uncensored.json") + monkeypatch.setattr(abliteration, "BUILTIN_DIRS", (str(tmp_path),)) + monkeypatch.setenv("VLLM_ABLITERATION", "Uncensored") + assert abliteration.direction_path() == path + + +def test_missing_direction_file_raises(tmp_path, monkeypatch): + monkeypatch.setattr(abliteration, "BUILTIN_DIRS", (str(tmp_path / "nope"),)) + monkeypatch.setenv("VLLM_ABLITERATION", "Uncensored") + with pytest.raises(FileNotFoundError): + abliteration.load_direction(HIDDEN) + + +def test_meta_device_turns_projection_off(tmp_path, monkeypatch): + """A meta-initialized module cannot keep the direction, so it stays off.""" + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path)) + module = torch.nn.Module() + with torch.device("meta"): + abliteration.register(module, HIDDEN, torch.float32) + assert module._abliteration_r.device.type == "meta" + assert module._abliteration_on is False + + +def test_legacy_env_name_is_accepted(tmp_path, monkeypatch): + """QWEN38_ABLITERATION still works, as existing deployments set it.""" + path = write_direction(tmp_path) + monkeypatch.delenv("VLLM_ABLITERATION", raising=False) + monkeypatch.setenv("QWEN38_ABLITERATION", path) + assert path == envs.VLLM_ABLITERATION + assert abliteration.direction_path() == path + + +def test_compile_projects_without_a_graph_break(tmp_path, monkeypatch): + """The projection stays traceable, so it does not split the compiled model.""" + monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path)) + + class Project(torch.nn.Module): + def __init__(self): + super().__init__() + abliteration.register(self, HIDDEN, torch.float32) + + def forward(self, x): + return abliteration.project_(self, x) + + module = Project().eval() + x = torch.randn(16, HIDDEN) + expected = abliteration.project_(module, x.clone()) + got = torch.compile(module, fullgraph=True)(x.clone()) + assert torch.allclose(got, expected, atol=1e-6) + + +def qwen3_5_config(dense: bool): + """A two-layer Qwen3.5 config, one full-attention and one linear layer.""" + if dense: + from transformers import Qwen3_5TextConfig as config_cls + + extra = {"intermediate_size": 256, "hidden_act": "silu"} + else: + from transformers import Qwen3_5MoeTextConfig as config_cls + + extra = { + "moe_intermediate_size": 32, + "shared_expert_intermediate_size": 32, + "num_experts": 2, + "num_experts_per_tok": 1, + "n_group": 1, + "topk_group": 1, + "norm_topk_prob": False, + "decoder_sparse_step": 1, + "mlp_only_layers": [], + } + return config_cls( + vocab_size=64, + hidden_size=HIDDEN, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + rms_norm_eps=1e-6, + layer_types=["full_attention", "linear_attention"], + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=2, + linear_num_value_heads=2, + num_nextn_predict_layers=1, + **extra, + ) + + +@pytest.mark.parametrize("dense", [True, False]) +def test_qwen3_5_model_projects_embeddings(tmp_path, monkeypatch, dense, dist_init): + """A real Qwen3.5 model registers the direction and drops it from embeddings.""" + from types import SimpleNamespace + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + + config = qwen3_5_config(dense) + path = write_direction(tmp_path, size=HIDDEN) + monkeypatch.setenv("VLLM_ABLITERATION", path) + + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace( + hf_config=config, + hf_text_config=config, + dtype=torch.bfloat16, + max_model_len=128, + is_diffusion=False, + head_dtype=None, + is_mm_prefix_lm=False, + multimodal_config=None, + quantization_config=None, + runner_type="generate", + is_moe=not dense, + logits_processors=None, + rswa_window=None, + ) + vllm_config.quant_config = None + + with set_current_vllm_config(vllm_config): + model = Qwen3_5Model(vllm_config=vllm_config) + with torch.no_grad(): + for param in model.parameters(): + param.data.normal_(0, 0.02) + + assert model._abliteration_on is True + for layer in model.layers: + assert layer._abliteration_on is True + + r = model._abliteration_r.float() + # The buffer carries the model dtype, so bf16 rounding keeps it near unit norm. + assert torch.allclose(r.norm(), torch.tensor(1.0), atol=2e-2) + + embeddings = model.embed_input_ids(torch.arange(6)).float() + residual = float((embeddings @ r).abs().max()) + scale = float(embeddings.abs().max()) + assert residual < 1e-2 * scale diff --git a/vllm/envs.py b/vllm/envs.py index 19e3250e2f..0513014486 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -227,6 +227,7 @@ if TYPE_CHECKING: VLLM_REGEX_COMPILATION_TIMEOUT_S: int = 5 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256 VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False + VLLM_ABLITERATION: str = "" VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost" VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 @@ -1189,6 +1190,15 @@ environment_variables: dict[str, Callable[[], Any]] = { os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower() in ("1", "true") ), + # Path to a refusal direction (a JSON file with a "refusal_direction" list, + # or a safetensors file with that tensor) projected out of every + # residual-stream write of the served model, or a bare name resolved under + # /opt/qwen38/abliteration and /data/abliteration. Empty or "0" is off. See + # vllm/model_executor/models/abliteration.py. QWEN38_ABLITERATION is the + # older name some deployments still set. + "VLLM_ABLITERATION": lambda: ( + os.getenv("VLLM_ABLITERATION") or os.getenv("QWEN38_ABLITERATION", "") + ), # We assume drivers can report p2p status correctly. # If the program hangs when using custom allreduce, # potantially caused by a bug in the driver (535 series), diff --git a/vllm/model_executor/models/abliteration.py b/vllm/model_executor/models/abliteration.py new file mode 100644 index 0000000000..d36ddff1c8 --- /dev/null +++ b/vllm/model_executor/models/abliteration.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional runtime refusal-direction removal ("abliteration"). + +``VLLM_ABLITERATION`` (or the legacy ``QWEN38_ABLITERATION``) names a JSON file +with a ``refusal_direction`` list, or a safetensors file with that tensor, of +length ``hidden_size``. A bare name resolves to ``.json`` in +``/opt/qwen38/abliteration``, then in ``/data/abliteration``. When set, every +write into the residual stream loses its component along the unit direction r: + + x <- x - r (r . x) + +This runs on the attention block outputs (full and linear attention), the MLP +output and the text token embeddings. For one direction it equals the weight +edit ``W' = W - r (r^T W)`` on ``self_attn.o_proj``, +``linear_attn.out_proj``, ``mlp.down_proj`` and the rows of ``embed_tokens`` +(Arditi et al., 2024), so the served weights, quantized or not, stay untouched. + +Unset, no projection kernel runs. The buffer always exists, holding zeros while +the switch is off, so the attribute a compiled or captured graph refers to is +always present, and the switch is part of the torch.compile cache key, so +toggling it never reuses a stale graph. +""" + +import json +import math +import os +from functools import lru_cache + +import torch + +from vllm import envs +from vllm.logger import init_logger + +logger = init_logger(__name__) + +ENV = "VLLM_ABLITERATION" +BUILTIN_DIRS = ("/opt/qwen38/abliteration", "/data/abliteration") +DIRECTION = "_abliteration_r" +ENABLED = "_abliteration_on" + +_shape_warned = False + + +def direction_path() -> str | None: + """The configured direction file, or None when the switch is off.""" + value = (envs.VLLM_ABLITERATION or "").strip() + if value in ("", "0"): + return None + if os.sep in value or value.endswith((".json", ".safetensors")): + return value + for directory in BUILTIN_DIRS: + path = os.path.join(directory, f"{value}.json") + if os.path.exists(path): + return path + return os.path.join(BUILTIN_DIRS[0], f"{value}.json") + + +@lru_cache(maxsize=4) +def _read(path: str) -> torch.Tensor: + """The direction as a unit float32 CPU tensor, read once per path.""" + if path.endswith(".json"): + with open(path) as f: + values = json.load(f)["refusal_direction"] + # A model built on the meta device would otherwise make this CPU tensor + # a meta one, and a meta tensor cannot be inspected or copied. + r = torch.tensor(values, dtype=torch.float32, device="cpu") + 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(dtype=torch.float32, device="cpu") + r = r.flatten() + 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") + logger.info("%s: refusal direction (%d values) from %s", ENV, r.numel(), path) + return r / norm + + +def load_direction(hidden_size: int) -> torch.Tensor | None: + """The unit refusal direction on CPU, or None when the switch is off.""" + path = direction_path() + if path is None: + return None + r = _read(path) + if r.numel() != hidden_size: + raise ValueError( + f"{ENV}: {path} holds {r.numel()} values; the model's hidden size " + f"is {hidden_size}" + ) + return r + + +def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None: + """Attach the direction as a non-persistent buffer, from a model constructor. + + Call it before CUDA-graph capture so the buffer lands on the model's device. + It holds zeros when the switch is off. A buffer that lands on the meta + device cannot hold a direction, which turns the projection off for that + module. + """ + r = load_direction(hidden_size) + buf = torch.zeros(hidden_size, dtype=dtype) + on = r is not None + if on and buf.device.type == "meta": + logger.warning( + "%s: %s is built on the meta device, which cannot hold a direction; " + "projection is off for this module", + ENV, + type(module).__name__, + ) + on = False + elif on: + buf.copy_(r) + module.register_buffer(DIRECTION, buf, persistent=False) + setattr(module, ENABLED, on) + + +def project_(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + """Remove the component of x along the direction, in place. No-op when off.""" + if not getattr(module, ENABLED, False): + return x + r = getattr(module, DIRECTION) + width = x.shape[-1] + if width != r.numel(): + _warn_shape(type(module).__name__, width, r.numel()) + return x + if r.dtype != x.dtype: + # Activations do not always carry the model dtype, and the rank-1 update + # needs one dtype throughout. + r = r.to(x.dtype) + if x.is_contiguous(): + flat = x.view(-1, width) + flat.addr_(flat @ r, r, alpha=-1.0) + else: + x.sub_(torch.outer(x.reshape(-1, width) @ r, r).view(x.shape)) + return x + + +def _warn_shape(name: str, width: int, size: int) -> None: + global _shape_warned + if not _shape_warned: + _shape_warned = True + logger.warning( + "%s: skipped a %d-wide tensor from %s; the direction has %d values", + ENV, + width, + name, + size, + ) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index f3c8fd007b..67c6fb8821 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -66,6 +66,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors from vllm.tokenizers.registry import cached_tokenizer_from_config +from . import abliteration from .interfaces import ( HasInnerState, IsHybrid, @@ -191,6 +192,7 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): self.post_attention_layernorm = Qwen3_5RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) + abliteration.register(self, config.hidden_size, model_config.dtype) self.layer_scale = getattr(config, "layer_scale", False) if self.layer_scale: @@ -254,6 +256,7 @@ class Qwen3_5Model(Qwen3NextModel): quant_config=self.quant_config, prefix=f"{prefix}.embed_tokens", ) + abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype) def get_layer(prefix: str): return Qwen3_5DecoderLayer( diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 763c6a5eb6..7d762ed158 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -37,6 +37,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors from vllm.transformers_utils.repo_utils import try_get_local_file +from . import abliteration from .interfaces import ( MultiModalEmbeddings, SupportsMultiModal, @@ -146,6 +147,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module): self.vocab_size, config.hidden_size, ) + abliteration.register(self, config.hidden_size, model_config.dtype) # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is # missing from hf_quant_config.json exclude_modules. Force unquantized. @@ -217,7 +219,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module): ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) + return abliteration.project_(self, self.embed_tokens(input_ids)) def forward( self, diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ee5c05d0a4..f63bb23a5f 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -61,6 +61,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType +from . import abliteration from .interfaces import ( EagleModelMixin, HasInnerState, @@ -537,6 +538,7 @@ class Qwen3NextDecoderLayer(nn.Module): self.post_attention_layernorm = Qwen3NextRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) + abliteration.register(self, config.hidden_size, model_config.dtype) self.layer_scale = getattr(config, "layer_scale", False) if self.layer_scale: @@ -584,6 +586,8 @@ class Qwen3NextDecoderLayer(nn.Module): else: raise ValueError("Invalid layer_type") + abliteration.project_(self, hidden_states) + if self.layer_scale: if len(hidden_states.shape) == 2: hidden_states = hidden_states * ( @@ -612,6 +616,8 @@ class Qwen3NextDecoderLayer(nn.Module): else: hidden_states = self.mlp(hidden_states) + abliteration.project_(self, hidden_states) + if self.layer_scale: if len(hidden_states.shape) == 2: hidden_states = hidden_states * ( @@ -661,6 +667,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin): self.vocab_size, config.hidden_size, ) + abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype) def get_layer(prefix: str): return Qwen3NextDecoderLayer( @@ -689,7 +696,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin): self.aux_hidden_state_layers: tuple[int, ...] = () def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) + return abliteration.project_(self, self.embed_tokens(input_ids)) @property def use_sequence_parallel(self) -> bool: