[watermarking] golden tests for backwards compatibility (#56809)
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Signed-off-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Raphaël Rialland
co-authored by
Simon Veitner
OpenAI Codex
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Simon Veitner
OpenAI Codex
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
b6761e8ded
commit
4ccfe12398
@@ -211,6 +211,7 @@ steps:
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/v1
|
||||
- tests/watermarking
|
||||
device: cpu-small
|
||||
commands:
|
||||
# split the test to avoid interference
|
||||
@@ -222,6 +223,7 @@ steps:
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/ec_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
- pytest -v -s watermarking
|
||||
|
||||
- label: ":nvidia: (H200 MIG 18GB) Extract Hidden States Integration"
|
||||
key: extract-hidden-states-integration
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,24 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Regenerate tests/watermarking/watermarking_goldens.json."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from tests.watermarking.golden_candidates import golden_payload, validate_golden_guards
|
||||
|
||||
GOLDENS_PATH = Path(__file__).with_name("watermarking_goldens.json")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
payload = golden_payload()
|
||||
validate_golden_guards(payload["candidates"])
|
||||
GOLDENS_PATH.write_text(
|
||||
json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
|
||||
)
|
||||
print(f"wrote {len(payload['candidates'])} candidates to {GOLDENS_PATH}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,885 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Deterministic watermarking golden candidates."""
|
||||
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import get_args
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config.watermarking import (
|
||||
WatermarkConfig,
|
||||
WatermarkContextScope,
|
||||
WatermarkingAlgorithm,
|
||||
WatermarkPRFName,
|
||||
derive_watermark_key,
|
||||
)
|
||||
from vllm.v1.watermarking import (
|
||||
DualKeyGumbelWatermarkDetector,
|
||||
GumbelWatermarkDetector,
|
||||
WatermarkDetection,
|
||||
WatermarkDetector,
|
||||
create_prf,
|
||||
create_watermarker,
|
||||
)
|
||||
from vllm.v1.worker.gpu.sample.watermark import repeated_context_mask
|
||||
|
||||
GOLDEN_SCHEMA_VERSION = 3
|
||||
|
||||
# JSON floats round-trip exactly, but libm differs by an ULP across builds.
|
||||
GOLDEN_FLOAT_RTOL = 1e-9
|
||||
GOLDEN_FLOAT_ATOL = 1e-300
|
||||
|
||||
P_VALUE_RATIO_MIN_DISTANCE = 1e-6
|
||||
REPETITIVE_MIN_KEY_PERTURBATION_TOKENS = 8
|
||||
REPETITIVE_MAX_P_VALUE_RATIO = 0.5
|
||||
ROUTING_BOUNDARY_MIN_DISTANCE = 1e-6
|
||||
|
||||
WATERMARK_CONFIG_FIELDS = (
|
||||
"algorithm",
|
||||
"alpha",
|
||||
"context_width",
|
||||
"deduplicate_contexts",
|
||||
"deduplicate_contexts_max_history",
|
||||
"prf",
|
||||
"allow_target_only_watermarking",
|
||||
)
|
||||
|
||||
# round(((i + 1) * (5**0.5 - 1) / 2) % 1.0, 4) for i in range(64)
|
||||
ROUTING_UNIFORMS: tuple[float, ...] = (
|
||||
0.618, 0.2361, 0.8541, 0.4721, 0.0902, 0.7082, 0.3262, 0.9443,
|
||||
0.5623, 0.1803, 0.7984, 0.4164, 0.0344, 0.6525, 0.2705, 0.8885,
|
||||
0.5066, 0.1246, 0.7426, 0.3607, 0.9787, 0.5967, 0.2148, 0.8328,
|
||||
0.4508, 0.0689, 0.6869, 0.305, 0.923, 0.541, 0.1591, 0.7771,
|
||||
0.3951, 0.0132, 0.6312, 0.2492, 0.8673, 0.4853, 0.1033, 0.7214,
|
||||
0.3394, 0.9574, 0.5755, 0.1935, 0.8115, 0.4296, 0.0476, 0.6656,
|
||||
0.2837, 0.9017, 0.5197, 0.1378, 0.7558, 0.3738, 0.9919, 0.6099,
|
||||
0.2279, 0.846, 0.464, 0.082, 0.7001, 0.3181, 0.9361, 0.5542,
|
||||
) # fmt: skip
|
||||
|
||||
|
||||
class GoldenFormatError(ValueError):
|
||||
"""The stored goldens do not have the shape this module expects."""
|
||||
|
||||
|
||||
class GoldenGuardError(ValueError):
|
||||
"""A candidate is too weak or too unstable to be frozen as a golden."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DeterministicGenerationFixture:
|
||||
vocabulary_size: int = 128
|
||||
num_tokens: int = 64
|
||||
logit_denominator: int = 16
|
||||
logit_modulus: int = 131
|
||||
dominant_token: int | None = None
|
||||
dominant_bias: int = 0
|
||||
prompt: tuple[int, ...] = ()
|
||||
routing_uniforms: tuple[float, ...] = ROUTING_UNIFORMS
|
||||
|
||||
def logits(self, position: int) -> torch.Tensor:
|
||||
token_ids = torch.arange(self.vocabulary_size, dtype=torch.int64)
|
||||
values = (token_ids * 37 + position * 17) % self.logit_modulus - 56
|
||||
if self.dominant_token is not None:
|
||||
values[self.dominant_token] += self.dominant_bias
|
||||
return (values.to(torch.float32) / self.logit_denominator).unsqueeze(0)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WatermarkingSchemeConfig:
|
||||
context_width: int = 4
|
||||
generation_alpha: float = 0.1
|
||||
detection_alpha: float = 0.1
|
||||
generation_deduplicate_contexts: WatermarkContextScope = "single_turn"
|
||||
generation_deduplicate_contexts_max_history: int | None = 8192
|
||||
detection_deduplicate_contexts: bool = True
|
||||
p_value_threshold: float = 0.01
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GenerationTrace:
|
||||
routing_draws: int
|
||||
key_b_routed: int | None
|
||||
key_b_effective: int | None
|
||||
dedup_skips: int
|
||||
partial_context_skips: int
|
||||
|
||||
|
||||
class _DeterministicSampler:
|
||||
def __init__(self, routing_uniforms: tuple[float, ...]) -> None:
|
||||
self.routing_uniforms = routing_uniforms
|
||||
self.routing_position = 0
|
||||
self.key_b_routed = 0
|
||||
|
||||
def __call__(self, logits: torch.Tensor) -> torch.Tensor:
|
||||
if logits.shape[-1] != 2:
|
||||
return logits.argmax(dim=-1)
|
||||
|
||||
probability_a = logits.softmax(dim=-1)[:, 0]
|
||||
uniform = self.routing_uniforms[
|
||||
self.routing_position % len(self.routing_uniforms)
|
||||
]
|
||||
self.routing_position += 1
|
||||
routed = (probability_a <= uniform).to(torch.int64)
|
||||
self.key_b_routed += int(routed.sum())
|
||||
return routed
|
||||
|
||||
|
||||
def routing_boundary(alpha: float) -> float:
|
||||
"""Return the realised key-A probability the deterministic sampler sees."""
|
||||
logits = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log()
|
||||
return float(logits.softmax(dim=-1)[0])
|
||||
|
||||
|
||||
def _create_gumbel_detector(
|
||||
key: int,
|
||||
config: WatermarkingSchemeConfig,
|
||||
prf: WatermarkPRFName,
|
||||
) -> WatermarkDetector:
|
||||
return GumbelWatermarkDetector(
|
||||
key=key,
|
||||
context_width=config.context_width,
|
||||
p_value_threshold=config.p_value_threshold,
|
||||
prf=prf,
|
||||
deduplicate_contexts=config.detection_deduplicate_contexts,
|
||||
)
|
||||
|
||||
|
||||
def _create_dual_key_gumbel_detector(
|
||||
key: int,
|
||||
config: WatermarkingSchemeConfig,
|
||||
prf: WatermarkPRFName,
|
||||
) -> WatermarkDetector:
|
||||
return DualKeyGumbelWatermarkDetector(
|
||||
key=key,
|
||||
context_width=config.context_width,
|
||||
p_value_threshold=config.p_value_threshold,
|
||||
prf=prf,
|
||||
deduplicate_contexts=config.detection_deduplicate_contexts,
|
||||
alpha=config.detection_alpha,
|
||||
)
|
||||
|
||||
|
||||
DETECTOR_FACTORIES = {
|
||||
"gumbel": _create_gumbel_detector,
|
||||
"dual_key_gumbel": _create_dual_key_gumbel_detector,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WatermarkingCandidate:
|
||||
id: str
|
||||
scheme: WatermarkingAlgorithm
|
||||
scheme_config: WatermarkingSchemeConfig
|
||||
prf: WatermarkPRFName
|
||||
key: int
|
||||
detection_key: int
|
||||
fixture: DeterministicGenerationFixture = DeterministicGenerationFixture()
|
||||
|
||||
def configuration(self) -> dict[str, object]:
|
||||
prf_version = getattr(create_prf(self.prf, self.key), "version", None)
|
||||
if not isinstance(prf_version, str):
|
||||
raise ValueError(f"PRF {self.prf} does not define a version")
|
||||
return {
|
||||
"scheme": self.scheme,
|
||||
"scheme_config": {
|
||||
"context_width": self.scheme_config.context_width,
|
||||
"generation_alpha": self.scheme_config.generation_alpha,
|
||||
"detection_alpha": self.scheme_config.detection_alpha,
|
||||
"generation_deduplicate_contexts": (
|
||||
self.scheme_config.generation_deduplicate_contexts
|
||||
),
|
||||
"generation_deduplicate_contexts_max_history": (
|
||||
self.scheme_config.generation_deduplicate_contexts_max_history
|
||||
),
|
||||
"detection_deduplicate_contexts": (
|
||||
self.scheme_config.detection_deduplicate_contexts
|
||||
),
|
||||
"p_value_threshold": self.scheme_config.p_value_threshold,
|
||||
},
|
||||
"prf": self.prf,
|
||||
"prf_version": prf_version,
|
||||
"key": str(self.key),
|
||||
"detection_key": str(self.detection_key),
|
||||
"fixture": {
|
||||
"vocabulary_size": self.fixture.vocabulary_size,
|
||||
"num_tokens": self.fixture.num_tokens,
|
||||
"logit_denominator": self.fixture.logit_denominator,
|
||||
"logit_modulus": self.fixture.logit_modulus,
|
||||
"dominant_token": self.fixture.dominant_token,
|
||||
"dominant_bias": self.fixture.dominant_bias,
|
||||
"prompt": list(self.fixture.prompt),
|
||||
},
|
||||
}
|
||||
|
||||
def resolved(self) -> dict[str, object]:
|
||||
"""Return the vLLM-side state this candidate resolves to.
|
||||
|
||||
An allowlist rather than the whole dataclass, so a new production field
|
||||
is a deliberate golden change: WATERMARK_CONFIG_FIELDS names the fields
|
||||
recorded here and test_goldens.py checks it still covers
|
||||
WatermarkConfig. Keys exceed 2**53 and are decimal strings.
|
||||
"""
|
||||
config = self._watermark_config()
|
||||
detector = self._detector()
|
||||
key_b_prf = getattr(detector, "key_b_prf", None)
|
||||
dual_key = self.scheme == "dual_key_gumbel"
|
||||
return {
|
||||
"watermark_config": {
|
||||
field: getattr(config, field) for field in WATERMARK_CONFIG_FIELDS
|
||||
},
|
||||
"derived_keys": {
|
||||
"key_a": str(derive_watermark_key(self.key, b"key_a"))
|
||||
if dual_key
|
||||
else None,
|
||||
"key_b": str(derive_watermark_key(self.key, b"key_b"))
|
||||
if dual_key
|
||||
else None,
|
||||
},
|
||||
"detector": {
|
||||
"type": type(detector).__name__,
|
||||
"context_width": detector.context_width,
|
||||
"p_value_threshold": detector.p_value_threshold,
|
||||
"deduplicate_contexts": detector.deduplicate_contexts,
|
||||
"alpha": getattr(detector, "alpha", None),
|
||||
"prf_key": str(detector.prf.key),
|
||||
"key_b_prf_key": None if key_b_prf is None else str(key_b_prf.key),
|
||||
},
|
||||
}
|
||||
|
||||
def generate(self) -> list[int]:
|
||||
return self.generate_with_trace()[0]
|
||||
|
||||
def generate_with_trace(self) -> tuple[list[int], GenerationTrace]:
|
||||
watermarker = create_watermarker(self._watermark_config())
|
||||
sampler = _DeterministicSampler(self.fixture.routing_uniforms)
|
||||
generated: list[int] = []
|
||||
dedup_skips = 0
|
||||
partial_context_skips = 0
|
||||
key_b_effective = 0
|
||||
|
||||
for position in range(self.fixture.num_tokens):
|
||||
context = self._context(generated)
|
||||
skip_mask = self._skip_mask(generated, context)
|
||||
skipped = skip_mask is not None and bool(skip_mask.item())
|
||||
if skipped:
|
||||
dedup_skips += 1
|
||||
if len(generated) < self.scheme_config.context_width:
|
||||
partial_context_skips += 1
|
||||
routed_before = sampler.key_b_routed
|
||||
sample = watermarker.sample(
|
||||
self.fixture.logits(position),
|
||||
context,
|
||||
sampler,
|
||||
skip_mask,
|
||||
)
|
||||
if sampler.key_b_routed > routed_before and not skipped:
|
||||
key_b_effective += 1
|
||||
generated.append(int(sample.token_ids.item()))
|
||||
|
||||
routes = self.scheme == "dual_key_gumbel" and (
|
||||
0 < self.scheme_config.generation_alpha < 1
|
||||
)
|
||||
trace = GenerationTrace(
|
||||
routing_draws=sampler.routing_position,
|
||||
key_b_routed=sampler.key_b_routed if routes else None,
|
||||
key_b_effective=key_b_effective if routes else None,
|
||||
dedup_skips=dedup_skips,
|
||||
partial_context_skips=partial_context_skips,
|
||||
)
|
||||
return generated, trace
|
||||
|
||||
def detect(self, token_ids: list[int]) -> WatermarkDetection:
|
||||
return self._detector().detect(token_ids)
|
||||
|
||||
def golden_entry(
|
||||
self,
|
||||
detection_token_ids: list[int] | None = None,
|
||||
) -> dict:
|
||||
"""Generate, detect and trace this candidate into one golden entry.
|
||||
|
||||
``detection_token_ids`` runs detection over stored tokens instead of
|
||||
the freshly generated ones, so a token drift and a detection drift are
|
||||
reported independently.
|
||||
"""
|
||||
token_ids, trace = self.generate_with_trace()
|
||||
scored = token_ids if detection_token_ids is None else detection_token_ids
|
||||
detection = self.detect(scored)
|
||||
threshold = self.scheme_config.p_value_threshold
|
||||
return {
|
||||
"configuration": self.configuration(),
|
||||
"resolved": self.resolved(),
|
||||
"generation": token_ids,
|
||||
"trace": _trace_payload(trace),
|
||||
"detection": {
|
||||
"score": detection.score,
|
||||
"p_value": detection.p_value,
|
||||
"p_value_ratio": detection.p_value / threshold,
|
||||
"num_scored_tokens": detection.num_scored_tokens,
|
||||
"is_watermarked": detection.is_watermarked,
|
||||
},
|
||||
}
|
||||
|
||||
def _watermark_config(self) -> WatermarkConfig:
|
||||
return WatermarkConfig(
|
||||
algorithm=self.scheme,
|
||||
key=self.key,
|
||||
alpha=self.scheme_config.generation_alpha,
|
||||
context_width=self.scheme_config.context_width,
|
||||
deduplicate_contexts=(self.scheme_config.generation_deduplicate_contexts),
|
||||
deduplicate_contexts_max_history=(
|
||||
self.scheme_config.generation_deduplicate_contexts_max_history
|
||||
),
|
||||
prf=self.prf,
|
||||
)
|
||||
|
||||
def _detector(self) -> WatermarkDetector:
|
||||
return DETECTOR_FACTORIES[self.scheme](
|
||||
self.detection_key,
|
||||
self.scheme_config,
|
||||
self.prf,
|
||||
)
|
||||
|
||||
def _context(self, generated: list[int]) -> torch.Tensor:
|
||||
width = self.scheme_config.context_width
|
||||
padded = [-1] * width + generated
|
||||
return torch.tensor([padded[-width:]], dtype=torch.int64)
|
||||
|
||||
def _skip_mask(
|
||||
self, generated: list[int], context: torch.Tensor
|
||||
) -> torch.Tensor | None:
|
||||
scope = self.scheme_config.generation_deduplicate_contexts
|
||||
if scope == "none":
|
||||
return None
|
||||
|
||||
prompt_len = len(self.fixture.prompt)
|
||||
storage = torch.zeros(
|
||||
(1, prompt_len + self.fixture.num_tokens),
|
||||
dtype=torch.int64,
|
||||
)
|
||||
if prompt_len:
|
||||
storage[0, :prompt_len] = torch.tensor(self.fixture.prompt)
|
||||
if generated:
|
||||
storage[0, prompt_len : prompt_len + len(generated)] = torch.tensor(
|
||||
generated
|
||||
)
|
||||
return repeated_context_mask(
|
||||
storage,
|
||||
torch.tensor([0]),
|
||||
torch.tensor([prompt_len]),
|
||||
torch.tensor([prompt_len + len(generated)]),
|
||||
context,
|
||||
self.scheme_config.generation_deduplicate_contexts_max_history,
|
||||
include_prompt=scope == "all",
|
||||
skip_partial_context=scope == "all",
|
||||
)
|
||||
|
||||
|
||||
def _trace_payload(trace: GenerationTrace) -> dict:
|
||||
return {
|
||||
"routing_draws": trace.routing_draws,
|
||||
"key_b_routed": trace.key_b_routed,
|
||||
"key_b_effective": trace.key_b_effective,
|
||||
"dedup_skips": trace.dedup_skips,
|
||||
"partial_context_skips": trace.partial_context_skips,
|
||||
}
|
||||
|
||||
|
||||
BALANCED_FIXTURE = DeterministicGenerationFixture()
|
||||
|
||||
# The first five tokens of BALANCED_FIXTURE under scope "all", so position 4 repeats.
|
||||
PROMPT_FIXTURE = DeterministicGenerationFixture(prompt=(46, 42, 38, 34, 122))
|
||||
|
||||
REPETITIVE_FIXTURE = DeterministicGenerationFixture(
|
||||
vocabulary_size=16,
|
||||
num_tokens=48,
|
||||
dominant_token=3,
|
||||
dominant_bias=52,
|
||||
)
|
||||
|
||||
MIDDLE_FIXTURE = DeterministicGenerationFixture(
|
||||
vocabulary_size=8,
|
||||
num_tokens=64,
|
||||
dominant_token=3,
|
||||
dominant_bias=16,
|
||||
)
|
||||
|
||||
# Few enough scored tokens that the p_value lands just under its threshold.
|
||||
NEAR_THRESHOLD_FIXTURE = DeterministicGenerationFixture(
|
||||
vocabulary_size=16,
|
||||
num_tokens=8,
|
||||
)
|
||||
|
||||
|
||||
def _candidate(
|
||||
id: str,
|
||||
scheme: WatermarkingAlgorithm,
|
||||
key: int,
|
||||
*,
|
||||
prf: WatermarkPRFName,
|
||||
detection_key: int | None = None,
|
||||
context_width: int = 4,
|
||||
generation_alpha: float = 0.1,
|
||||
detection_alpha: float | None = None,
|
||||
generation_deduplicate_contexts: WatermarkContextScope = "single_turn",
|
||||
generation_deduplicate_contexts_max_history: int | None = 8192,
|
||||
detection_deduplicate_contexts: bool = True,
|
||||
p_value_threshold: float = 0.01,
|
||||
fixture: DeterministicGenerationFixture = BALANCED_FIXTURE,
|
||||
) -> WatermarkingCandidate:
|
||||
return WatermarkingCandidate(
|
||||
id=id,
|
||||
scheme=scheme,
|
||||
scheme_config=WatermarkingSchemeConfig(
|
||||
context_width=context_width,
|
||||
generation_alpha=generation_alpha,
|
||||
detection_alpha=(
|
||||
generation_alpha if detection_alpha is None else detection_alpha
|
||||
),
|
||||
generation_deduplicate_contexts=generation_deduplicate_contexts,
|
||||
generation_deduplicate_contexts_max_history=(
|
||||
generation_deduplicate_contexts_max_history
|
||||
),
|
||||
detection_deduplicate_contexts=detection_deduplicate_contexts,
|
||||
p_value_threshold=p_value_threshold,
|
||||
),
|
||||
prf=prf,
|
||||
key=key,
|
||||
detection_key=key if detection_key is None else detection_key,
|
||||
fixture=fixture,
|
||||
)
|
||||
|
||||
|
||||
WATERMARKING_CANDIDATES = (
|
||||
_candidate("gumbel-philox-key0-cw1", "gumbel", 0, prf="philox", context_width=1),
|
||||
_candidate("gumbel-philox-key42-cw4", "gumbel", 42, prf="philox"),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-wrong-key",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
detection_key=43,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key2p32-cw5-no-dedup",
|
||||
"gumbel",
|
||||
2**32,
|
||||
prf="philox",
|
||||
context_width=5,
|
||||
generation_deduplicate_contexts="none",
|
||||
detection_deduplicate_contexts=False,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-keymax-cw16-all",
|
||||
"gumbel",
|
||||
2**64 - 1,
|
||||
prf="philox",
|
||||
context_width=16,
|
||||
generation_deduplicate_contexts="all",
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-all",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts="all",
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-all-prompt",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts="all",
|
||||
fixture=PROMPT_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-mid-history8",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts_max_history=8,
|
||||
fixture=MIDDLE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-mid-history-none",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts_max_history=None,
|
||||
fixture=MIDDLE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-mid-no-dedup",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts="none",
|
||||
generation_deduplicate_contexts_max_history=None,
|
||||
detection_deduplicate_contexts=False,
|
||||
fixture=MIDDLE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-mid-all",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts="all",
|
||||
generation_deduplicate_contexts_max_history=None,
|
||||
fixture=MIDDLE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-key42-cw4-near-threshold",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
p_value_threshold=0.0015,
|
||||
fixture=NEAR_THRESHOLD_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-repeated-cw1",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
context_width=1,
|
||||
fixture=REPETITIVE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"gumbel-philox-repeated-cw4-all",
|
||||
"gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts="all",
|
||||
fixture=REPETITIVE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-key0-cw1-alpha0",
|
||||
"dual_key_gumbel",
|
||||
0,
|
||||
prf="philox",
|
||||
context_width=1,
|
||||
generation_alpha=0.0,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-key0-cw4-alpha0.1",
|
||||
"dual_key_gumbel",
|
||||
0,
|
||||
prf="philox",
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-key42-cw4-alpha-mismatch",
|
||||
"dual_key_gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_alpha=0.1,
|
||||
detection_alpha=0.2,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-key2p32-cw5-alpha0.5",
|
||||
"dual_key_gumbel",
|
||||
2**32,
|
||||
prf="philox",
|
||||
context_width=5,
|
||||
generation_alpha=0.5,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-keymax-cw16-alpha1",
|
||||
"dual_key_gumbel",
|
||||
2**64 - 1,
|
||||
prf="philox",
|
||||
context_width=16,
|
||||
generation_alpha=1.0,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-keymax-cw4-alpha0.4",
|
||||
"dual_key_gumbel",
|
||||
2**64 - 1,
|
||||
prf="philox",
|
||||
generation_alpha=0.4,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-key7-cw4-near-threshold",
|
||||
"dual_key_gumbel",
|
||||
7,
|
||||
prf="philox",
|
||||
p_value_threshold=0.003,
|
||||
fixture=NEAR_THRESHOLD_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-repeated-cw1",
|
||||
"dual_key_gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
context_width=1,
|
||||
fixture=REPETITIVE_FIXTURE,
|
||||
),
|
||||
_candidate(
|
||||
"dual-key-gumbel-philox-repeated-cw4-all",
|
||||
"dual_key_gumbel",
|
||||
42,
|
||||
prf="philox",
|
||||
generation_deduplicate_contexts="all",
|
||||
fixture=REPETITIVE_FIXTURE,
|
||||
),
|
||||
)
|
||||
|
||||
DEDUPLICATION_TWINS = (
|
||||
(
|
||||
"gumbel-philox-key42-cw4-mid-history-none",
|
||||
"gumbel-philox-key42-cw4-mid-no-dedup",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def validate_golden_guards(entries: Mapping[str, dict]) -> None:
|
||||
candidates_by_id = {
|
||||
candidate.id: candidate for candidate in WATERMARKING_CANDIDATES
|
||||
}
|
||||
combinations = {
|
||||
(candidate.scheme, candidate.prf) for candidate in WATERMARKING_CANDIDATES
|
||||
}
|
||||
expected_combinations = configured_algorithm_prf_combinations()
|
||||
if combinations != expected_combinations:
|
||||
missing_combinations = sorted(expected_combinations - combinations)
|
||||
unexpected_combinations = sorted(combinations - expected_combinations)
|
||||
raise GoldenGuardError(
|
||||
"candidate scheme/PRF coverage mismatch: "
|
||||
f"missing {missing_combinations}, unexpected {unexpected_combinations}"
|
||||
)
|
||||
|
||||
missing = sorted(set(candidates_by_id) - set(entries))
|
||||
if missing:
|
||||
raise GoldenGuardError(f"golden entries are missing candidates: {missing}")
|
||||
|
||||
for candidate in WATERMARKING_CANDIDATES:
|
||||
entry = entries[candidate.id]
|
||||
ratio = entry["detection"]["p_value_ratio"]
|
||||
if abs(ratio - 1.0) < P_VALUE_RATIO_MIN_DISTANCE:
|
||||
raise GoldenGuardError(
|
||||
f"{candidate.id}: p_value is within "
|
||||
f"{P_VALUE_RATIO_MIN_DISTANCE:g} of p_value_threshold "
|
||||
f"(ratio {ratio!r}), so is_watermarked is not reproducible"
|
||||
)
|
||||
|
||||
if candidate.fixture == REPETITIVE_FIXTURE:
|
||||
alternate_tokens = replace(candidate, key=candidate.key ^ 1).generate()
|
||||
changed_tokens = sum(
|
||||
actual != alternate
|
||||
for actual, alternate in zip(
|
||||
entry["generation"], alternate_tokens, strict=True
|
||||
)
|
||||
)
|
||||
if changed_tokens < REPETITIVE_MIN_KEY_PERTURBATION_TOKENS:
|
||||
raise GoldenGuardError(
|
||||
f"{candidate.id}: key ^ 1 changes only {changed_tokens} tokens; "
|
||||
f"expected at least {REPETITIVE_MIN_KEY_PERTURBATION_TOKENS}"
|
||||
)
|
||||
full_context_skips = (
|
||||
entry["trace"]["dedup_skips"] - entry["trace"]["partial_context_skips"]
|
||||
)
|
||||
if full_context_skips < 1:
|
||||
raise GoldenGuardError(
|
||||
f"{candidate.id}: repetitive fixture has no full-context "
|
||||
"deduplication skip"
|
||||
)
|
||||
if ratio > REPETITIVE_MAX_P_VALUE_RATIO:
|
||||
raise GoldenGuardError(
|
||||
f"{candidate.id}: p_value_ratio {ratio!r} exceeds "
|
||||
f"{REPETITIVE_MAX_P_VALUE_RATIO}"
|
||||
)
|
||||
|
||||
alpha = candidate.scheme_config.generation_alpha
|
||||
if candidate.scheme == "dual_key_gumbel" and 0 < alpha < 1:
|
||||
boundary = routing_boundary(alpha)
|
||||
distance = min(
|
||||
abs(uniform - boundary)
|
||||
for uniform in candidate.fixture.routing_uniforms
|
||||
)
|
||||
if distance < ROUTING_BOUNDARY_MIN_DISTANCE:
|
||||
raise GoldenGuardError(
|
||||
f"{candidate.id}: routing uniform is only {distance!r} from "
|
||||
f"the realised boundary; expected at least "
|
||||
f"{ROUTING_BOUNDARY_MIN_DISTANCE:g}"
|
||||
)
|
||||
|
||||
for enabled_id, disabled_id in DEDUPLICATION_TWINS:
|
||||
enabled = candidates_by_id[enabled_id]
|
||||
disabled = candidates_by_id[disabled_id]
|
||||
expected_disabled = replace(
|
||||
enabled,
|
||||
id=disabled_id,
|
||||
scheme_config=replace(
|
||||
enabled.scheme_config,
|
||||
generation_deduplicate_contexts="none",
|
||||
detection_deduplicate_contexts=False,
|
||||
),
|
||||
)
|
||||
if disabled != expected_disabled:
|
||||
raise GoldenGuardError(
|
||||
f"{enabled_id} and {disabled_id} must differ only in generation "
|
||||
"and detection deduplication"
|
||||
)
|
||||
|
||||
enabled_entry = entries[enabled_id]
|
||||
disabled_entry = entries[disabled_id]
|
||||
if enabled_entry["trace"]["dedup_skips"] < 1:
|
||||
raise GoldenGuardError(f"{enabled_id}: expected at least one dedup skip")
|
||||
if disabled_entry["trace"]["dedup_skips"] != 0:
|
||||
raise GoldenGuardError(f"{disabled_id}: expected zero dedup skips")
|
||||
if enabled_entry["generation"] == disabled_entry["generation"]:
|
||||
raise GoldenGuardError(
|
||||
f"{enabled_id} and {disabled_id}: generations must differ"
|
||||
)
|
||||
enabled_scored = enabled_entry["detection"]["num_scored_tokens"]
|
||||
disabled_scored = disabled_entry["detection"]["num_scored_tokens"]
|
||||
if enabled_scored == disabled_scored:
|
||||
raise GoldenGuardError(
|
||||
f"{enabled_id} and {disabled_id}: detector scored-token counts "
|
||||
"must differ"
|
||||
)
|
||||
|
||||
|
||||
def golden_payload() -> dict:
|
||||
candidates = {
|
||||
candidate.id: candidate.golden_entry() for candidate in WATERMARKING_CANDIDATES
|
||||
}
|
||||
return {
|
||||
"schema_version": GOLDEN_SCHEMA_VERSION,
|
||||
"candidates": candidates,
|
||||
}
|
||||
|
||||
|
||||
def floats_match(actual: float, expected: float) -> bool:
|
||||
return math.isclose(
|
||||
actual,
|
||||
expected,
|
||||
rel_tol=GOLDEN_FLOAT_RTOL,
|
||||
abs_tol=GOLDEN_FLOAT_ATOL,
|
||||
)
|
||||
|
||||
|
||||
def _relative_difference(actual: float, expected: float) -> float:
|
||||
if expected == 0.0:
|
||||
return 0.0 if actual == 0.0 else math.inf
|
||||
return abs(actual - expected) / abs(expected)
|
||||
|
||||
|
||||
def _compare_float(
|
||||
field: str,
|
||||
actual: float,
|
||||
expected: float,
|
||||
differences: list[str],
|
||||
) -> None:
|
||||
if not floats_match(actual, expected):
|
||||
differences.append(
|
||||
f"{field}: expected {expected!r}, got {actual!r}, "
|
||||
f"relative difference {_relative_difference(actual, expected):.3e}"
|
||||
)
|
||||
|
||||
|
||||
def _compare_exact(
|
||||
field: str,
|
||||
actual: object,
|
||||
expected: object,
|
||||
differences: list[str],
|
||||
) -> None:
|
||||
if actual != expected:
|
||||
differences.append(f"{field}: expected {expected!r}, got {actual!r}")
|
||||
|
||||
|
||||
def _compare_mapping(
|
||||
field: str,
|
||||
actual: Mapping[str, object],
|
||||
expected: Mapping[str, object],
|
||||
differences: list[str],
|
||||
) -> None:
|
||||
for key in sorted(set(actual) | set(expected)):
|
||||
if actual.get(key) != expected.get(key):
|
||||
differences.append(
|
||||
f"{field}.{key}: expected {expected.get(key)!r}, "
|
||||
f"got {actual.get(key)!r}"
|
||||
)
|
||||
|
||||
|
||||
def compare_entries(produced: dict, golden: dict) -> list[str]:
|
||||
"""Return one line per field that differs, empty when the two agree.
|
||||
|
||||
Configuration, resolved state, tokens, trace,
|
||||
num_scored_tokens and is_watermarked are compared exactly; score, p_value
|
||||
and p_value_ratio are compared with GOLDEN_FLOAT_RTOL.
|
||||
"""
|
||||
differences: list[str] = []
|
||||
for field in ("configuration", "resolved", "trace"):
|
||||
_compare_mapping(field, produced[field], golden[field], differences)
|
||||
|
||||
if produced["generation"] != golden["generation"]:
|
||||
first_difference = next(
|
||||
(
|
||||
index
|
||||
for index, (actual, expected) in enumerate(
|
||||
zip(produced["generation"], golden["generation"])
|
||||
)
|
||||
if actual != expected
|
||||
),
|
||||
min(len(produced["generation"]), len(golden["generation"])),
|
||||
)
|
||||
differences.append(
|
||||
f"generation: first differs at index {first_difference} "
|
||||
f"({len(produced['generation'])} tokens generated, "
|
||||
f"{len(golden['generation'])} in the golden): "
|
||||
f"expected {golden['generation']}, got {produced['generation']}"
|
||||
)
|
||||
|
||||
detection = produced["detection"]
|
||||
expected_detection = golden["detection"]
|
||||
for field in ("score", "p_value", "p_value_ratio"):
|
||||
_compare_float(
|
||||
f"detection.{field}",
|
||||
detection[field],
|
||||
expected_detection[field],
|
||||
differences,
|
||||
)
|
||||
for field in ("num_scored_tokens", "is_watermarked"):
|
||||
_compare_exact(
|
||||
f"detection.{field}",
|
||||
detection[field],
|
||||
expected_detection[field],
|
||||
differences,
|
||||
)
|
||||
return differences
|
||||
|
||||
|
||||
def compare_golden(candidate: WatermarkingCandidate, golden: dict) -> list[str]:
|
||||
"""Return one line per field of ``candidate`` that drifted from ``golden``."""
|
||||
return compare_entries(candidate.golden_entry(golden["generation"]), golden)
|
||||
|
||||
|
||||
def read_goldens(path: Path) -> dict[str, dict]:
|
||||
"""Read the goldens file and return its candidate mapping."""
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
version = payload.get("schema_version")
|
||||
if version != GOLDEN_SCHEMA_VERSION:
|
||||
raise GoldenFormatError(
|
||||
f"{path}: schema_version {version!r}, expected {GOLDEN_SCHEMA_VERSION}"
|
||||
)
|
||||
return payload["candidates"]
|
||||
|
||||
|
||||
def configured_algorithms() -> set[str]:
|
||||
return set(get_args(WatermarkingAlgorithm))
|
||||
|
||||
|
||||
def configured_prfs() -> set[str]:
|
||||
return set(get_args(WatermarkPRFName))
|
||||
|
||||
|
||||
def configured_algorithm_prf_combinations() -> set[tuple[str, str]]:
|
||||
return {
|
||||
(algorithm, prf)
|
||||
for algorithm in configured_algorithms()
|
||||
for prf in configured_prfs()
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Check the watermarking goldens still reproduce."""
|
||||
|
||||
import dataclasses
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.watermarking.golden_candidates import (
|
||||
DETECTOR_FACTORIES,
|
||||
GOLDEN_FLOAT_RTOL,
|
||||
WATERMARK_CONFIG_FIELDS,
|
||||
WATERMARKING_CANDIDATES,
|
||||
compare_golden,
|
||||
configured_algorithm_prf_combinations,
|
||||
configured_algorithms,
|
||||
read_goldens,
|
||||
routing_boundary,
|
||||
validate_golden_guards,
|
||||
)
|
||||
from vllm.config.watermarking import WatermarkConfig
|
||||
|
||||
GOLDENS_PATH = Path(__file__).with_name("watermarking_goldens.json")
|
||||
|
||||
CANDIDATES_BY_ID = {candidate.id: candidate for candidate in WATERMARKING_CANDIDATES}
|
||||
|
||||
NEGATIVE_CANDIDATE_ID = "gumbel-philox-key42-cw4-wrong-key"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def goldens() -> dict[str, dict]:
|
||||
return read_goldens(GOLDENS_PATH)
|
||||
|
||||
|
||||
def test_goldens_cover_every_candidate(goldens: dict[str, dict]):
|
||||
candidate_ids = set(CANDIDATES_BY_ID)
|
||||
|
||||
assert len(candidate_ids) == len(WATERMARKING_CANDIDATES)
|
||||
assert set(goldens) == candidate_ids, "golden contract does not match candidates"
|
||||
candidate_combinations = {
|
||||
(candidate.scheme, candidate.prf) for candidate in WATERMARKING_CANDIDATES
|
||||
}
|
||||
assert candidate_combinations == configured_algorithm_prf_combinations()
|
||||
assert set(DETECTOR_FACTORIES) == configured_algorithms()
|
||||
|
||||
|
||||
def test_golden_fixture_guards(goldens: dict[str, dict]):
|
||||
validate_golden_guards(goldens)
|
||||
|
||||
|
||||
def test_resolved_records_every_watermark_config_field():
|
||||
assert {field.name for field in dataclasses.fields(WatermarkConfig)} == (
|
||||
set(WATERMARK_CONFIG_FIELDS) | {"key"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"candidate",
|
||||
WATERMARKING_CANDIDATES,
|
||||
ids=lambda candidate: candidate.id,
|
||||
)
|
||||
def test_watermarking_candidate_golden(candidate, goldens: dict[str, dict]):
|
||||
differences = compare_golden(candidate, goldens[candidate.id])
|
||||
|
||||
assert not differences, "\n".join(
|
||||
[
|
||||
f"{candidate.id}: golden drift "
|
||||
f"(tolerance {GOLDEN_FLOAT_RTOL:g} on score and p_value)",
|
||||
*differences,
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_negative_row_is_not_watermarked(goldens: dict[str, dict]):
|
||||
candidate = CANDIDATES_BY_ID[NEGATIVE_CANDIDATE_ID]
|
||||
|
||||
assert candidate.detection_key != candidate.key
|
||||
assert goldens[NEGATIVE_CANDIDATE_ID]["detection"]["is_watermarked"] is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"candidate",
|
||||
[
|
||||
candidate
|
||||
for candidate in WATERMARKING_CANDIDATES
|
||||
if candidate.scheme == "dual_key_gumbel"
|
||||
and 0 < candidate.scheme_config.generation_alpha < 1
|
||||
],
|
||||
ids=lambda candidate: candidate.id,
|
||||
)
|
||||
def test_routing_sequence_resolves_alpha(candidate):
|
||||
alpha = candidate.scheme_config.generation_alpha
|
||||
uniforms = candidate.fixture.routing_uniforms
|
||||
boundary = routing_boundary(alpha)
|
||||
share = sum(1 for uniform in uniforms if boundary <= uniform) / len(uniforms)
|
||||
|
||||
assert abs(share - alpha) <= 0.03
|
||||
@@ -276,9 +276,9 @@ def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
|
||||
"vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append
|
||||
)
|
||||
|
||||
sampler.add_request(0, 1, SamplingParams(temperature=0))
|
||||
sampler.add_request(0, 1, SamplingParams(temperature=1))
|
||||
sampler.add_request(0, 1, SamplingParams(temperature=0, watermarking=False))
|
||||
sampler.add_request(0, SamplingParams(temperature=0))
|
||||
sampler.add_request(0, SamplingParams(temperature=1))
|
||||
sampler.add_request(0, SamplingParams(temperature=0, watermarking=False))
|
||||
|
||||
assert messages == [
|
||||
(
|
||||
@@ -945,6 +945,7 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):
|
||||
speculator = object.__new__(StubSpeculator)
|
||||
speculator.model = StubModel()
|
||||
speculator.use_fp64_gumbel = False
|
||||
speculator.acceptance_estimator = None
|
||||
draft_watermarker = object.__new__(DraftWatermarker)
|
||||
draft_watermarker.watermarker = StubWatermarker()
|
||||
draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]])
|
||||
@@ -978,6 +979,7 @@ def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
|
||||
speculator.seeds = torch.zeros(2, dtype=torch.int64)
|
||||
speculator._step_cols = torch.tensor([0])
|
||||
speculator.use_fp64_gumbel = False
|
||||
speculator.acceptance_estimator = None
|
||||
watermark_logits: list[torch.Tensor] = []
|
||||
|
||||
def sample(logits, sampled, idx_map, temperature):
|
||||
@@ -1012,6 +1014,7 @@ def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch
|
||||
speculator.seeds = torch.zeros(2, dtype=torch.int64)
|
||||
speculator._step_cols = torch.tensor([0])
|
||||
speculator.use_fp64_gumbel = False
|
||||
speculator.acceptance_estimator = None
|
||||
speculator.draft_watermarker = None
|
||||
monkeypatch.setattr(
|
||||
"vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user