From 5d6754ff6b103e3bc11e917ab1028e9cf031f29f Mon Sep 17 00:00:00 2001 From: MOVZX Date: Thu, 15 Oct 2026 22:56:06 +0700 Subject: [PATCH] fix: keep the refusal-direction projection dynamo-safe With VLLM_ABLITERATION set, boot crashed: dynamo traces abliteration.register from Qwen3NextGatedDeltaNet.__init__, and buf.copy_(r) fails on the meta device during the compile-time model construction pass. Disable tracing for register and guard the copy so an unmaterializable buffer turns projection off with a warning. Log activation at registration instead of inside project_: project_ runs in the compiled region, and a logger call there aborts dynamo. The width guard uses the same escape hatch. Also project the embedding that Qwen3NextMultiTokenPredictor builds for itself, which the first pass missed, and add tests for the 2560-value Flash-Next direction against a 5120 model and for a half-width activation. Co-authored-by: Hinata Signed-off-by: MOVZX --- tests/models/test_abliteration.py | 35 ++++++++++++++ vllm/model_executor/models/abliteration.py | 50 +++++++++++++++----- vllm/model_executor/models/qwen3_next_mtp.py | 4 +- 3 files changed, 75 insertions(+), 14 deletions(-) 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,