FEAT Add generic GCG model adapters and an open-weight SLM compatibility matrix

Is your feature request related to a problem? Please describe.

GCG loads models through AutoModelForCausalLM, but its gradient path still hard-codes exact classes for GPT-2/GPT-J, GPT-NeoX, Llama, Mistral/Mixtral, and Phi-3. A newly supported Hugging Face causal model can therefore load successfully and then fail when GCG tries to obtain its input embeddings.

Prompt slicing also assumes a tokenizer with a usable chat template and character-to-token mapping. There is no explicit preflight check or maintained compatibility matrix, so unsupported tokenizer behavior may fail late or silently compute incorrect goal, control, or target slices.

We want GCG to support the leading open-weight small language model families without losing its multi-model, multi-prompt, held-out evaluation, or progressive optimization capabilities.

Describe the solution you'd like

Introduce small model and tokenizer adapter interfaces for GCG.

The model adapter should:

use the standard Hugging Face get_input_embeddings() contract instead of architecture-specific isinstance branches where possible;
expose the embedding matrix, device, dtype, vocabulary size, and relevant cache capabilities;
validate that the model supports differentiable causal-LM input-embedding optimization;
produce actionable capability errors before worker processes and long-running optimization begin.

The tokenizer adapter should:

render conversations using the tokenizer's chat template;
locate goal, control, assistant-role, and target token slices robustly;
validate decode/encode round trips and vocabulary compatibility;
expose a stable tokenizer/vocabulary identity that later multi-model aggregation work can use safely;
allow narrowly scoped family-specific behavior only when the standard Hugging Face interfaces are insufficient.

Create and document an initial compatibility matrix for representative causal text models/tokenizers from these open-weight SLM families: Qwen, Phi, Gemma, Llama, Mistral/Mixtral/Nemo, SmolLM, and applicable DeepSeek-distilled models. Keep the matrix architecture-based so new checkpoints can be added without copying model-specific attack code.

Acceptance criteria:

The hard-coded embedding helper branches are replaced by the adapter contract or reduced to documented exceptions.
Every supported family has tokenizer/template/slice tests using a small representative tokenizer and a tiny or randomly initialized compatible model where practical.
At least one opt-in GPU smoke test exercises the complete gradient and candidate-scoring path for representative modern architectures.
Unsupported models fail during preflight with a clear explanation of the missing capability.
Adding support for another standard Hugging Face causal architecture normally requires a compatibility test or adapter registration, not edits to the core GCG optimization loop.
Existing single-model, arbitrary multi-model, single-prompt, multi-prompt, held-out evaluation, and progressive model/goal behavior are preserved.
Support must not be restricted to a single model, exactly two models, or a single shared prompt.

Describe alternatives you've considered, if relevant

We could continue adding architecture-specific isinstance branches, but that does not scale with the Hugging Face model ecosystem and leaves tokenizer assumptions implicit. Replacing PyRIT's implementation with a single-model GCG library would simplify compatibility but would lose PyRIT's multi-model, multi-prompt, and progressive optimization properties.

Additional context

Track this as a focused child of #960. Coordinate performance-related cache work with #962, but keep prefix caching, quantization, model sharding, and heterogeneous-tokenizer gradient aggregation as separate changes so this issue remains reviewable.

I would like to take this.

I reproduced the embeddings half on main (9616789) first, so here is what I am working from.

The three accessors in pyrit/executor/promptgen/gcg/attack/base/attack_manager.py (get_embedding_layer, get_embedding_matrix, get_embeddings) each dispatch through an isinstance chain over six hard-coded classes and end in raise ValueError(f"Unknown model type: {type(model)}"). Since GCG loads through AutoModelForCausalLM, anything outside that list loads fine and then fails at the gradient step.

Building tiny random configs for a spread of architectures, no weights or GPU needed:

model_type       has get_input_embeddings   GCG accessors
gpt2             True                       layer=ok          matrix=ok          embed=ok
llama            True                       layer=ok          matrix=ok          embed=ok
qwen3            True                       layer=ValueError  matrix=ValueError  embed=ValueError
gemma3_text      True                       layer=ValueError  matrix=ValueError  embed=ValueError
olmo2            True                       layer=ValueError  matrix=ValueError  embed=ValueError
starcoder2       True                       layer=ValueError  matrix=ValueError  embed=ValueError
mpt              True                       layer=ValueError  matrix=ValueError  embed=ValueError

Every one of them implements PreTrainedModel.get_input_embeddings(), so the information GCG needs is already reachable through the public API.

Proposed approach

Fall back to model.get_input_embeddings() in all three accessors instead of raising, so any AutoModelForCausalLM works by default. Keep the existing branches where there is a genuine quirk rather than a naming difference: GPT-2/GPT-J and GPT-NeoX currently call .half() in get_embeddings while the Llama/Mistral/Phi-3 paths do not, and I do not want to silently change that dtype behaviour for existing users.
Add unit tests over several architectures using tiny configs like the table above, so the generic path is covered on CPU in CI rather than only for the six listed classes.
Preflight check and compatibility matrix as described, once the accessor part is agreed.

Two questions before I open anything:

Would you prefer this as one PR, or the accessor fix first and the preflight plus matrix separately? The accessor change stands alone and is easy to review in isolation.
On the .half() asymmetry: is that deliberate for the older models, or incidental? If it is incidental I would rather leave it exactly as-is in this PR and raise it separately than fold a behaviour change into a compatibility fix.

I think either way is fine. Heads up it may be a couple weeks before we get to reviewing this because of an upcoming release.

Thanks for the review and the merge.

#2501 only covered the embedding accessors, but I wrote Closes #2489 in the description, so this issue auto-closed when it merged. That was my mistake. Could you reopen it?

Still outstanding from the original scope:

the tokenizer adapter (chat template and character-to-token mapping)
the preflight compatibility check
the open-weight SLM compatibility matrix

I would like to take those next if the issue stays open.

devangpratap

devangpratap you're quite right! Please feel free to start with those. However, I'd recommend keeping changesets as small as possible even if that means creating multiple PRs.

devangpratap

Thanks again for the work here. I reviewed the current main state against the original acceptance criteria.

Completed

#2501 replaced the architecture-specific embedding lookup with Hugging Face's get_input_embeddings() contract and added tiny-model accessor coverage across legacy and newer architectures.
#2679 fixed goal/control/target slice placement so repeated target or control text is located in the correct conversation turn. The current code also explicitly requires a fast tokenizer and chat template and rejects unsafe content/token boundaries.

Remaining

Add explicit model/tokenizer adapter contracts, including device/dtype/vocabulary/cache capability exposure, stable tokenizer/vocabulary identity, encode/decode round-trip checks, and model/tokenizer vocabulary compatibility validation.
Run capability preflight before worker startup and long optimization, with actionable errors for missing or non-differentiable input embeddings and unsupported tokenizer behavior.
Document the architecture-based compatibility matrix for Qwen, Phi, Gemma, Llama, Mistral/Mixtral/Nemo, SmolLM, and applicable DeepSeek-distilled models, with representative tokenizer/template/slice tests.
Add representative modern-architecture coverage for the complete gradient and candidate-scoring path. Current cross-architecture tests cover embedding helpers; the real CPU integration path is GPT-2, and the opt-in GPU AML path is Llama-2.

devangpratap, do you still intend to continue with the remaining pieces? If so, do you need any help splitting or reviewing the work? Otherwise, would you prefer that we make the issue available for someone else to pick up? If we do not hear back by October 22, 2026, we may make it available to other contributors.

Yes, I'd like to keep going. I'll start with the preflight and keep each piece in its own PR.

One question for the modern-architecture coverage: are tiny randomly initialized configs built offline okay, or would you rather the tests load real tiny checkpoints from the Hub?

Thanks for continuing this work. With #3338 merged, the capability-preflight portion is complete, but #2489 still has remaining scope.

Completed

#2501: replaced architecture-specific embedding lookup with Hugging Face's get_input_embeddings() contract and added tiny-model coverage across legacy and newer architectures.
#2679: corrected goal/control/assistant/target slice placement and strengthened template/content-boundary validation.
#3338: added preflight before worker startup for fast-tokenizer support, floating-point input embeddings, and actual prompt/special-token IDs fitting the embedding matrix; it also changed the public default to use_fast=True.

Remaining

Add explicit model and tokenizer adapter contracts that expose embedding matrix, device, dtype, vocabulary size/identity, and relevant cache capabilities.
Add stable tokenizer/vocabulary identity and explicit encode/decode round-trip validation for safe multi-model use.
Document an architecture-based compatibility matrix for Qwen, Phi, Gemma, Llama, Mistral/Mixtral/Nemo, SmolLM, and applicable DeepSeek-distilled models, with representative tokenizer/template/slice tests for each supported family.
Add an opt-in GPU smoke test covering the complete gradient and candidate-scoring path on at least one representative modern architecture. Current real CPU gradient integration is GPT-2, Qwen has candidate-evaluation unit coverage, and the AML end-to-end test remains Llama-2.

Suggested next steps

A focused adapter-contract PR covering model/tokenizer metadata, stable identity, and round-trip validation.
A separate compatibility-matrix and representative family-test PR.
A final opt-in modern-architecture GPU smoke-test PR.

Tiny randomly initialized configs built offline are appropriate for the family compatibility tests where practical; the opt-in GPU smoke test should exercise a real representative checkpoint and the full path.

Navigation Menu

Repository navigation

FEAT Add generic GCG model adapters and an open-weight SLM compatibility matrix #2489

Is your feature request related to a problem? Please describe.

Describe the solution you'd like

Describe alternatives you've considered, if relevant

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

FEAT Add generic GCG model adapters and an open-weight SLM compatibility matrix #2489

Description

Is your feature request related to a problem? Please describe.

Describe the solution you'd like

Describe alternatives you've considered, if relevant

Additional context

Activity

devangpratap commented on Aug 26, 2026

romanlutz commented on Aug 27, 2026

devangpratap commented on Sep 10, 2026

romanlutz commented on Sep 16, 2026

romanlutz commented on Oct 18, 2026

devangpratap commented on Oct 19, 2026

romanlutz commented on Oct 22, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions