From d86d5aef10bdcf64c0340fc5dadf3a4b1f1ec39f Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:49:49 -0400 Subject: [PATCH 1/2] FEAT add GCG capability preflight before workers start Check each model and tokenizer pair in get_workers before any worker process starts: a tokenizers-backed tokenizer, a floating point input embedding matrix from get_input_embeddings(), and a tokenizer vocabulary that fits inside it. Each failure raises a ValueError naming the model and the fix. Before, a bad embedding matrix failed inside a worker subprocess and the parent hung on its result queue. Contributes to #2489 --- .../gcg/attack/base/attack_manager.py | 46 +++++++++++++- .../executor/promptgen/gcg/test_gcg_core.py | 61 ++++++++++++++++++- 2 files changed, 105 insertions(+), 2 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 1b46944e7f..56417a4c91 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -2390,6 +2390,47 @@ def _execute_task(*, model: Any, task: ModelWorkerTask) -> Any: return task.operation(*task.args, **task.kwargs) +def _check_gcg_capabilities(*, model: Any, tokenizer: Any, model_path: str) -> None: + """ + Check that GCG can optimize against a model and tokenizer before any worker starts. + + Args: + model (Any): The loaded causal language model. + tokenizer (Any): The tokenizer paired with the model. + model_path (str): The model path, used in error messages. + + Raises: + ValueError: If the tokenizer is not fast, the model exposes no floating point + input embedding matrix, or the tokenizer has ids outside that matrix. + """ + 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." + ) + if len(tokenizer) > weight.shape[0]: + raise ValueError( + f"The tokenizer for {model_path!r} has {len(tokenizer)} tokens but the model's input embeddings " + f"have only {weight.shape[0]} rows. Check that the tokenizer matches the model, or call " + "model.resize_token_embeddings(len(tokenizer)) after adding tokens." + ) + + def get_workers(params: Any, evaluation: bool = False) -> tuple[list[ModelWorker], list[ModelWorker]]: """ Load and optionally start training and held-out model workers. @@ -2398,7 +2439,8 @@ def get_workers(params: Any, evaluation: bool = False) -> tuple[list[ModelWorker tuple[list[ModelWorker], list[ModelWorker]]: Training and held-out workers. Raises: - ValueError: If a tokenizer does not define a chat template. + ValueError: If a tokenizer does not define a chat template, or a model and + tokenizer pair fails the GCG capability check. """ tokenizers = [] for i in range(len(params.tokenizer_paths)): @@ -2447,6 +2489,8 @@ def get_workers(params: Any, evaluation: bool = False) -> tuple[list[ModelWorker ) for i in range(len(params.model_paths)) ] + for model_path, worker in zip(params.model_paths, workers, strict=True): + _check_gcg_capabilities(model=worker.model, tokenizer=worker.tokenizer, model_path=model_path) if not evaluation: for worker in workers: worker.start() diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index 2e622e56d3..91670a7297 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path +from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock, call, patch, sentinel @@ -1086,7 +1087,9 @@ def test_starts_workers_for_training(self) -> None: tokenizer = MagicMock() tokenizer.pad_token = "" tokenizer.chat_template = "{{ messages[0]['content'] }}" - worker = MagicMock() + tokenizer.__len__.return_value = 8 + worker = MagicMock(tokenizer=tokenizer) + worker.model.get_input_embeddings.return_value = torch.nn.Embedding(8, 4) with ( patch.object(attack_manager_mod.AutoTokenizer, "from_pretrained", return_value=tokenizer), @@ -1098,6 +1101,62 @@ def test_starts_workers_for_training(self) -> None: assert train_workers == [worker] assert test_workers == [] + def test_capability_check_runs_before_workers_start(self) -> None: + params = MagicMock() + params.tokenizer_paths = ["fake/chat-model"] + params.tokenizer_kwargs = [{}] + params.model_paths = ["fake/chat-model"] + params.model_kwargs = [{}] + params.devices = ["cpu"] + params.token = "" + + tokenizer = MagicMock() + tokenizer.pad_token = "" + tokenizer.chat_template = "{{ messages[0]['content'] }}" + tokenizer.is_fast = False + worker = MagicMock(tokenizer=tokenizer) + + with ( + patch.object(attack_manager_mod.AutoTokenizer, "from_pretrained", return_value=tokenizer), + patch.object(attack_manager_mod, "ModelWorker", return_value=worker), + pytest.raises(ValueError, match="'fake/chat-model' .* is not backed by the Hugging Face tokenizers"), + ): + attack_manager_mod.get_workers(params, evaluation=False) + + worker.start.assert_not_called() + + +class TestCheckGcgCapabilities: + 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_accepts_tokenizer_that_exactly_fills_embeddings(self) -> None: + model = MagicMock() + model.get_input_embeddings.return_value = torch.nn.Embedding(len(_fast_tokenizer()), 4) + attack_manager_mod._check_gcg_capabilities(model=model, tokenizer=_fast_tokenizer(), model_path="fake/model") + + @pytest.mark.parametrize( + ("embeddings", "match"), + [ + (NotImplementedError(), "does not expose an input embedding matrix"), + (None, "does not expose an input embedding matrix"), + (SimpleNamespace(weight=torch.zeros(256, 4, dtype=torch.int8)), "are torch.int8"), + (torch.nn.Embedding(4, 4), "has 27 tokens but the model's input embeddings have only 4 rows"), + ], + ) + def test_rejects_unusable_embeddings(self, embeddings: Any, match: str) -> None: + model = MagicMock() + if isinstance(embeddings, Exception): + model.get_input_embeddings.side_effect = embeddings + else: + model.get_input_embeddings.return_value = embeddings + with pytest.raises(ValueError, match=match): + attack_manager_mod._check_gcg_capabilities( + model=model, tokenizer=_fast_tokenizer(), model_path="fake/model" + ) + def test_model_worker_uses_model_device_dispatch() -> None: model = MagicMock() From 6926922445c952133a079b662b81c56283773034 Mon Sep 17 00:00:00 2001 From: devangpratap Date: Tue, 20 Oct 2026 09:24:20 -0400 Subject: [PATCH 2/2] Check only token ids GCG uses and default to use_fast=True The vocabulary check compared len(tokenizer) to the embedding rows, which rejects Gemma 3 text checkpoints: their image placeholder is an added token past the rows that the chat template never emits. Check the base vocabulary, the special token ids, and a probe conversation rendered as in _update_ids instead. Default GCGModelConfig.tokenizer_kwargs to use_fast=True so the default matches the fast tokenizer requirement, and test get_workers with the default model config. --- .../gcg/attack/base/attack_manager.py | 20 ++++++--- pyrit/executor/promptgen/gcg/config.py | 7 ++-- .../executor/promptgen/gcg/test_config.py | 2 +- .../executor/promptgen/gcg/test_gcg_core.py | 42 ++++++++++--------- 4 files changed, 41 insertions(+), 30 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 56417a4c91..6baeb445af 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -2401,7 +2401,7 @@ def _check_gcg_capabilities(*, model: Any, tokenizer: Any, model_path: str) -> N Raises: ValueError: If the tokenizer is not fast, the model exposes no floating point - input embedding matrix, or the tokenizer has ids outside that matrix. + input embedding matrix, or GCG would use token ids outside that matrix. """ if not tokenizer.is_fast: raise ValueError( @@ -2423,11 +2423,21 @@ def _check_gcg_capabilities(*, model: Any, tokenizer: Any, model_path: str) -> N 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." ) - if len(tokenizer) > weight.shape[0]: + # 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(prompt, add_special_tokens=False).input_ids, + *(i for i in special_ids if i is not None), + ) + if max_id >= weight.shape[0]: raise ValueError( - f"The tokenizer for {model_path!r} has {len(tokenizer)} tokens but the model's input embeddings " - f"have only {weight.shape[0]} rows. Check that the tokenizer matches the model, or call " - "model.resize_token_embeddings(len(tokenizer)) after adding tokens." + 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." ) diff --git a/pyrit/executor/promptgen/gcg/config.py b/pyrit/executor/promptgen/gcg/config.py index 483f6b389f..dd7bf2e667 100644 --- a/pyrit/executor/promptgen/gcg/config.py +++ b/pyrit/executor/promptgen/gcg/config.py @@ -41,7 +41,7 @@ def _default_model_kwargs() -> dict[str, Any]: def _default_tokenizer_kwargs() -> dict[str, Any]: - return {"use_fast": False} + return {"use_fast": True} @dataclass @@ -58,9 +58,8 @@ class GCGModelConfig: ``AutoModelForCausalLM.from_pretrained``. Defaults to ``{"low_cpu_mem_usage": True, "use_cache": False}``. tokenizer_kwargs (dict[str, Any]): Extra keyword arguments forwarded to - ``AutoTokenizer.from_pretrained``. Defaults to ``{"use_fast": False}`` - because several legacy chat tokenizers (Llama-2, Vicuna) ship slow - tokenizers only. + ``AutoTokenizer.from_pretrained``. Defaults to ``{"use_fast": True}`` + because GCG needs the character offsets only fast tokenizers provide. """ name: str diff --git a/tests/unit/executor/promptgen/gcg/test_config.py b/tests/unit/executor/promptgen/gcg/test_config.py index abae513cab..76c2c75027 100644 --- a/tests/unit/executor/promptgen/gcg/test_config.py +++ b/tests/unit/executor/promptgen/gcg/test_config.py @@ -81,7 +81,7 @@ def test_minimal_config_constructs_with_defaults() -> None: assert config.models[0].name == _LLAMA_2 assert config.models[0].device == "cuda:0" assert config.models[0].model_kwargs == {"low_cpu_mem_usage": True, "use_cache": False} - assert config.models[0].tokenizer_kwargs == {"use_fast": False} + assert config.models[0].tokenizer_kwargs == {"use_fast": True} assert config.test_models == [] assert config.algorithm.n_steps == 500 assert config.algorithm.batch_size == 512 diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index 91670a7297..e914ef8a02 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -34,6 +34,7 @@ get_embedding_matrix, get_embeddings, ) +from pyrit.executor.promptgen.gcg.config import GCGModelConfig # noqa: E402 gcg_attack_mod = pytest.importorskip( "pyrit.executor.promptgen.gcg.attack.gcg.gcg_attack", @@ -1074,29 +1075,28 @@ def test_raises_when_tokenizer_has_no_chat_template(self) -> None: with pytest.raises(ValueError, match="no chat_template configured"): get_workers(params) - def test_starts_workers_for_training(self) -> None: + def test_starts_workers_for_training_with_default_model_config(self, tmp_path: Path) -> None: + _fast_tokenizer().save_pretrained(tmp_path) + model_config = GCGModelConfig(name=str(tmp_path), device="cpu") params = MagicMock() - params.tokenizer_paths = ["fake/chat-model"] - params.tokenizer_kwargs = [{}] - params.model_paths = ["fake/chat-model"] - params.model_kwargs = [{}] - params.devices = ["cpu"] - params.token = "" + params.tokenizer_paths = [model_config.name] + params.tokenizer_kwargs = [model_config.tokenizer_kwargs] + params.model_paths = [model_config.name] + params.model_kwargs = [model_config.model_kwargs] + params.devices = [model_config.device] + params.token = None params.num_train_models = 1 - tokenizer = MagicMock() - tokenizer.pad_token = "" - tokenizer.chat_template = "{{ messages[0]['content'] }}" - tokenizer.__len__.return_value = 8 - worker = MagicMock(tokenizer=tokenizer) - worker.model.get_input_embeddings.return_value = torch.nn.Embedding(8, 4) + worker = MagicMock(model=_tiny_model("llama")) - with ( - patch.object(attack_manager_mod.AutoTokenizer, "from_pretrained", return_value=tokenizer), - patch.object(attack_manager_mod, "ModelWorker", return_value=worker), - ): + def build_worker(model_path: str, token: Any, model_kwargs: Any, tokenizer: Any, device: str) -> MagicMock: + worker.tokenizer = tokenizer + return worker + + with patch.object(attack_manager_mod, "ModelWorker", side_effect=build_worker): train_workers, test_workers = attack_manager_mod.get_workers(params, evaluation=False) + assert worker.tokenizer.is_fast worker.start.assert_called_once_with() assert train_workers == [worker] assert test_workers == [] @@ -1132,9 +1132,11 @@ def test_accepts_tiny_model_with_fast_tokenizer(self) -> None: model=_tiny_model("llama"), tokenizer=_fast_tokenizer(), model_path="fake/model" ) - def test_accepts_tokenizer_that_exactly_fills_embeddings(self) -> None: + 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. model = MagicMock() - model.get_input_embeddings.return_value = torch.nn.Embedding(len(_fast_tokenizer()), 4) + model.get_input_embeddings.return_value = torch.nn.Embedding(25, 4) attack_manager_mod._check_gcg_capabilities(model=model, tokenizer=_fast_tokenizer(), model_path="fake/model") @pytest.mark.parametrize( @@ -1143,7 +1145,7 @@ def test_accepts_tokenizer_that_exactly_fills_embeddings(self) -> None: (NotImplementedError(), "does not expose an input embedding matrix"), (None, "does not expose an input embedding matrix"), (SimpleNamespace(weight=torch.zeros(256, 4, dtype=torch.int8)), "are torch.int8"), - (torch.nn.Embedding(4, 4), "has 27 tokens but the model's input embeddings have only 4 rows"), + (torch.nn.Embedding(24, 4), "produces token id 24 but the model's input embeddings have only 24 rows"), ], ) def test_rejects_unusable_embeddings(self, embeddings: Any, match: str) -> None: