vLLM/tests/models/test_abliteration.py

# 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()
    abliteration.clear_logs()
    yield
    abliteration._read.cache_clear()
    abliteration.clear_logs()


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_first_projection_logs_that_it_is_active(tmp_path, monkeypatch, caplog):
    """The boot log proves the switch is live, not just loaded."""
    import logging

    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
    module = registered(torch.nn.Linear(HIDDEN, HIDDEN))
    with caplog.at_level(
        logging.INFO, logger="vllm.model_executor.models.abliteration"
    ):
        abliteration.project_(module, torch.randn(2, HIDDEN))
    assert "active, projected a 8-wide tensor" in caplog.text


def test_flash_next_direction_is_rejected_by_the_27b_model(tmp_path, monkeypatch):
    """Flash-Next stores 2560 values; the dense 27B stream is 5120 wide."""
    path = write_direction(tmp_path, name="flash_next.json", size=2560)
    monkeypatch.setenv("VLLM_ABLITERATION", path)
    with pytest.raises(
        ValueError, match="holds 2560 values; the model's hidden size is 5120"
    ):
        registered(torch.nn.Module(), size=5120)


def test_half_width_activation_is_skipped(tmp_path, monkeypatch):
    """A 2560-wide tensor cannot be projected by the 5120 direction."""
    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path, size=5120))
    module = registered(torch.nn.Module(), size=5120)
    x = torch.randn(4, 2560)
    before = x.clone()
    abliteration.project_(module, x)
    assert torch.equal(x, before)


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