From e8930962ff609f2a92b806af02e51eb2793e8061 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:32:14 -0700 Subject: [PATCH 1/2] FIX: Prevent GCG model gradient accumulation Compute coordinate gradients with input-only autograd, keep model ownership inside persistent workers, and stream prompt aggregation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5c847c5a-6917-4e2b-81a6-e511d9cbbe6e --- .../gcg/attack/base/attack_manager.py | 89 +++++--- .../promptgen/gcg/attack/gcg/gcg_attack.py | 21 +- .../executor/promptgen/gcg/test_gcg_core.py | 203 +++++++++++++++++- 3 files changed, 270 insertions(+), 43 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 04150ae408..e545ef3f00 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -3,13 +3,14 @@ from __future__ import annotations -import gc import json import logging import math import random import time from copy import deepcopy +from dataclasses import dataclass +from enum import Enum from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -383,12 +384,10 @@ def logits(self, model: Any, test_controls: Any = None, return_ids: bool = False if return_ids: del locs, test_ids - gc.collect() return model(input_ids=ids, attention_mask=attn_mask).logits, ids del locs, test_ids logits = model(input_ids=ids, attention_mask=attn_mask).logits del ids - gc.collect() return logits def target_loss(self, logits: torch.Tensor, ids: torch.Tensor) -> torch.Tensor: @@ -598,7 +597,14 @@ def grad(self, model: Any) -> torch.Tensor: Returns: torch.Tensor: Aggregated prompt gradients. """ - return torch.stack([prompt.grad(model) for prompt in self._prompts]).sum(dim=0) + first_gradient = self._prompts[0].grad(model) + if len(self._prompts) == 1: + return first_gradient + result_dtype = first_gradient.dtype + gradient = first_gradient.float() if result_dtype in (torch.float16, torch.bfloat16) else first_gradient.clone() + for prompt in self._prompts[1:]: + gradient.add_(prompt.grad(model).to(dtype=gradient.dtype)) + return gradient.to(dtype=result_dtype) def logits(self, model: Any, test_controls: Any = None, return_ids: bool = False) -> Any: """ @@ -888,7 +894,6 @@ def control_weight_fn(_: int) -> float: steps += 1 start = time.time() - torch.cuda.empty_cache() control, loss = self.step( batch_size=batch_size, topk=topk, @@ -940,14 +945,14 @@ def test( Jailbreak, exact-match, and loss results. """ for j, worker in enumerate(workers): - worker(prompts[j], "test", worker.model) + worker(prompts[j], ModelWorkerOperation.TEST) model_tests = np.array([worker.results.get() for worker in workers]) model_tests_jb = model_tests[..., 0].tolist() model_tests_mb = model_tests[..., 1].tolist() model_tests_loss: list[list[float]] = [] if include_loss: for j, worker in enumerate(workers): - worker(prompts[j], "test_loss", worker.model) + worker(prompts[j], ModelWorkerOperation.TEST_LOSS) model_tests_loss = [worker.results.get() for worker in workers] return model_tests_jb, model_tests_mb, model_tests_loss @@ -1781,6 +1786,26 @@ def run( return total_jb, total_em, test_total_jb, test_total_em, total_outputs, test_total_outputs +class ModelWorkerOperation(str, Enum): + """A model operation supported by ``ModelWorker``.""" + + GRAD = "grad" + LOGITS = "logits" + CONTRAST_LOGITS = "contrast_logits" + TEST = "test" + TEST_LOSS = "test_loss" + + +@dataclass(frozen=True) +class ModelWorkerTask: + """A spawn-safe model worker task that excludes the worker-owned model.""" + + obj: Any + operation: ModelWorkerOperation | Callable[..., Any] + args: tuple[Any, ...] + kwargs: dict[str, Any] + + class ModelWorker: """Run model operations in a dedicated multiprocessing worker.""" @@ -1802,33 +1827,30 @@ def __init__( move_to_device = cast("Callable[[torch.device], PreTrainedModel]", model.to) self.model = move_to_device(torch.device(device)).eval() self.tokenizer = tokenizer - self.tasks: mp.JoinableQueue[Any] = mp.JoinableQueue() + self.tasks: mp.JoinableQueue[ModelWorkerTask | None] = mp.JoinableQueue() self.results: mp.JoinableQueue[Any] = mp.JoinableQueue() self.process: mp.Process | None = None @staticmethod - def run(model: Any, tasks: mp.JoinableQueue[Any], results: mp.JoinableQueue[Any]) -> None: + def run( + model: Any, + tasks: mp.JoinableQueue[ModelWorkerTask | None], + results: mp.JoinableQueue[Any], + ) -> None: """Process queued model operations until a stop sentinel arrives.""" + model.requires_grad_(False) + model.zero_grad(set_to_none=True) while True: task = tasks.get() if task is None: + tasks.task_done() break - ob, fn, args, kwargs = task - if fn == "grad": + if task.operation is ModelWorkerOperation.GRAD: with torch.enable_grad(): # type: ignore[no-untyped-call, unused-ignore] - results.put(ob.grad(*args, **kwargs)) + results.put(ModelWorker._execute_task(model=model, task=task)) else: with torch.no_grad(): - if fn == "logits": - results.put(ob.logits(*args, **kwargs)) - elif fn == "contrast_logits": - results.put(ob.contrast_logits(*args, **kwargs)) - elif fn == "test": - results.put(ob.test(*args, **kwargs)) - elif fn == "test_loss": - results.put(ob.test_loss(*args, **kwargs)) - else: - results.put(fn(*args, **kwargs)) + results.put(ModelWorker._execute_task(model=model, task=task)) tasks.task_done() def start(self) -> ModelWorker: @@ -1856,16 +1878,35 @@ def stop(self) -> ModelWorker: torch.cuda.empty_cache() return self - def __call__(self, ob: Any, fn: str, *args: Any, **kwargs: Any) -> ModelWorker: + def __call__( + self, + ob: Any, + operation: ModelWorkerOperation | Callable[..., Any], + *args: Any, + **kwargs: Any, + ) -> ModelWorker: """ Queue an operation for execution by this worker. Returns: ModelWorker: This worker. """ - self.tasks.put((deepcopy(ob), fn, args, kwargs)) + self.tasks.put(ModelWorkerTask(obj=deepcopy(ob), operation=operation, args=args, kwargs=kwargs)) return self + @staticmethod + def _execute_task(*, model: Any, task: ModelWorkerTask) -> Any: + """ + Execute a task with the persistent model when the operation requires one. + + Returns: + Any: The operation result. + """ + if isinstance(task.operation, ModelWorkerOperation): + method = getattr(task.obj, task.operation.value) + return method(model, *task.args, **task.kwargs) + return task.operation(*task.args, **task.kwargs) + def get_workers(params: Any, evaluation: bool = False) -> tuple[list[ModelWorker], list[ModelWorker]]: """ diff --git a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py index cdacecd7e6..916260e267 100644 --- a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py +++ b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import gc import logging from typing import Any @@ -12,6 +11,7 @@ from pyrit.executor.promptgen.gcg.attack.base.attack_manager import ( AttackPrompt, + ModelWorkerOperation, MultiPromptAttack, PromptManager, get_embedding_matrix, @@ -48,7 +48,7 @@ def token_gradients( torch.Tensor: The gradients of each token in the input_slice with respect to the loss. Raises: - RuntimeError: If backpropagation does not produce token gradients. + RuntimeError: If autograd does not produce token gradients. """ embed_weights = get_embedding_matrix(model) one_hot = torch.zeros( @@ -70,11 +70,10 @@ def token_gradients( targets = input_ids[target_slice] loss = nn.CrossEntropyLoss()(logits[0, loss_slice, :], targets) - loss.backward() - - if one_hot.grad is None: - raise RuntimeError("Model backward pass did not produce token gradients") - return one_hot.grad.clone() + coordinate_gradient = torch.autograd.grad(loss, one_hot, allow_unused=True)[0] + if coordinate_gradient is None: + raise RuntimeError("Autograd did not produce token gradients") + return coordinate_gradient class GCGAttackPrompt(AttackPrompt): @@ -273,7 +272,7 @@ def step( loss_function = self._resolve_loss(target_weight=target_weight, control_weight=control_weight) for j, worker in enumerate(self.workers): - worker(self.prompts[j], "grad", worker.model) + worker(self.prompts[j], ModelWorkerOperation.GRAD) # Aggregate gradients grad = None @@ -324,7 +323,6 @@ def step( ) ) del grad, control_cand - gc.collect() # Search loss = torch.zeros(len(control_cands) * batch_size).to(main_device) @@ -336,7 +334,7 @@ def step( prompt_indices = progress if progress is not None else range(len(self.prompts[0])) for i in prompt_indices: for k, worker in enumerate(self.workers): - worker(self.prompts[k][i], "logits", worker.model, cand, return_ids=True) + worker(self.prompts[k][i], ModelWorkerOperation.LOGITS, cand, return_ids=True) logits, ids = zip(*[worker.results.get() for worker in self.workers], strict=True) loss[j * batch_size : (j + 1) * batch_size] += sum( loss_function.compute_loss( @@ -348,7 +346,6 @@ def step( for k, (logit, token_ids) in enumerate(zip(logits, ids, strict=True)) ) del logits, ids - gc.collect() if progress is not None: progress.set_description( @@ -359,9 +356,7 @@ def step( model_idx = min_idx // batch_size batch_idx = min_idx % batch_size next_control, cand_loss = control_cands[model_idx][batch_idx], loss[min_idx] - del control_cands, loss - gc.collect() current_length = self._get_control_length(control=next_control) if current_length is not None: diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index 4d8332e185..148e67c980 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -1,8 +1,11 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import pickle +from copy import deepcopy +from dataclasses import dataclass from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch, sentinel import pytest @@ -18,6 +21,8 @@ EvaluateAttack = attack_manager_mod.EvaluateAttack IndividualPromptAttack = attack_manager_mod.IndividualPromptAttack ModelWorker = attack_manager_mod.ModelWorker +ModelWorkerOperation = attack_manager_mod.ModelWorkerOperation +ModelWorkerTask = attack_manager_mod.ModelWorkerTask ProgressiveMultiPromptAttack = attack_manager_mod.ProgressiveMultiPromptAttack get_embedding_layer = attack_manager_mod.get_embedding_layer get_embedding_matrix = attack_manager_mod.get_embedding_matrix @@ -38,6 +43,54 @@ LengthPreservingFilter = default_implementations_mod.LengthPreservingFilter +@dataclass +class _TinyModelOutput: + logits: torch.Tensor + + +class _TinyCausalLM(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = torch.nn.Embedding(8, 4) + self.projection = torch.nn.Linear(4, 8, bias=False) + + @property + def device(self) -> torch.device: + return self.embedding.weight.device + + def forward(self, *, inputs_embeds: torch.Tensor) -> _TinyModelOutput: + return _TinyModelOutput(logits=self.projection(inputs_embeds.cumsum(dim=1))) + + +def _backward_coordinate_gradient( + *, + model: _TinyCausalLM, + input_ids: torch.Tensor, + input_slice: slice, + target_slice: slice, + loss_slice: slice, +) -> torch.Tensor: + embedding_weights = model.embedding.weight + one_hot = torch.zeros( + input_ids[input_slice].shape[0], + embedding_weights.shape[0], + device=model.device, + dtype=embedding_weights.dtype, + ) + one_hot.scatter_(1, input_ids[input_slice].unsqueeze(1), torch.ones(one_hot.shape[0], 1)) + one_hot.requires_grad_() + input_embeddings = (one_hot @ embedding_weights).unsqueeze(0) + embeddings = model.embedding(input_ids.unsqueeze(0)).detach() + full_embeddings = torch.cat( + [embeddings[:, : input_slice.start, :], input_embeddings, embeddings[:, input_slice.stop :, :]], dim=1 + ) + logits = model(inputs_embeds=full_embeddings).logits + loss = torch.nn.CrossEntropyLoss()(logits[0, loss_slice, :], input_ids[target_slice]) + loss.backward() + assert one_hot.grad is not None + return one_hot.grad.clone() + + class TestGetFilteredCands: """Tests for MultiPromptAttack.get_filtered_cands.""" @@ -582,6 +635,75 @@ def test_model_worker_uses_model_device_dispatch() -> None: assert worker.model is evaluated_model +def test_model_worker_task_payload_excludes_model() -> None: + worker = object.__new__(ModelWorker) + worker.model = sentinel.model + worker.tasks = MagicMock() + prompt = {"prompt": "value"} + + worker(prompt, ModelWorkerOperation.GRAD, 42, option=True) + + task = worker.tasks.put.call_args.args[0] + assert isinstance(task, ModelWorkerTask) + assert task.obj == prompt + assert task.obj is not prompt + assert task.operation is ModelWorkerOperation.GRAD + assert task.args == (42,) + assert task.kwargs == {"option": True} + assert not hasattr(task, "model") + assert all(argument is not worker.model for argument in task.args) + assert all(value is not worker.model for value in task.kwargs.values()) + assert pickle.loads(pickle.dumps(task)) == task + + +@pytest.mark.parametrize( + ("operation", "method_name"), + [ + (ModelWorkerOperation.GRAD, "grad"), + (ModelWorkerOperation.LOGITS, "logits"), + (ModelWorkerOperation.CONTRAST_LOGITS, "contrast_logits"), + (ModelWorkerOperation.TEST, "test"), + (ModelWorkerOperation.TEST_LOSS, "test_loss"), + ], +) +def test_model_worker_run_uses_worker_owned_model(operation: ModelWorkerOperation, method_name: str) -> None: + model = MagicMock() + target = MagicMock() + getattr(target, method_name).return_value = sentinel.result + task = ModelWorkerTask( + obj=target, + operation=operation, + args=(sentinel.argument,), + kwargs={"option": True}, + ) + tasks = MagicMock() + tasks.get.side_effect = [task, None] + results = MagicMock() + + ModelWorker.run(model, tasks, results) + + model.requires_grad_.assert_called_once_with(False) + model.zero_grad.assert_called_once_with(set_to_none=True) + getattr(target, method_name).assert_called_once_with(model, sentinel.argument, option=True) + results.put.assert_called_once_with(sentinel.result) + assert tasks.task_done.call_count == 2 + + +def test_multi_prompt_test_dispatches_without_model_payload() -> None: + attack = object.__new__(MultiPromptAttack) + worker = MagicMock() + worker.results.get.side_effect = [[(True, 1)], [0.25]] + prompt = MagicMock() + + result = attack.test([worker], [prompt], include_loss=True) + + assert result == ([[True]], [[1]], [[0.25]]) + assert worker.call_args_list == [ + call(prompt, ModelWorkerOperation.TEST), + call(prompt, ModelWorkerOperation.TEST_LOSS), + ] + + class _Queue: def __init__(self, items: list[Any]) -> None: self._items = list(items) @@ -836,6 +958,15 @@ def test_step_default_path_matches_legacy_behavior(self) -> None: assert actual_control == legacy_controls[0] assert actual_loss == pytest.approx(legacy_loss[0].item()) + grad_args, grad_kwargs = worker.calls[0] + assert grad_args == (prompt_manager, ModelWorkerOperation.GRAD) + assert grad_kwargs == {} + logits_args, logits_kwargs = worker.calls[1] + assert logits_args[0] is prompt + assert logits_args[1] is ModelWorkerOperation.LOGITS + assert len(logits_args) == 3 + assert all(argument is not worker.model for argument in logits_args) + assert logits_kwargs == {"return_ids": True} def test_step_uses_custom_protocol_implementations_when_supplied(self) -> None: gradient = torch.randn(3, 6) @@ -1073,17 +1204,42 @@ def test_attack_prompt_logits_builds_attention_mask() -> None: assert torch.equal(model.call_args.kwargs["attention_mask"], torch.ones(1, 4, dtype=torch.long)) -def test_prompt_manager_grad_sums_prompt_gradients() -> None: +def test_prompt_manager_grad_streams_and_sums_prompt_gradients() -> None: prompt_manager = object.__new__(PromptManager) first_prompt = MagicMock() first_prompt.grad.return_value = torch.tensor([1.0, 2.0]) second_prompt = MagicMock() second_prompt.grad.return_value = torch.tensor([3.0, 4.0]) - prompt_manager._prompts = [first_prompt, second_prompt] + third_prompt = MagicMock() + third_prompt.grad.return_value = torch.tensor([5.0, 6.0]) + prompt_manager._prompts = [first_prompt, second_prompt, third_prompt] + model = MagicMock() + + with patch.object(torch, "stack", side_effect=AssertionError("prompt gradients must be streamed")): + result = prompt_manager.grad(model) + + assert torch.equal(result, torch.tensor([9.0, 12.0])) + for prompt in prompt_manager._prompts: + prompt.grad.assert_called_once_with(model) + + +def test_prompt_manager_grad_preserves_fp16_reduction_precision() -> None: + prompt_manager = object.__new__(PromptManager) + prompt_gradients = [ + torch.tensor([10000.0], dtype=torch.float16), + torch.tensor([1.0], dtype=torch.float16), + torch.tensor([-10000.0], dtype=torch.float16), + ] + prompt_manager._prompts = [MagicMock() for _ in prompt_gradients] + for prompt, gradient in zip(prompt_manager._prompts, prompt_gradients, strict=True): + prompt.grad.return_value = gradient result = prompt_manager.grad(MagicMock()) - assert torch.equal(result, torch.tensor([4.0, 6.0])) + expected = torch.stack(prompt_gradients).sum(dim=0) + assert torch.equal(result, expected) + assert result.item() == 1.0 + assert result.dtype is torch.float16 def test_multi_prompt_run_anneals_and_accepts_lower_loss() -> None: @@ -1142,7 +1298,41 @@ def test_gcg_step_requires_worker() -> None: attack.step() -def test_token_gradients_raises_when_backward_produces_no_gradient() -> None: +def test_token_gradients_matches_backward_without_model_parameter_gradients() -> None: + torch.manual_seed(2026) + backward_model = _TinyCausalLM() + input_only_model = deepcopy(backward_model) + input_ids = torch.tensor([0, 1, 2, 3, 4]) + input_slice = slice(1, 3) + target_slice = slice(3, 5) + loss_slice = slice(2, 4) + expected = _backward_coordinate_gradient( + model=backward_model, + input_ids=input_ids, + input_slice=input_slice, + target_slice=target_slice, + loss_slice=loss_slice, + ) + + with ( + patch.object(gcg_attack_mod, "get_embedding_matrix", side_effect=lambda model: model.embedding.weight), + patch.object(gcg_attack_mod, "get_embeddings", side_effect=lambda model, ids: model.embedding(ids)), + ): + actual = token_gradients( + input_only_model, + input_ids, + input_slice=input_slice, + target_slice=target_slice, + loss_slice=loss_slice, + ) + + assert torch.equal(actual, expected) + assert not actual.requires_grad + assert actual.grad_fn is None + assert all(parameter.grad is None for parameter in input_only_model.parameters()) + + +def test_token_gradients_raises_when_coordinate_gradient_missing() -> None: model = MagicMock() model.device = torch.device("cpu") model.return_value.logits = torch.randn(1, 3, 4) @@ -1153,7 +1343,8 @@ def test_token_gradients_raises_when_backward_produces_no_gradient() -> None: patch.object(gcg_attack_mod, "get_embedding_matrix", return_value=torch.ones(4, 2)), patch.object(gcg_attack_mod, "get_embeddings", return_value=torch.ones(1, 3, 2)), patch.object(gcg_attack_mod.nn, "CrossEntropyLoss", return_value=loss_function), - pytest.raises(RuntimeError, match="did not produce token gradients"), + patch.object(gcg_attack_mod.torch.autograd, "grad", return_value=(None,)), + pytest.raises(RuntimeError, match="Autograd did not produce token gradients"), ): token_gradients( model, From fec3feb2205f9b4614765099bfad6c646fd2fc1c Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:57:49 -0700 Subject: [PATCH 2/2] DOC: Explain worker task memory savings Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5c847c5a-6917-4e2b-81a6-e511d9cbbe6e --- .../executor/promptgen/gcg/attack/base/attack_manager.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index e545ef3f00..efd0fc7f30 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -1798,7 +1798,14 @@ class ModelWorkerOperation(str, Enum): @dataclass(frozen=True) class ModelWorkerTask: - """A spawn-safe model worker task that excludes the worker-owned model.""" + """ + A spawn-safe worker payload that excludes the worker-owned model. + + Typed model operations receive the worker's persistent model during dispatch, + so each queued task serializes only the prompt payload and operation arguments + rather than serializing the full model again. ``obj`` is the prompt or prompt + manager that receives the operation. + """ obj: Any operation: ModelWorkerOperation | Callable[..., Any]