diff --git a/tests/models/test_abliteration.py b/tests/models/test_abliteration.py index 5e9435ee9a..b052cd60c3 100644 --- a/tests/models/test_abliteration.py +++ b/tests/models/test_abliteration.py @@ -16,8 +16,10 @@ 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): @@ -97,6 +99,39 @@ def test_hidden_size_mismatch_raises(tmp_path, monkeypatch): 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})) diff --git a/vllm/model_executor/models/abliteration.py b/vllm/model_executor/models/abliteration.py index d36ddff1c8..dba47574dd 100644 --- a/vllm/model_executor/models/abliteration.py +++ b/vllm/model_executor/models/abliteration.py @@ -39,6 +39,14 @@ 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: @@ -92,29 +100,44 @@ def load_direction(hidden_size: int) -> torch.Tensor | None: 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 buffer that lands on the meta - device cannot hold a direction, which turns the projection off for that - module. + 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 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) + 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: @@ -138,6 +161,7 @@ def project_(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: return x +@torch._dynamo.disable def _warn_shape(name: str, width: int, size: int) -> None: global _shape_warned if not _shape_warned: diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index 4ae8f21840..b076e1ff6b 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.models import abliteration from vllm.model_executor.models.qwen3_next import ( Qwen3NextDecoderLayer, Qwen3NextModel, @@ -66,6 +67,7 @@ class Qwen3NextMultiTokenPredictor(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 the checkpoint quant exclude list (its `ignore` glob @@ -113,7 +115,7 @@ class Qwen3NextMultiTokenPredictor(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,