FIX Propagate GCG random_seed to all RNG sources for deterministic runs

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

@@ -1016,17 +1016,22 @@ 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.
         """
+        py_rng = random.Random(random_seed)
+        models = getattr(self, "models", None)
+        device = models[0].device if models else "cpu"
+        self._torch_gen = torch.Generator(device=device).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:
 
@@ -1400,6 +1405,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.
@@ -1431,6 +1437,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.
@@ -1499,6 +1507,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
@@ -1656,6 +1665,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.
@@ -1687,6 +1697,8 @@ 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.
@@ -1741,6 +1753,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

@@ -104,6 +104,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 +115,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 +129,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 +203,19 @@ 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_gen: torch.Generator | None = getattr(self, "_torch_gen", None)
+        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:
+            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

@@ -260,9 +260,11 @@ async def _setup_async(self, *, context: GCGContext) -> None:
         self._ensure_spawn_start_method()
         context.memory_labels = combine_dict({}, context.memory_labels)
 
+        np_rng = np.random.default_rng(self._algorithm.random_seed)
         context.targets, context.test_targets = self._apply_target_augmentation(
             train_targets=context.targets,
             test_targets=context.test_targets,
+            np_rng=np_rng,
         )
 
         log_gpu_memory(step=0)
@@ -318,6 +320,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 +384,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: