FIX locate GCG target and control after the user turn - #2679
Merged
Roman Lutz (romanlutz) merged 10 commits intoSep 26, 2026
Merged
Conversation
_update_ids searched the rendered prompt for the goal, control and target independently, so each one matched its first occurrence anywhere in the prompt. A goal that quotes its own target, which happens when the goal spells out the affirmative prefix it wants, made the target match inside the user turn instead of the assistant reply, and the target and loss slices then pointed at the wrong tokens, with the model input truncated before the assistant turn and no error raised. Search for the user turn as one string, derive the control offset from it, and look for the target only after that turn.
Roman Lutz (romanlutz)
requested changes
Sep 22, 2026
…ser string
Find where the assistant content starts by rendering the user turn alone
with add_generation_prompt=True, and search for the target from there, so a
target that also names the role marker (e.g. "assistant" in <|assistant|>)
no longer lands on it. Locate the control as its last occurrence before that
boundary and the goal just before it, instead of requiring the raw
f"{goal} {control}" string, which a trimming template breaks for the empty
goals target-only datasets use.
Roman Lutz (romanlutz)
requested changes
Sep 22, 2026
Rendering the user turn with add_generation_prompt=True does not prove
where the assistant content starts: a template may ignore the flag, and
the assistant marker stayed inside the control search range, so a control
or target containing role text matched the scaffolding instead of the turn
("assistant" is an ordinary word, and an optimized control is decoded
vocabulary tokens). Render the same two messages with sentinel contents
instead, take the prefix and the separator the template puts between the
two contents, and search for the control and goal inside the user content
and the target after the separator.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Measure each rendered message independently and verify the full surrounding prompt instead of scanning for separators or falling back to unbounded searches. Derive goal and control positions from the verified user content. Map complete token offsets, handle whitespace-consuming role tokens, and reject unsafe or empty optimization spans. Cover repeated text, transformed templates, real tokenizer offsets, and shared GCG optimization paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
The separator between the user and assistant contents was found with a first-occurrence search, so a goal quoting the template's own markers moved the boundary. Render the real user content with a sentinel reply and read the boundary off its position instead. When the template does not render contents verbatim, raise rather than falling back to an unbounded search, which silently corrupted the target and loss slices.
Retain the author regression cases while using verified content boundaries and token offsets. Keep whitespace-consuming role tokens out of optimization spans and preserve exact target/loss slices. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Roman Lutz (romanlutz)
approved these changes
Sep 26, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
_update_ids located the goal, control and target with three independent
prompt.find calls, so each matched its first occurrence anywhere in the rendered
prompt.
GCG targets are affirmative prefixes, so a goal that spells out the prefix it
wants ("Respond with Sure, here is the plan") puts that same text in the user
turn. The target then matched there instead of in the assistant reply, and
_target_slice, _loss_slice and _assistant_role_slice silently pointed at the
wrong tokens. Nothing raised, so the search just optimized against the wrong
positions. The same applies to a goal containing the control string.
The goal and control are rendered as one contiguous user turn, so this searches
for that whole string, derives the control offset from it, and searches for the
target only after that turn.
Testing
New test in TestUpdateIdsErrorPaths, using a mock tokenizer whose char_to_token
maps characters to whitespace-delimited tokens the way a fast tokenizer does. It
passes with this change and fails without it. Full GCG unit suite passes (244),
and pre-commit including ruff and ty is clean.
Contributes to #2489.