FEAT GCG: add model and tokenizer adapter contracts with control round-trip validation

pyrit/executor/promptgen/gcg/attack/base/attack_manager.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import inspect
 import json
 import logging
 import math
@@ -26,6 +25,7 @@
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
 from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
 
+from pyrit.executor.promptgen.gcg.attack.base.model_adapters import GCGModelAdapter, GCGTokenizerAdapter
 from pyrit.executor.promptgen.gcg.attack.base.progressive_schedule import (
     ProgressiveScheduleController,
     ProgressiveScheduleState,
@@ -785,21 +785,15 @@ def loss(
             torch.Tensor: One scalar loss per candidate control.
         """
         selective_loss = cast("Any", loss_function)
-        try:
-            forward_parameters = inspect.signature(model.forward).parameters
-        except (TypeError, ValueError):
-            forward_parameters = {}
+        model_adapter = GCGModelAdapter(model)
 
-        supports_selective_logits = "logits_to_keep" in forward_parameters
-        supports_prefix_cache = "past_key_values" in forward_parameters and "use_cache" in forward_parameters
-
-        if supports_selective_logits and hasattr(selective_loss, "get_required_logit_positions"):
+        if model_adapter.supports_selective_logits and hasattr(selective_loss, "get_required_logit_positions"):
             logit_positions = selective_loss.get_required_logit_positions(
                 target_slice=self._target_slice,
                 control_slice=self._control_slice,
                 device=model.device,
             )
-            if use_prefix_cache and supports_prefix_cache and self._control_slice.start > 1:
+            if use_prefix_cache and model_adapter.supports_prefix_cache and self._control_slice.start > 1:
                 cached_loss = self._loss_with_prefix_cache(
                     model,
                     test_controls,
@@ -2408,46 +2402,48 @@ def _check_gcg_capabilities(
 
     Raises:
         ValueError: If the tokenizer is not fast, the model exposes no floating point
-            input embedding matrix, or GCG would use token ids outside that matrix.
+            input embedding matrix, GCG would use token ids outside that matrix, or an
+            initial control does not survive the decode/encode round trip GCG applies every step.
     """
-    if not tokenizer.is_fast:
-        raise ValueError(
-            f"The tokenizer for {model_path!r} ({type(tokenizer).__name__}) is not backed by the Hugging Face "
-            "tokenizers library, so it cannot return the character-to-token offsets GCG uses to locate the goal, "
-            "control and target. Load it with the default tokenizers backend."
-        )
-    try:
-        weight = getattr(model.get_input_embeddings(), "weight", None)
-    except NotImplementedError:
-        weight = None
-    if not isinstance(weight, torch.Tensor):
-        raise ValueError(
-            f"Model {model_path!r} does not expose an input embedding matrix through get_input_embeddings(). "
-            "GCG computes token gradients through that matrix."
-        )
-    if not weight.is_floating_point():
-        raise ValueError(
-            f"The input embeddings of {model_path!r} are {weight.dtype}, so GCG cannot compute token gradients "
-            "through them. Load the model without quantizing the embedding layer."
-        )
+    tokenizer_adapter = GCGTokenizerAdapter(tokenizer)
+    tokenizer_adapter.validate(model_path=model_path)
+    model_adapter = GCGModelAdapter(model)
+    model_adapter.validate(model_path=model_path)
+    prompts = list(prompts)
     # Added tokens the chat template never emits (e.g. Gemma 3's image token) may sit past the
     # embedding rows, so check only the ids GCG uses: the base vocabulary sampled by
     # get_nonascii_toks, the special token ids, and a conversation rendered as in _update_ids.
     messages = [{"role": "user", "content": "Say it ! !"}, {"role": "assistant", "content": "Sure"}]
     prompt = tokenizer.apply_chat_template(messages, tokenize=False)
     special_ids = (tokenizer.bos_token_id, tokenizer.eos_token_id, tokenizer.pad_token_id, tokenizer.unk_token_id)
     max_id = max(
-        tokenizer.vocab_size - 1,
+        tokenizer_adapter.vocab_size - 1,
         *tokenizer(prompt, add_special_tokens=False).input_ids,
         *(i for i in special_ids if i is not None),
     )
     for attack_prompt in prompts:
         max_id = max(max_id, int(attack_prompt.input_ids.max()))
-    if max_id >= weight.shape[0]:
+    if max_id >= model_adapter.num_embeddings:
         raise ValueError(
             f"The tokenizer for {model_path!r} produces token id {max_id} but the model's input embeddings "
-            f"have only {weight.shape[0]} rows. Check that the tokenizer matches the model."
+            f"have only {model_adapter.num_embeddings} rows. Check that the tokenizer matches the model."
+        )
+    for attack_prompt in prompts:
+        # GCG decodes each accepted candidate as get_filtered_cands does and renders it again through
+        # control_str, so the initial control must come back as the same tokens.
+        control, control_toks = attack_prompt.control, attack_prompt.control_toks
+        attack_prompt.control_str = tokenizer.decode(
+            control_toks, skip_special_tokens=True, clean_up_tokenization_spaces=False
         )
+        round_tripped = attack_prompt.control
+        same_tokens = torch.equal(attack_prompt.control_toks, control_toks)
+        attack_prompt.control_str = control
+        if not same_tokens:
+            raise ValueError(
+                f"The initial control {control!r} does not survive a decode/encode round trip with the tokenizer "
+                f"for {model_path!r}: it comes back as {round_tripped!r}, so GCG would not optimize the "
+                "tokens it started from. Use a control_init that encodes back to the same tokens."
+            )
 
 
 def start_workers(*, params: Any, workers: list[ModelWorker]) -> None:

pyrit/executor/promptgen/gcg/attack/base/model_adapters.py

@@ -0,0 +1,128 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Model and tokenizer contracts GCG relies on, read through standard Hugging Face interfaces."""
+
+import hashlib
+import inspect
+import json
+from typing import Any
+
+import torch
+
+
+class GCGModelAdapter:
+    """The parts of a causal language model that GCG reads, without architecture-specific branches."""
+
+    def __init__(self, model: Any) -> None:
+        """
+        Wrap a loaded causal language model.
+
+        Args:
+            model (Any): The loaded causal language model.
+        """
+        self.model = model
+        try:
+            forward_parameters = inspect.signature(model.forward).parameters
+        except (TypeError, ValueError):
+            forward_parameters = {}
+        self.supports_selective_logits = "logits_to_keep" in forward_parameters
+        self.supports_prefix_cache = "past_key_values" in forward_parameters and "use_cache" in forward_parameters
+
+    @property
+    def embedding_matrix(self) -> torch.Tensor:
+        """The input embedding matrix GCG computes token gradients through."""
+        weight: torch.Tensor = self.model.get_input_embeddings().weight
+        return weight
+
+    @property
+    def device(self) -> torch.device:
+        """The device of the input embedding matrix."""
+        return self.embedding_matrix.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """The dtype of the input embedding matrix."""
+        return self.embedding_matrix.dtype
+
+    @property
+    def num_embeddings(self) -> int:
+        """The number of token ids the model can embed."""
+        return int(self.embedding_matrix.shape[0])
+
+    def validate(self, *, model_path: str) -> None:
+        """
+        Check that GCG can compute token gradients through the input embeddings.
+
+        Args:
+            model_path (str): The model path, used in error messages.
+
+        Raises:
+            ValueError: If the model exposes no floating point input embedding matrix.
+        """
+        try:
+            weight = self.embedding_matrix
+        except (AttributeError, NotImplementedError):
+            weight = None
+        if not isinstance(weight, torch.Tensor):
+            raise ValueError(
+                f"Model {model_path!r} does not expose an input embedding matrix through get_input_embeddings(). "
+                "GCG computes token gradients through that matrix."
+            )
+        if not weight.is_floating_point():
+            raise ValueError(
+                f"The input embeddings of {model_path!r} are {weight.dtype}, so GCG cannot compute token gradients "
+                "through them. Load the model without quantizing the embedding layer."
+            )
+
+
+class GCGTokenizerAdapter:
+    """The parts of a tokenizer that GCG reads when rendering prompts and sharing candidates across models."""
+
+    def __init__(self, tokenizer: Any) -> None:
+        """
+        Wrap a tokenizer.
+
+        Args:
+            tokenizer (Any): The tokenizer paired with a GCG model.
+        """
+        self.tokenizer = tokenizer
+
+    @property
+    def vocab_size(self) -> int:
+        """The size of the base vocabulary, without added tokens."""
+        return int(self.tokenizer.vocab_size)
+
+    def vocabulary_identity(self, *, num_embeddings: int) -> str:
+        """
+        Return a stable identity for the token ids a model with ``num_embeddings`` rows can sample.
+
+        Two tokenizers with the same identity map every sampleable id to the same token, so
+        their gradients can be summed. Added tokens past the embedding rows are ignored since
+        they can never be sampled.
+
+        Args:
+            num_embeddings (int): The number of rows in the paired model's input embedding matrix.
+
+        Returns:
+            str: A SHA-256 hex digest of the sampleable ``(token, id)`` pairs.
+        """
+        vocab = sorted((i, token) for token, i in self.tokenizer.get_vocab().items() if i < num_embeddings)
+        return hashlib.sha256(json.dumps(vocab).encode()).hexdigest()
+
+    def validate(self, *, model_path: str) -> None:
+        """
+        Check that the tokenizer can return the character offsets GCG uses to locate prompt slices.
+
+        Args:
+            model_path (str): The model path, used in error messages.
+
+        Raises:
+            ValueError: If the tokenizer is not backed by the Hugging Face tokenizers library.
+        """
+        if not self.tokenizer.is_fast:
+            raise ValueError(
+                f"The tokenizer for {model_path!r} ({type(self.tokenizer).__name__}) is not backed by the Hugging "
+                "Face tokenizers library, so it cannot return the character-to-token offsets GCG uses to locate the "
+                "goal, control and target. Load it with the default tokenizers backend."
+            )

pyrit/executor/promptgen/gcg/attack/gcg/candidate_proposer.py

@@ -13,6 +13,7 @@
     ModelWorkerOperation,
     PromptManager,
 )
+from pyrit.executor.promptgen.gcg.attack.base.model_adapters import GCGTokenizerAdapter
 from pyrit.executor.promptgen.gcg.extension_protocols import CandidateFilter, SamplingStrategy
 
 
@@ -78,6 +79,8 @@ def __init__(
         self._sample_fn = sample_fn
         self._filter_fn = filter_fn
         self._main_device = main_device
+        # Vocabulary identity per (worker index, embedding rows), so each vocabulary is hashed once per run.
+        self._vocabulary_identities: dict[tuple[int, int], str] = {}
 
     def propose_candidates(
         self,
@@ -114,15 +117,20 @@ def propose_candidates(
         control_cands: list[list[str]] = []
         group_worker_indices: list[int] = []
         grad: torch.Tensor | None = None
-        vocab: dict[str, int] | None = None
+        vocab: str | None = None
 
         # Collect and aggregate gradients across workers
         for j, worker in enumerate(self._workers):
             new_grad: torch.Tensor = worker.results.get().to(self._main_device)
             new_grad = new_grad / new_grad.norm(dim=-1, keepdim=True)
             # Each embedding row is a sampleable token ID, including added tokens.
             # Tokens outside the embedding matrix cannot be sampled and need not match.
-            new_vocab = {token: i for token, i in worker.tokenizer.get_vocab().items() if i < new_grad.shape[-1]}
+            key = (j, new_grad.shape[-1])
+            if key not in self._vocabulary_identities:
+                self._vocabulary_identities[key] = GCGTokenizerAdapter(worker.tokenizer).vocabulary_identity(
+                    num_embeddings=new_grad.shape[-1]
+                )
+            new_vocab = self._vocabulary_identities[key]
 
             if grad is None:
                 grad = torch.zeros_like(new_grad)

tests/unit/executor/promptgen/gcg/test_gcg_core.py

@@ -34,6 +34,10 @@
     get_embedding_matrix,
     get_embeddings,
 )
+from pyrit.executor.promptgen.gcg.attack.base.model_adapters import (  # noqa: E402
+    GCGModelAdapter,
+    GCGTokenizerAdapter,
+)
 from pyrit.executor.promptgen.gcg.config import GCGAlgorithmConfig, GCGModelConfig, GCGStrategyConfig  # noqa: E402
 from pyrit.executor.promptgen.gcg.extension_protocols import SuffixInitializer  # noqa: E402
 
@@ -1175,12 +1179,87 @@ def test_capability_check_runs_before_workers_start(self) -> None:
         worker.start.assert_not_called()
 
 
+class TestModelAdapters:
+    def test_model_adapter_exposes_the_embedding_contract(self) -> None:
+        model = _tiny_model("llama")
+        weight = model.get_input_embeddings().weight
+        adapter = GCGModelAdapter(model)
+
+        assert adapter.embedding_matrix is weight
+        assert (adapter.device, adapter.dtype, adapter.num_embeddings) == (weight.device, weight.dtype, 256)
+        assert adapter.supports_selective_logits
+        assert adapter.supports_prefix_cache
+
+    def test_model_adapter_reports_forward_without_cache_or_selective_logits(self) -> None:
+        model = MagicMock()
+        model.forward = lambda input_ids, attention_mask=None: None
+        adapter = GCGModelAdapter(model)
+
+        assert not adapter.supports_selective_logits
+        assert not adapter.supports_prefix_cache
+
+    @pytest.mark.parametrize(
+        ("added_token", "num_embeddings", "same_identity"),
+        [
+            (None, 27, True),
+            # new_domain_term gets id 27, so it is only sampleable with a 28th embedding row.
+            ("new_domain_term", 27, True),
+            ("new_domain_term", 28, False),
+        ],
+    )
+    def test_vocabulary_identity_covers_only_sampleable_ids(
+        self, added_token: str | None, num_embeddings: int, same_identity: bool
+    ) -> None:
+        other = _fast_tokenizer()
+        if added_token is not None:
+            other.add_tokens([added_token])
+
+        identity = GCGTokenizerAdapter(_fast_tokenizer()).vocabulary_identity(num_embeddings=num_embeddings)
+        other_identity = GCGTokenizerAdapter(other).vocabulary_identity(num_embeddings=num_embeddings)
+
+        assert (identity == other_identity) is same_identity
+
+    def test_vocabulary_identity_differs_for_a_different_vocabulary(self) -> None:
+        words = GCGTokenizerAdapter(_fast_tokenizer()).vocabulary_identity(num_embeddings=27)
+        bytes_ = GCGTokenizerAdapter(_fast_tokenizer(byte_level=True)).vocabulary_identity(num_embeddings=27)
+
+        assert words != bytes_
+
+
 class TestCheckGcgCapabilities:
+    def test_rejects_control_that_does_not_round_trip(self) -> None:
+        from tokenizers import decoders
+
+        tokenizer = _fast_tokenizer()
+        # Decoding without spaces turns the control "hello world" into the unknown word "helloworld".
+        tokenizer.backend_tokenizer.decoder = decoders.Fuse()
+        prompt = AttackPrompt(goal="Say it", target="Sure", tokenizer=tokenizer, control_init="hello world")
+
+        with pytest.raises(ValueError, match="'hello world' does not survive .* comes back as 'helloworld'"):
+            attack_manager_mod._check_gcg_capabilities(
+                model=_tiny_model("llama"), tokenizer=tokenizer, model_path="fake/model", prompts=[prompt]
+            )
+
     def test_accepts_tiny_model_with_fast_tokenizer(self) -> None:
         attack_manager_mod._check_gcg_capabilities(
             model=_tiny_model("llama"), tokenizer=_fast_tokenizer(), model_path="fake/model"
         )
 
+    def test_round_trip_check_leaves_the_prompt_unchanged(self) -> None:
+        tokenizer = _fast_tokenizer()
+        prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !")
+        control, input_ids, control_slice = prompt.control, prompt.input_ids.tolist(), prompt._control_slice
+
+        attack_manager_mod._check_gcg_capabilities(
+            model=_tiny_model("llama"), tokenizer=tokenizer, model_path="fake/model", prompts=[prompt]
+        )
+
+        assert (prompt.control, prompt.input_ids.tolist(), prompt._control_slice) == (
+            control,
+            input_ids,
+            control_slice,
+        )
+
     def test_accepts_unused_added_tokens_beyond_embeddings(self) -> None:
         # <start_of_turn> and <end_of_turn> (ids 25 and 26) are added tokens the ChatML
         # template never emits, like Gemma 3's image token; the rendered prompt tops out at 24.