FIX Propagate GCG random_seed to all RNG sources for deterministic runs

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

@@ -98,6 +98,17 @@ class ProgressiveScheduleState:
     stop_inner_on_success: bool = False
 
 
+@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]
+
+
 class NpEncoder(json.JSONEncoder):
     """Encode NumPy scalar and array values for JSON output."""
 
@@ -995,17 +1006,31 @@ 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:
+            workers = getattr(self, "workers", [])
+            try:
+                sampling_device = workers[0].model.device
+                self._torch_gens = {
+                    i: torch.Generator(device=sampling_device).manual_seed(random_seed + i) for i in range(len(workers))
+                }
+            except (TypeError, AttributeError, IndexError):
+                self._torch_gens = {i: torch.Generator().manual_seed(random_seed + i) for i in range(len(workers))}
 
         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 +1403,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 +1435,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 +1446,25 @@ def run(
         # not keep looking current.
         self.last_schedule_state = None
 
+        rng_bundle = getattr(self, "_rng_bundle", None)
+        if rng_bundle is None:
+            derived_seeds = {i: random_seed + i for i in range(len(self.workers))}
+            try:
+                sampling_device = self.workers[0].model.device
+                torch_gens = {
+                    i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i])
+                    for i in range(len(self.workers))
+                }
+            except (TypeError, AttributeError):
+                torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(self.workers))}
+            rng_bundle = RngBundle(
+                np_rng=np.random.default_rng(random_seed),
+                py_rng=random.Random(random_seed),
+                torch_gens=torch_gens,
+                base_seed=random_seed,
+                derived_seeds=derived_seeds,
+            )
+
         _update_attack_log_params(
             logfile=self.logfile,
             params={
@@ -1432,6 +1479,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,
             },
         )
 
@@ -1462,6 +1511,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 +1527,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 +1685,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 +1717,31 @@ 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)
+        if rng_bundle is None:
+            derived_seeds = {i: random_seed + i for i in range(len(self.workers))}
+            try:
+                sampling_device = self.workers[0].model.device
+                torch_gens = {
+                    i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i])
+                    for i in range(len(self.workers))
+                }
+            except (TypeError, AttributeError):
+                torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(self.workers))}
+            rng_bundle = RngBundle(
+                np_rng=np.random.default_rng(random_seed),
+                py_rng=random.Random(random_seed),
+                torch_gens=torch_gens,
+                base_seed=random_seed,
+                derived_seeds=derived_seeds,
+            )
+
         _update_attack_log_params(
             logfile=self.logfile,
             params={
@@ -1683,6 +1756,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,
             },
         )
 
@@ -1703,6 +1778,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 +1795,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