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.

[ValbuenaVC]

Can we better document how this is relevant to the PR's goal of reducing overall memory usage? You added this in the PR body ("using typed, spawn-safe worker task payloads that exclude the model for grad, logits, test, and test-loss operations") but the description as-is is vague. Something like:

Suggested change

	"""A spawn-safe model worker task that excludes the worker-owned model."""
	"""
	A spawn-safe model worker task that excludes the worker-owned model.
	By assigning each worker a task with its own ModelWorkerOperation and
	tensor object / prompt (obj), the lifecycle and memory footprint of each operation
	is bounded, preventing memory overuse.
	"""

That exact description is probably missing important context especially for obj which is just typed as Any but I think we should make sure it's explicit why we introduce this abstraction to handle memory usage