FIX Propagate GCG random_seed to all RNG sources for deterministic runs

pyrit/executor/promptgen/gcg/attack/base/attack_manager.py

@@ -995,17 +995,30 @@ def run(
         log_first: bool = False,
         filter_cand: bool = True,
         verbose: bool = True,
+        random_seed: int = 42,
     ) -> tuple[str, float, int]:
         """
         Run iterative optimization.
 
         Returns:
             tuple[str, float, int]: The final control, loss, and step count.
         """
+        rng_bundle = getattr(self, "_rng_bundle", None)
+        py_rng = rng_bundle.py_rng if rng_bundle else random.Random(random_seed)
+        if rng_bundle:
+            self._torch_gens = rng_bundle.torch_gens
+        else:
+            try:
+                self._torch_gens = {
+                    i: torch.Generator(device=self.models[i].device).manual_seed(random_seed + i)
+                    for i in range(len(self.workers))
+                }
+            except (AttributeError, TypeError):
+                self._torch_gens = {0: torch.Generator().manual_seed(random_seed)}
 
         def acceptance_probability(e: float, e_prime: float, k: int) -> bool:
             temperature = max(1 - float(k + 1) / (n_steps + anneal_from), 1.0e-7)
-            return e_prime < e or math.exp(-(e_prime - e) / temperature) >= random.random()
+            return e_prime < e or math.exp(-(e_prime - e) / temperature) >= py_rng.random()
 
         if target_weight is None:
 
@@ -1378,6 +1391,7 @@ def run(
         stop_on_success: bool = True,
         verbose: bool = True,
         filter_cand: bool = True,
+        random_seed: int = 42,
     ) -> tuple[str, int]:
         """
         Execute the progressive multi-prompt attack.
@@ -1409,6 +1423,8 @@ def run(
                 Whether to print verbose output (default is True)
             filter_cand (bool, optional):
                 Whether to filter candidates whose lengths changed after re-tokenization (default is True)
+            random_seed (int, optional):
+                Seed for deterministic random number generation (default is 42)
 
         Returns:
             tuple[str, int]: The final control suffix and completed step count.
@@ -1418,6 +1434,8 @@ def run(
         # not keep looking current.
         self.last_schedule_state = None
 
+        rng_bundle = getattr(self, "_rng_bundle", None)
+
         _update_attack_log_params(
             logfile=self.logfile,
             params={
@@ -1432,6 +1450,8 @@ def run(
                 "anneal": anneal,
                 "incr_control": incr_control,
                 "stop_on_success": stop_on_success,
+                "random_seed": random_seed,
+                "derived_seeds": rng_bundle.derived_seeds if rng_bundle else {},
             },
         )
 
@@ -1462,6 +1482,7 @@ def run(
             )
             if schedule.goals_admitted == len(self.goals) and schedule.workers_admitted == len(self.workers):
                 schedule.stop_inner_on_success = False
+            attack._rng_bundle = rng_bundle
             inner_result: tuple[str, float, int] = attack.run(
                 n_steps=n_steps - schedule.steps_completed,
                 batch_size=batch_size,
@@ -1477,6 +1498,7 @@ def run(
                 test_steps=test_steps,
                 filter_cand=filter_cand,
                 verbose=verbose,
+                random_seed=random_seed,
             )
             control, inner_loss, inner_steps = inner_result
             schedule.loss = inner_loss
@@ -1634,6 +1656,7 @@ def run(
         stop_on_success: bool = True,
         verbose: bool = True,
         filter_cand: bool = True,
+        random_seed: int = 42,
     ) -> tuple[str, int]:
         """
         Execute the individual-prompt attack.
@@ -1665,10 +1688,14 @@ def run(
                 Whether to print verbose output (default is True)
             filter_cand (bool, optional):
                 Whether to filter candidates (default is True)
+            random_seed (int, optional):
+                Seed for deterministic random number generation (default is 42)
 
         Returns:
             tuple[str, int]: The final control suffix and configured step count.
         """
+        rng_bundle = getattr(self, "_rng_bundle", None)
+
         _update_attack_log_params(
             logfile=self.logfile,
             params={
@@ -1683,6 +1710,8 @@ def run(
                 "anneal": anneal,
                 "incr_control": incr_control,
                 "stop_on_success": stop_on_success,
+                "random_seed": random_seed,
+                "derived_seeds": rng_bundle.derived_seeds if rng_bundle else {},
             },
         )
 
@@ -1703,6 +1732,7 @@ def run(
                 self.test_targets,
                 self.test_workers,
             )
+            attack._rng_bundle = rng_bundle
             attack.run(
                 n_steps=n_steps,
                 batch_size=batch_size,
@@ -1719,6 +1749,7 @@ def run(
                 log_first=True,
                 filter_cand=filter_cand,
                 verbose=verbose,
+                random_seed=random_seed,
             )
 
         return self.control, n_steps

pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import inspect
 import logging
 from typing import Any
 
@@ -104,6 +105,7 @@ def sample_control(
         topk: int = 256,
         temp: float = 1.0,
         allow_non_ascii: bool = True,
+        torch_generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         """
         Sample new control token candidates based on gradients.
@@ -114,6 +116,7 @@ def sample_control(
             topk (int): Number of top gradient positions to sample from. Defaults to 256.
             temp (float): Temperature for sampling. Currently unused but kept for API compatibility. Defaults to 1.0.
             allow_non_ascii (bool): Whether to allow non-ASCII tokens. Defaults to True.
+            torch_generator (torch.Generator | None): Optional generator for deterministic sampling.
 
         Returns:
             torch.Tensor: Batch of new candidate control token sequences.
@@ -127,7 +130,9 @@ def sample_control(
             torch.int64
         )
         new_token_val = torch.gather(
-            top_indices[new_token_pos], 1, torch.randint(0, topk, (batch_size, 1), device=grad.device)
+            top_indices[new_token_pos],
+            1,
+            torch.randint(0, topk, (batch_size, 1), device=grad.device, generator=torch_generator),
         )
         return original_control_toks.scatter_(1, new_token_pos.unsqueeze(-1), new_token_val)
 
@@ -199,15 +204,22 @@ def _sample_control_candidates(
     ) -> torch.Tensor:
         sampler = self._resolve_sampling()
         prompt_manager = self.prompts[worker_index]
-        return sampler.sample_candidates(
-            gradient=gradient,
-            control_tokens=prompt_manager.control_toks,
-            batch_size=batch_size,
-            top_k=topk,
-            temperature=temp,
-            allow_non_ascii=allow_non_ascii,
-            non_ascii_tokens=prompt_manager.disallowed_toks,
-        )
+        torch_gens = getattr(self, "_torch_gens", None) or {}
+        torch_gen = torch_gens.get(worker_index)
+        kwargs: dict[str, Any] = {
+            "gradient": gradient,
+            "control_tokens": prompt_manager.control_toks,
+            "batch_size": batch_size,
+            "top_k": topk,
+            "temperature": temp,
+            "allow_non_ascii": allow_non_ascii,
+            "non_ascii_tokens": prompt_manager.disallowed_toks,
+        }
+        if torch_gen is not None:
+            sig = inspect.signature(sampler.sample_candidates)
+            if "torch_generator" in sig.parameters:
+                kwargs["torch_generator"] = torch_gen
+        return sampler.sample_candidates(**kwargs)
 
     def _filter_control_candidates(
         self,

pyrit/executor/promptgen/gcg/default_implementations.py

@@ -57,6 +57,7 @@ def sample_candidates(
         temperature: float,
         allow_non_ascii: bool,
         non_ascii_tokens: torch.Tensor,
+        torch_generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         """
         Sample ``batch_size`` candidate suffix token sequences.
@@ -79,6 +80,8 @@ def sample_candidates(
                 the top-k.
             non_ascii_tokens (torch.Tensor): Token ids to exclude when
                 ``allow_non_ascii`` is False.
+            torch_generator (torch.Generator | None): Optional generator for
+                deterministic sampling.
 
         Returns:
             torch.Tensor: Candidate suffix token sequences with shape
@@ -99,7 +102,7 @@ def sample_candidates(
         new_token_val = torch.gather(
             top_indices[new_token_pos],
             1,
-            torch.randint(0, top_k, (batch_size, 1), device=gradient.device),
+            torch.randint(0, top_k, (batch_size, 1), device=gradient.device, generator=torch_generator),
         )
         return original_control_tokens.scatter_(1, new_token_pos.unsqueeze(-1), new_token_val)
 

pyrit/executor/promptgen/gcg/extension_protocols.py

@@ -76,6 +76,7 @@ def sample_candidates(
         temperature: float,
         allow_non_ascii: bool,
         non_ascii_tokens: torch.Tensor,
+        torch_generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         """
         Sample ``batch_size`` candidate suffix token sequences.
@@ -101,6 +102,9 @@ def sample_candidates(
             non_ascii_tokens (torch.Tensor): Token ids to exclude when
                 ``allow_non_ascii`` is False, shape ``(num_disallowed,)``
                 and integer dtype.
+            torch_generator (torch.Generator | None): Optional random number
+                generator for deterministic sampling. When provided, all
+                random tensor operations should use this generator.
 
         Returns:
             torch.Tensor: Candidate suffix token sequences with shape

pyrit/executor/promptgen/gcg/generator.py

@@ -36,12 +36,14 @@
 import asyncio
 import json
 import logging
+import random
 import time
 from dataclasses import dataclass, field
 from functools import partial
 from typing import Any, overload
 
 import numpy as np
+import torch
 import torch.multiprocessing as mp
 from pydantic import Field
 
@@ -69,6 +71,17 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class RngBundle:
+    """Per-run RNG state bundle for deterministic GCG execution."""
+
+    np_rng: np.random.Generator
+    py_rng: random.Random
+    torch_gens: dict[int, torch.Generator]
+    base_seed: int
+    derived_seeds: dict[int, int]
+
+
 @dataclass
 class GCGContext(PromptGeneratorStrategyContext):
     """
@@ -97,6 +110,7 @@ class GCGContext(PromptGeneratorStrategyContext):
     test_workers: list[Any] = field(default_factory=list)
     attack: Any | None = None
     logfile_path: str | None = None
+    rng_bundle: RngBundle | None = None
 
 
 class GCGResult(PromptGeneratorStrategyResult):
@@ -260,17 +274,33 @@ async def _setup_async(self, *, context: GCGContext) -> None:
         self._ensure_spawn_start_method()
         context.memory_labels = combine_dict({}, context.memory_labels)
 
-        context.targets, context.test_targets = self._apply_target_augmentation(
-            train_targets=context.targets,
-            test_targets=context.test_targets,
-        )
-
         log_gpu_memory(step=0)
         log_train_goals(train_goals=context.goals)
 
         params = self._to_attack_params(context=context)
         context.workers, context.test_workers = await asyncio.to_thread(get_workers, params)
 
+        seed = self._algorithm.random_seed
+        derived_seeds = {i: seed + i for i in range(len(context.workers))}
+        try:
+            torch_gens = {
+                i: torch.Generator(device=context.workers[i].model.device).manual_seed(derived_seeds[i])
+                for i in range(len(context.workers))
+            }
+        except (TypeError, AttributeError):
+            torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(context.workers))}
+        context.rng_bundle = RngBundle(
+            np_rng=np.random.default_rng(seed),
+            py_rng=random.Random(seed),
+            torch_gens=torch_gens,
+            base_seed=seed,
+            derived_seeds=derived_seeds,
+        )
+
+        context.targets, context.test_targets = self._apply_target_augmentation(
+            train_targets=context.targets, test_targets=context.test_targets, np_rng=context.rng_bundle.np_rng
+        )
+
     async def _perform_async(self, *, context: GCGContext) -> GCGResult:
         """
         Build the attack, run the optimization loop, and read the result back.
@@ -303,6 +333,8 @@ async def _perform_async(self, *, context: GCGContext) -> GCGResult:
             logfile_path=context.logfile_path,
         )
 
+        context.attack._rng_bundle = context.rng_bundle
+
         await asyncio.to_thread(
             context.attack.run,
             n_steps=self._algorithm.n_steps,
@@ -318,6 +350,7 @@ async def _perform_async(self, *, context: GCGContext) -> GCGResult:
             verbose=self._output.verbose,
             filter_cand=self._algorithm.filter_cand,
             allow_non_ascii=self._algorithm.allow_non_ascii,
+            random_seed=self._algorithm.random_seed,
         )
 
         return self._read_result(logfile_path=context.logfile_path, memory_labels=context.memory_labels)
@@ -381,22 +414,25 @@ def _apply_target_augmentation(
         *,
         train_targets: list[str],
         test_targets: list[str],
+        np_rng: np.random.Generator | None = None,
     ) -> tuple[list[str], list[str]]:
         """
         Randomly substitute equivalent target phrasings for diversity.
 
         Returns:
             tuple[list[str], list[str]]: Augmented training and held-out targets.
         """
+        if np_rng is None:
+            np_rng = np.random.default_rng()
 
         def _shorten(s: str) -> str:
             return s.replace("Sure, h", "H")
 
         def _contract(s: str) -> str:
             return s.replace("Sure, here is", "Sure, here's")
 
-        train_targets = [_shorten(t) if np.random.random() < 0.5 else _contract(t) for t in train_targets]
-        test_targets = [_shorten(t) if np.random.random() < 0.5 else _contract(t) for t in test_targets]
+        train_targets = [_shorten(t) if np_rng.random() < 0.5 else _contract(t) for t in train_targets]
+        test_targets = [_shorten(t) if np_rng.random() < 0.5 else _contract(t) for t in test_targets]
         return train_targets, test_targets
 
     def _to_attack_params(self, *, context: GCGContext) -> Any: