FIX: Prevent GCG model gradient accumulation

[romanlutz]

Description

GCG retained a full gradient tensor for every trainable model parameter while computing coordinate gradients, roughly doubling model-related CUDA memory. The multiprocessing path also serialized a model with every worker operation even though each worker already owned a persistent model.

This change addresses those root causes by:

• computing only the input coordinate gradient with torch.autograd.grad(loss, one_hot)
• freezing and clearing the persistent worker-owned model once at worker startup
• using typed, spawn-safe worker task payloads that exclude the model for grad, logits, test, and test-loss operations
• streaming multi-prompt gradient aggregation, with FP32 accumulation for FP16/BF16 inputs to preserve cancellation-sensitive contributions
• removing gc.collect() and per-step torch.cuda.empty_cache() calls from GCG hot paths

A fresh-process TinyLlama 1.1B FP16 benchmark on an RTX 2000 Ada 8 GB ran 8 steps across 3 repeats in both local and multiprocessing modes (96 measured steps). All paired loss traces were bitwise identical and live allocations remained flat across every local run.

Metric	origin/main	This change
Local peak live CUDA	4456.72 MiB	2289.72 MiB (-48.62%)
Local nvidia-smi peak	5077 MiB	2731 MiB (-46.21%)
Local steady-step latency	3.316 s	0.726 s (-78.11%)
Multiprocessing nvidia-smi peak	4832 MiB	2830 MiB (-41.43%)
Multiprocessing steady-step latency	5.571 s	2.715 s (-51.26%)
Model parameter-gradient bytes	2098.18 MiB	0

The live and peak allocation reductions independently confirm that this is not merely allocator-cache trimming. Benchmark scripts and raw results were kept outside the repository.

Closes #961

Tests and Documentation

• uv run pytest tests\unit\executor\promptgen\gcg\test_gcg_core.py -q (56 passed)
• uv run ruff format --check on the changed files
• uv run ruff check on the changed files
• uv run ty check on the changed files
• Repository pre-commit hooks

No documentation changes are needed because this preserves the existing GCG API and attack behavior.

[Copilot]

Pull request overview

This PR reduces GCG CUDA memory usage and improves throughput by preventing unnecessary model parameter-gradient materialization and by avoiding repeated model serialization in the multiprocessing worker pipeline, while keeping GCG behavior/API consistent.

Changes:

• Compute coordinate gradients via torch.autograd.grad(loss, one_hot) to avoid accumulating per-parameter gradients.
• Introduce typed worker operations/tasks (ModelWorkerOperation, ModelWorkerTask) so queued multiprocessing work no longer carries a model payload per operation.
• Stream prompt-gradient aggregation with FP32 accumulation for FP16/BF16 inputs, and remove gc.collect() / per-step torch.cuda.empty_cache() from hot paths.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py	Switches token-gradient computation to autograd.grad and updates worker dispatch to use typed operations without model payloads.
pyrit/executor/promptgen/gcg/attack/base/attack_manager.py	Adds ModelWorkerOperation/ModelWorkerTask, updates worker execution to use the persistent worker-owned model, and streams gradient reduction with FP32 accumulation where needed.
tests/unit/executor/promptgen/gcg/test_gcg_core.py	Expands unit coverage to validate spawn-safe task payloads, worker-owned model usage, streamed gradient reduction, and parity vs backward-based gradients.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.