vLLM/vllm/model_executor/models/abliteration.py

# 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 ``<name>.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
_register_logged = False


def clear_logs() -> None:
    """Forget the once-only log flags, for tests."""
    global _register_logged, _shape_warned
    _register_logged = False
    _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


@torch._dynamo.disable
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 construction pass that cannot hold
    data, the meta device or a fake-tensor compile pass, turns the projection off
    for that module and says so.
    """
    r = load_direction(hidden_size)
    buf = torch.zeros(hidden_size, dtype=dtype)
    on = r is not None
    if on:
        try:
            buf.copy_(r)
        except (RuntimeError, NotImplementedError):
            # Meta or fake tensor construction: no room for a direction here.
            logger.warning(
                "%s: %s is constructed on the %s device, which cannot hold a "
                "direction; projection is off for this module",
                ENV,
                type(module).__name__,
                buf.device.type,
            )
            on = False
    module.register_buffer(DIRECTION, buf, persistent=False)
    setattr(module, ENABLED, on)
    # Logging from project_ would sit inside a torch.compile region, which cannot
    # trace a logger call, so the proof that the switch is live goes here.
    global _register_logged
    if on and not _register_logged:
        _register_logged = True
        logger.info(
            "%s: projection is active, first registered on %s (hidden size %d)",
            ENV,
            type(module).__name__,
            hidden_size,
        )


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


@torch._dynamo.disable
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,
        )