MAINT Complete GCG optimization phase extraction after #2416

[romanlutz]

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

#2416 and merged #2467 made GCG optimization state explicit through OptimizationRunState and ProgressiveScheduleState. They also extracted stopping, checkpoint logging, final progressive evaluation, and final candidate selection.

The remaining optimization phases are still concentrated in two large methods:

• GCGMultiPromptAttack.step() combines worker gradient dispatch and collection, gradient normalization/grouping, candidate sampling, candidate filtering, VRAM-bounded logits evaluation, loss accumulation, progress output, candidate selection, and tensor cleanup.
• ProgressiveMultiPromptAttack.run() combines construction of each inner attack with goal and worker admission, remaining-budget calculation, loss-measurement tracking, control-weight scheduling, stop-policy changes, and finalization.

This makes each phase difficult to test independently and raises the risk of changing candidate ordering, worker/result pairing, memory behavior, or exact-budget scheduling while working on adjacent GCG improvements.

This issue is a structural follow-up to #2416, not a performance or RNG implementation. Faster algorithms remain tracked by #962. Deterministic RNG propagation remains tracked by #2490 and open PR #2502.

Describe the solution you'd like

Complete the behavior-preserving phase extraction as three independently reviewable work items under this issue.

1. Extract candidate proposal

Move gradient dispatch/collection, per-worker normalization, grouping of incompatible gradient shapes, candidate sampling, and candidate filtering out of GCGMultiPromptAttack.step() behind a typed internal result such as a candidate-group or candidate-batch model.

The extracted contract must preserve:

• worker and model-group ordering;
• the current normalization and same-shape aggregation behavior;
• one candidate group per compatible gradient group;
• the existing SamplingStrategy and CandidateFilter extension protocols;
• batch_size, topk, temperature, and non-ASCII behavior;
• current error behavior when workers or aggregate gradients are unavailable.

2. Extract VRAM-bounded candidate evaluation

Move the candidate-group × prompt × worker logits/loss loop into a focused evaluator that returns losses paired with the exact candidate groups it evaluated. Keep candidate batches step-local and preserve the current sequential group evaluation that bounds VRAM.

The extracted contract must preserve:

• ModelWorkerOperation.LOGITS dispatch and result-queue pairing;
• target/control slices passed to LossFunction;
• loss accumulation across prompts and workers;
• flat ordering expected by _select_best_candidate();
• progress reporting without making it part of the evaluator's data model;
• timely release of logits, token IDs, gradients, candidate tensors, and aggregate loss tensors;
• output suffix and normalized loss equivalence.

This boundary should make future work under #962 possible without implementing a faster algorithm here.

3. Model progressive admission transitions

Move the transition decisions in ProgressiveMultiPromptAttack.run() into explicit methods or a small typed schedule controller. The run method should orchestrate inner attacks; the schedule component should decide what state follows a completed inner result.

The extracted contract must preserve:

• progressive goal admission before worker admission;
• remaining-step-budget calculation and exact-budget termination;
• no admission or sentinel reset after the budget is exhausted;
• measured-loss tracking without treating a legitimate non-finite loss as unmeasured;
• stop_inner_on_success transitions;
• control-weight increments and their threshold;
• final held-out evaluation only after all goals and workers are admitted;
• last_schedule_state, public return values, and existing log behavior.

Each work item should be independently reviewable and leave public GCG configuration, extension protocols, results, and logging formats unchanged. The implementation may land in multiple pull requests that reference this umbrella issue.

Describe alternatives you've considered, if relevant

One broad rewrite could replace both loops at once, but it would combine candidate semantics, device/memory behavior, and progressive scheduling in a high-risk diff. Conversely, opening three unrelated issues would lose the shared #2416 follow-up context. This umbrella keeps one architectural goal while defining three contribution-sized phases.

Extracting annealing and accepted-candidate state is intentionally deferred. #2490 and #2502 currently modify that path, so a separate transition extraction should be reconsidered after the RNG work lands rather than creating overlapping work now.

The issue also should not implement probe sampling, alternative loss algorithms, vectorized scoring, or other speedups from #962. These extractions provide testable boundaries for that work.

Additional context

Suggested order:

1. Candidate proposal, which establishes the typed grouping contract.
2. Candidate evaluation, which consumes those groups and preserves the VRAM boundary.
3. Progressive scheduling, which is independent of the first two and may proceed in parallel if desired.

Required regression coverage should include:

• same-shape and mixed-shape worker gradients;
• exact candidate-group and candidate ordering;
• sampling/filter extension implementations;
• multiple prompts and workers with exact logits/result pairing;
• loss accumulation and normalization;
• bounded sequential candidate-group evaluation and cleanup on success/failure;
• goal-only, worker-only, and combined progressive schedules;
• exact-budget exhaustion at each admission boundary;
• legitimate infinite/non-finite measured losses;
• control-weight threshold transitions;
• early success and final held-out evaluation;
• seeded trajectory and final result parity with the current implementation.

Use deterministic fakes and queues rather than timing-based assertions. Run the complete GCG unit suite and relevant integration tests after each phase. Any performance numbers should be reported under #962, and deterministic RNG acceptance criteria remain owned by #2490/#2502.

[Kaap10]

Hi Roman Lutz (@romanlutz),

I’ve gone through the issue context from #2416 / PR #2467 and closely examined the current implementation in GCGMultiPromptAttack.step().

I’d love to take this up and tackle Item 1: Extract candidate proposal as the first independent PR.

Proposed Implementation Plan for Item 1:

1. Typed Proposal Boundary: Extract gradient dispatch, per-worker $L_2$ normalization, incompatible-shape grouping, token sampling, and candidate filtering into a dedicated GCGCandidateProposer returning a typed, immutable CandidateProposalBatch.
2. Preserved Invariants:
   • Retain exact worker ordering and model-group grouping behavior during shape transitions.
   • Strictly preserve the SamplingStrategy and CandidateFilter extension protocols (batch_size, topk, temperature, allow_non_ascii).
   • Preserve device management by guaranteeing cross-worker gradient tensors are migrated to main_device before accumulation.
3. Comprehensive Deterministic Testing:
   • Add dedicated unit tests in tests/unit/executor/promptgen/gcg/test_gcg_proposal.py using deterministic fakes/mocks (covering single-worker, mixed-shape multi-worker, empty worker validation, and custom filter protocol injection).
4. Zero Regressions: Keep backward-compatible delegation shims on GCGMultiPromptAttack to ensure zero breakage for existing tests or downstream code.

Could you please assign Item 1 to me? I will open an independently reviewable PR once the implementation and test matrix are complete.

[Kaap10]

Hi Roman Lutz (@romanlutz),

I have opened PR #2671 to address Item 1: Extract candidate proposal:

• Extracted proposal logic into a modular GCGCandidateProposer with an immutable CandidateProposalBatch dataclass.
• Preserved legacy shims (_sample_control_candidates, _filter_control_candidates) in GCGAttack for 100% backward compatibility.
• Added comprehensive unit tests covering deterministic sampling, filtering, and cross-device gradient transfers.

Please let me know if you have any feedback or if any changes are required. Once this PR is reviewed and merged, I will proceed with Item 2: Extract VRAM-bounded candidate evaluation in the next PR.

Thanks!

[Kaap10]

Hi Roman Lutz (@romanlutz),

Thanks for reviewing and merging #2671, and for adding the last-worker tests - noted that group indexing nuance for the next parts!

Should I go ahead and pick up Item 2 (Candidate Evaluation) next?
Thanks!

[romanlutz]

Sure if you can!

[Kaap10]

Just put up #2700 for Item 2! Extracted the candidate evaluation phase into GCGCandidateEvaluator with full unit test coverage. Whenever you have a moment, would love your feedback!

[Kaap10]

Thanks for merging #2700, Roman Lutz (@romanlutz)!

I'm ready to tackle Item 3 (Model progressive admission transitions) next as the final phase.

[romanlutz]

Appreciate all the rapid iterations 🙂