Repository navigation
FEAT add GCG capability preflight before workers start - #3338
devangpratap wants to merge 3 commits into
Conversation
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 microsoft#2489
| 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." | ||
| ) |
There was a problem hiding this comment.
This check rejects Gemma 3 text checkpoints (e.g. gemma-3-1b-it), a family #2489 targets. Their tokenizer has 262145 entries for 262144 embedding rows. The extra entry is the image placeholder token (image_soft_token), which the chat template never emits.
len(tokenizer) counts every added token, but GCG only uses chat-template-rendered ids, gradient-sampled candidates bounded by the rows, and get_nonascii_toks, which covers range(3, vocab_size) plus bos/eos/pad/unk. With the real tokenizer and a randomly initialized Gemma3ForCausalLM, all ids stayed in range. The error's remedies don't apply: the tokenizer does match, and GCGGenerator users never hold the model to call resize_token_embeddings.
Suggest checking tokenizer.vocab_size <= weight.shape[0], and that bos/eos/pad/unk and the ids of a probe conversation rendered as in _update_ids are below it. Avoid all_special_ids; it includes Gemma's image token. Please add a regression test with an unused added special token beyond the rows.
There was a problem hiding this comment.
Thanks, switched to your suggestion: vocab_size, the four special ids, and a probe conversation rendered as in _update_ids. Checked it on the real gemma-3-1b-it tokenizer. The regression test uses the test tokenizer's two unused added tokens past the rows.
| 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. | ||
| """ |
There was a problem hiding this comment.
This makes the normal GCGModelConfig(name=...) path fail every time because _default_tokenizer_kwargs() still returns {"use_fast": False}. For example, the minimal config shown in config.py passes that value to AutoTokenizer, then this preflight rejects the resulting tokenizer before workers start. Could we change the GCG default to use_fast=True (and update the stale Llama/Vicuna explanation) and add a test that builds params from the default model config? Otherwise this PR's new requirement is incompatible with the public API defaults.
There was a problem hiding this comment.
Changed the default to use_fast=True, updated the docstring, and the worker start test now loads a real tokenizer with the default GCGModelConfig kwargs. One note: transformers 5 ignores use_fast in AutoTokenizer, so the old default still gave fast tokenizers (I checked Llama-2, Vicuna, TinyLlama and Gemma 3 with use_fast=False). The change mainly makes the default match what GCG needs.
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.
This branch has not been deployed
3 participants
Description
Contributes to #2489 (the capability preflight piece).
get_workersnow checks every model and tokenizer pair after loading and before any worker process starts. Each failure is aValueErrorthat names the model and says what to change:get_input_embeddings()is missing, raisesNotImplementedError, or has no weight tensoradd_tokenswithoutresize_token_embeddingsUntil now, a bad embedding matrix or a token id past its rows only failed inside a worker subprocess once optimization started. The worker exited and the parent waited on its result queue forever. The tokenizer requirement was already enforced by
AttackPrompt, but only after every model was loaded and every worker started. That check in_update_idsstays, sinceAttackPromptcan be built withoutget_workers.The vocabulary check is the minimum the preflight needs; the adapter contract piece can extend it.
Tests and Documentation
TestCheckGcgCapabilities: a tiny random Llama with the offline fast tokenizer passes, as does a tokenizer that exactly fills the embedding matrix; a missing, raising,int8, and too-small embedding each raise with the expected message.test_capability_check_runs_before_workers_start: a failing pair raises fromget_workersandworker.start()is never called.uv run pytest tests/unit/executor/promptgen/gcg/ -q: 410 passed, 0 skipped.Qwen/Qwen2.5-0.5B-InstructandHuggingFaceTB/SmolLM2-135M-Instruct: both pass, and both raise the vocabulary error after adding tokens past the embedding size.pre-commit runon the changed files passes.