Skip to content

FEAT GCG: add model and tokenizer adapter contracts with control round-trip validation - #3618

Open
devangpratap wants to merge 1 commit into
microsoft:mainfrom
devangpratap:feat-gcg-model-tokenizer-adapters
Open

devangpratap wants to merge 1 commit into
microsoft:mainfrom
devangpratap:feat-gcg-model-tokenizer-adapters

Conversation

@devangpratap

Copy link
Copy Markdown
Contributor

Description

Adds explicit model and tokenizer contracts for GCG in a new model_adapters.py, read only through standard Hugging Face interfaces:

  • GCGModelAdapter: the input embedding matrix, its device, dtype and row count, and whether forward supports selective logits (logits_to_keep) and a prefix cache (past_key_values + use_cache). validate() holds the existing missing and non-float embedding errors.
  • GCGTokenizerAdapter: base vocabulary size, a stable vocabulary identity (SHA-256 of the sorted (id, token) pairs below the embedding row count), and validate() with the existing fast-tokenizer error.

These now back the existing code paths: _check_gcg_capabilities validates through both adapters, loss() reads the capability flags from the model adapter, and GCGCandidateProposer groups worker gradients by vocabulary identity instead of comparing full vocab dicts. The identity is computed once per worker and embedding size, not on every step.

The preflight also gets a round-trip check: GCG decodes each accepted candidate the way get_filtered_cands does and renders it again through control_str, so the initial control must come back as the same tokens. If it does not, the run fails before worker startup with the control and what it turned into, instead of silently optimizing different tokens. The check restores the prompt afterwards. With #3581 in, the default control_init round-trips and the identity is stable on the real tokenizers of Gemma 2, Gemma 3, Qwen2.5, SmolLM2, DeepSeek-R1-Distill-Qwen, Llama 3.2, Phi-4, TinyLlama, Phi-3 and Mistral v0.3.

Existing errors are unchanged. The one new failure is a control_init that does not encode back to its own tokens, which now raises before workers start.

PR 6 of the #2489 series. The compatibility matrix and family tests are next.

Tests and Documentation

  • TestModelAdapters: the embedding contract and capability flags on a tiny Llama, the flags on a forward without cache or logits_to_keep, identity ignoring added tokens past the embedding rows, and identity differing between vocabularies.
  • test_rejects_control_that_does_not_round_trip: a tokenizer whose decoder drops spaces turns "hello world" into "helloworld", and the preflight rejects it.
  • test_round_trip_check_leaves_the_prompt_unchanged: a control that passes keeps its control, input_ids and control slice.
  • test_vocabularies_are_read_once_per_worker_across_steps: three proposer steps read each worker's vocabulary once.
  • uv run pytest tests/unit/executor/promptgen/gcg/: 454 passed, 0 skipped.
  • pre-commit run on the changed files passes.
  • No documentation changes here; the compatibility matrix PR adds the GCG docs.

…d-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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant