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"),