Refactoring GCG Adversarial Attacks in Microsoft PyRIT: Architecture, VRAM Optimization and Clean Systems Design

TL;DR

Refactored Microsoft PyRIT’s monolithic 1,000-line Greedy Coordinate Gradient (GCG) adversarial optimization loop into 3 typed, modular components (CandidateProposer, CandidateEvaluator, and ProgressiveScheduleController).

Bounded GPU VRAM consumption to $O(1)$ by introducing sequential candidate evaluation with immediate tensor cleanup del logits, ids), completely eliminating multi-model Out-Of-Memory (OOM) crashes.

Maintained 100% backward compatibility with zero breaking changes, supported by 300+ deterministic automated tests running across Windows, macOS, and Linux runners without requiring active GPUs.

• Tracking Issue: Closes #2665
• Merged PRs: #2671, #2700, #2720

[1] Introduction

Large language models (LLMs) are deployed with safety guardrails to prevent harmful, dangerous, or unauthorized outputs. In AI security and red teaming, testing whether these guardrails can be bypassed is an essential part of defensive engineering.

One of the most effective white-box algorithms for this is Greedy Coordinate Gradient (GCG). It searches for specific suffix tokens that, when appended to a prompt, guide the neural network into bypassing its safety constraints.

While GCG is mathematically powerful, running it reliably at scale is a difficult engineering challenge. In production red teaming frameworks, GCG must handle multiple models, diverse tokenizers, and heavy GPU tensor computations.

This article walks through how I refactored the GCG optimization engine inside Microsoft PyRIT (Python Risk Identification Tool for Generative AI) across three merged pull requests under Issue #2665, resolving critical VRAM spikes and transforming a monolithic loop into a maintainable, typed architecture.

---

[2] Background: How GCG Works

To understand the engineering constraints, it helps to understand how GCG operates under the hood:

1. Target Loss Formulation: The attacker specifies a target phrase (for example: "Sure, here is how to..."). The goal is to minimize the model loss on this target string.
2. Gradient Computation: The algorithm computes the gradient of the loss with respect to the one-hot token embeddings of the attack suffix.
3. Candidate Proposal: Using the gradient values, the algorithm identifies the top candidate tokens most likely to decrease the loss at each token position.
4. Candidate Evaluation: The algorithm evaluates a large batch of these candidate sequences across the target models using forward passes.
5. Selection: The candidate producing the lowest loss is selected, and the loop repeats until the model produces the target output.

In theory, this is a clean optimization loop. In practice, running hundreds of candidates across multiple neural networks quickly creates memory and architectural bottlenecks.

---

[3] System Bottlenecks in the Original Architecture

Before the refactoring, the optimization logic was concentrated inside two oversized methods: GCGMultiPromptAttack.step() and ProgressiveMultiPromptAttack.run().

This monolithic design introduced three major engineering problems:

1. High Coupling and Test Fragility

A single function was responsible for dispatching worker tasks across GPUs, aggregating incompatible gradient tensor shapes, sampling replacement tokens, executing model forward passes, and deciding step budgets. Because every phase was interwoven, writing fast unit tests was impossible without spinning up heavy model instances.

2. VRAM Spikes and Out-Of-Memory Crashes

Evaluating candidate tokens requires computing model logits across prompts and models. The original implementation held intermediate tensors in GPU memory throughout the step loop. During multi-prompt or multi-model runs, this caused sudden Out-Of-Memory (OOM) failures on GPU clusters.

3. Entangled Scheduling State

GCG uses a progressive curriculum: it starts on a single prompt and single model, then gradually admits more prompts, more models, and increases control weights. These state transitions were managed through nested conditionals, leading to edge cases where non-finite losses (NaN or Inf) could inadvertently trigger early exit.

---

[4] Architecture Redesign - Phase Extraction Strategy

Under tracking issue #2665, Microsoft maintainer Roman Lutz outlined the architectural requirements: decouple the pipeline into three independently reviewable, typed modules while maintaining 100% backward compatibility for all existing public APIs.

The refactoring was executed across three distinct engineering phases.

---

[5] Phase 1: Decoupling Candidate Proposal

Reference -> PR #2671

The first phase focused on extracting gradient aggregation, normalization, and candidate sampling logic out of GCGMultiPromptAttack.step().

Technical Implementation

First, I created an immutable data container using @dataclass(frozen=True, slots=True) to cleanly pass proposed candidates between modules:

@dataclass(frozen=True, slots=True)
class CandidateProposalBatch:
    """Grouped candidate control tokens and filtered text strings across compatible workers."""
    control_candidates_by_group: list[list[str]]
    group_worker_indices: list[int]

    @property
    def num_groups(self) -> int:
        return len(self.control_candidates_by_group)

Inside GCGCandidateProposer, we dispatch gradient computations across worker threads and normalize the gradients with an \(L_2\) norm so models with different loss magnitudes contribute equally:

# Pull gradient from worker and migrate to primary device
new_grad: torch.Tensor = worker.results.get().to(self._main_device)

# L2 normalization across embedding dimensions
new_grad = new_grad / new_grad.norm(dim=-1, keepdim=True)

Next came handling heterogeneous models. If an attack targets both LLaMA and Mistral simultaneously, their vocabularies and gradient tensor shapes differ. Instead of crashing, the proposer groups compatible shapes together:

# If vocabulary or tensor shape differs, flush the preceding group
if grad.shape != new_grad.shape or vocab != new_vocab:
    with torch.no_grad():
        sampled = self._sample_group(worker_idx=j - 1, gradient=grad, batch_size=batch_size, topk=topk)
        filtered = self._filter_group(worker_idx=j - 1, control_cand=sampled)
        control_cands.append(filtered)
        group_worker_indices.append(j - 1)
    grad = new_grad
    vocab = new_vocab
else:
    grad += new_grad

By extracting this into GCGCandidateProposer, we gained the ability to write 6 fast, deterministic unit tests for gradient aggregation and shape grouping without allocating real neural networks.

---

[6] Phase 2: VRAM-Bounded Candidate Evaluation

Reference -> PR #2700

The second phase addressed GPU memory exhaustion by isolating candidate evaluation into a dedicated, memory-efficient component: GCGCandidateEvaluator.

Technical Implementation

When evaluating 512 or 1,024 candidate suffixes across multiple models, storing all prediction logits simultaneously quickly exceeds the 24GB or 80GB VRAM limits of modern GPUs.

To fix this, the evaluation loop evaluates prompts sequentially and explicitly frees temporary tensors from GPU memory immediately after their losses are accumulated:

# pyrit/executor/promptgen/gcg/attack/gcg/candidate_evaluator.py

with torch.no_grad():
    for j, cand in enumerate(control_candidates_by_group):
        for i in range(num_prompts):
            # Dispatch logits computation to worker models
            for k, worker in enumerate(self._workers):
                worker(self._prompts[k][i], ModelWorkerOperation.LOGITS, cand, return_ids=True)

            # Retrieve worker logits and token IDs
            logits, ids = zip(*[worker.results.get() for worker in self._workers], strict=True)

            # Accumulate loss on main_device across all models
            loss[j * batch_size : (j + 1) * batch_size] += sum(
                self._loss_function.compute_loss(
                    logits=logit,
                    token_ids=token_ids,
                    target_slice=self._prompts[k][i]._target_slice,
                    control_slice=self._prompts[k][i]._control_slice,
                ).to(self._main_device)
                for k, (logit, token_ids) in enumerate(zip(logits, ids, strict=True))
            )

            # Explicit cleanup: delete bulky intermediate tensors immediately
            del logits, ids

Why does del logits, ids matter here? In Python, temporary variables inside loops can remain referenced until the next iteration. In PyTorch, logits tensors for large batch sizes occupy gigabytes of memory. Calling del logits, ids informs Python's garbage collector and the PyTorch caching allocator that these memory blocks can be reused immediately, bounding VRAM consumption to a constant size.

Additionally, the evaluator accounts for numerical edge cases such as NaN or Inf losses, ensuring stability during extreme optimization steps.

---

[7] Phase 3: Progressive Scheduling and State Machine

Reference -> PR #2720

The final phase extracted multi-model curriculum scheduling and budget accounting into ProgressiveScheduleController.

Technical Implementation

Instead of maintaining nested if/else flags inside ProgressiveMultiPromptAttack.run(), we created an explicit state machine and state container:

class ScheduleTransitionAction(Enum):
    """Action to take following an inner attack result in progressive scheduling."""
    CONTINUE = auto()
    FINALIZE_AND_STOP = auto()


@dataclass
class ProgressiveScheduleState:
    """Typed schedule state tracking curriculum admission and steps."""
    goals_admitted: int
    workers_admitted: int
    steps_completed: int = 0
    loss: float = float("inf")
    stop_inner_on_success: bool = False

The controller enforces an invariant in multi-model attacks: goals must be admitted before target model workers are expanded.

# pyrit/executor/promptgen/gcg/attack/base/progressive_schedule.py

def update_on_success(self, state: ProgressiveScheduleState) -> ScheduleTransitionAction:
    """Advance the progressive curriculum when an attack threshold is met."""
    if state.goals_admitted < self._total_goals:
        # Admit next attack goal
        state.goals_admitted += 1
        return ScheduleTransitionAction.CONTINUE
    elif state.workers_admitted < self._total_workers:
        # Admit next target model worker
        state.workers_admitted += 1
        return ScheduleTransitionAction.CONTINUE
    
    # All goals and workers admitted: finalize attack
    return ScheduleTransitionAction.FINALIZE_AND_STOP

The controller also manages exact step budgeting, guaranteeing the attack terminates cleanly when the user's iteration limit is reached without leaving hanging tasks.

---

[8] Code Quality, Reviews, and Maintainer Collaboration

Working on core infrastructure in a major open-source repository requires close collaboration with maintainers. Throughout this project, Microsoft maintainer Roman Lutz reviewed each pull request.

During the execution of Phase 2, Roman acknowledged the development velocity:

In Phase 3 (PR #2720), Roman identified an important compatibility detail regarding Python class definitions. When ProgressiveScheduleState was moved into the new module, an existing import in attack_manager.py shadowed the class. As a result, caller isinstance() checks returned False:

"The original ProgressiveScheduleState definition still shadows the new import. That means controller.state is a different class from the one exposed through the existing import path... Dataclass equality also fails between the two classes even when every field matches... A regression assertion using the existing import path would catch this."

I resolved this by removing the duplicate definition, re-exporting the imported class directly, and adding explicit test assertions covering both isinstance and dataclass equality:

# tests/unit/executor/promptgen/gcg/test_run_state.py

assert isinstance(attack.last_schedule_state, attack_manager.ProgressiveScheduleState)
assert attack.last_schedule_state == expected_state

Roman verified the fix and approved the PR:

"Confirmed in 2138122 that both import paths now expose the same class, restoring isinstance checks and dataclass equality for the returned state. This addresses the compatibility issue."

With all three PRs merged, Issue #2665 was officially closed, setting the stage for future performance optimizations:

---

[9] Results and Engineering Impact

The refactoring achieved significant measurable improvements:

• Clean Separation of Concerns: Replaced a 1,000-line monolithic loop with three modular, single-responsibility components.
• Eliminated VRAM Spikes: Sequential evaluation paired with immediate tensor disposal resolved GPU memory exhaustion across multi-model benchmarks.
• Zero Breaking Changes: Maintained complete backward compatibility for all public methods, parameter signatures, and configuration classes.
• Comprehensive Test Suite: Over 300 automated unit and integration tests now run deterministically across Windows, macOS, and Linux runners without requiring active GPUs.

---

[10] My Takeaways

1. Decouple Complex Algorithms Early: Large research algorithms tend to accumulate responsibilities over time. Separating proposal, evaluation, and scheduling makes codebases easier to maintain and extend.
2. Deterministic Cleanup in Accelerated Workloads: In PyTorch applications, relying solely on Python automatic garbage collection is insufficient for high-throughput loops. Explicitly deleting intermediate tensors prevents memory fragmentation.
3. Guard Class Contracts During Refactoring: Moving a dataclass or type definition can silently break isinstance checks if legacy export paths are not strictly validated with regression tests.

---

Thanks for Reading :)

Connect with me on LinkedIn