From 50efe6f19ec8029f657f667c2f0725b8d55b5448 Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:31:10 -0400 Subject: [PATCH] FEAT GCG: add model and tokenizer adapter contracts with control round-trip validation GCGModelAdapter exposes the input embedding matrix, its device, dtype and row count, and whether forward supports logits_to_keep and a prefix cache. GCGTokenizerAdapter exposes the base vocabulary size and a stable vocabulary identity. The preflight validates through both, the proposer groups gradients by vocabulary identity, and the preflight now rejects an initial control that does not encode back to its own tokens. --- .../gcg/attack/base/attack_manager.py | 64 ++++----- .../gcg/attack/base/model_adapters.py | 128 ++++++++++++++++++ .../gcg/attack/gcg/candidate_proposer.py | 12 +- .../executor/promptgen/gcg/test_gcg_core.py | 79 +++++++++++ .../promptgen/gcg/test_gcg_proposal.py | 21 +++ 5 files changed, 268 insertions(+), 36 deletions(-) create mode 100644 pyrit/executor/promptgen/gcg/attack/base/model_adapters.py diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index b0a70d57be..95330a5d80 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/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,28 +2402,14 @@ 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. @@ -2437,17 +2417,33 @@ def _check_gcg_capabilities( 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: diff --git a/pyrit/executor/promptgen/gcg/attack/base/model_adapters.py b/pyrit/executor/promptgen/gcg/attack/base/model_adapters.py new file mode 100644 index 0000000000..69b7c33e68 --- /dev/null +++ b/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." + ) diff --git a/pyrit/executor/promptgen/gcg/attack/gcg/candidate_proposer.py b/pyrit/executor/promptgen/gcg/attack/gcg/candidate_proposer.py index e78043e6b4..13b95edfb4 100644 --- a/pyrit/executor/promptgen/gcg/attack/gcg/candidate_proposer.py +++ b/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,7 +117,7 @@ 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): @@ -122,7 +125,12 @@ def propose_candidates( 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) diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index d0419df180..fd23ec71f7 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/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: # and (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. diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_proposal.py b/tests/unit/executor/promptgen/gcg/test_gcg_proposal.py index 831c7146f6..925eb53b62 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_proposal.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_proposal.py @@ -248,6 +248,27 @@ def test_gradients_are_summed_only_for_the_same_base_vocabulary( expected_norm = 3 - len(group_worker_indices) assert torch.allclose(norms, torch.full_like(norms, float(expected_norm))) + def test_vocabularies_are_read_once_per_worker_across_steps(self) -> None: + workers = [_create_mock_worker((1, 3, 4)), _create_mock_worker((1, 3, 4))] + sampling = MagicMock(spec=SamplingStrategy) + sampling.sample_candidates.return_value = torch.tensor([[10]]) + candidate_filter = MagicMock(spec=CandidateFilter) + candidate_filter.filter_candidates.return_value = ["res"] + proposer = GCGCandidateProposer( + workers=workers, + prompts=[_create_mock_prompt_manager(), _create_mock_prompt_manager()], + sampling=sampling, + candidate_filter=candidate_filter, + main_device=torch.device("cpu"), + ) + + for _ in range(3): + batch = proposer.propose_candidates(batch_size=1, topk=2, current_control_str="test") + assert batch.group_worker_indices == [1] + + for worker in workers: + worker.tokenizer.get_vocab.assert_called_once_with() + @pytest.mark.parametrize( ("added_tokens", "embedding_rows", "representatives", "expected_candidates"),