FEAT add GCG capability preflight before workers start

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

@@ -2390,6 +2390,57 @@ 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 GCG would use token 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."
+        )
+    # 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} 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."
+        )
+
+
 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 +2449,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 +2499,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()

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

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

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
 
@@ -33,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",
@@ -1073,30 +1075,89 @@ 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 = [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
+
+        worker = MagicMock(model=_tiny_model("llama"))
+
+        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 == []
+
+    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 = ""
-        params.num_train_models = 1
 
         tokenizer = MagicMock()
         tokenizer.pad_token = "<pad>"
         tokenizer.chat_template = "{{ messages[0]['content'] }}"
-        worker = MagicMock()
+        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"),
         ):
-            train_workers, test_workers = attack_manager_mod.get_workers(params, evaluation=False)
+            attack_manager_mod.get_workers(params, evaluation=False)
 
-        worker.start.assert_called_once_with()
-        assert train_workers == [worker]
-        assert test_workers == []
+        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_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.
+        model = MagicMock()
+        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(
+        ("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(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:
+        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: