fix: replace home-cooked set_seed function with Transformers builtin

src/heretic/main.py

@@ -82,7 +82,6 @@ from .utils import (
     load_prompts,
     print,
     print_memory_usage,
-    set_seed,
     upload_reproduce_folder,
 )
 
@@ -257,7 +256,7 @@ def run():
     if settings.seed is None:
         settings.seed = random.randint(0, 2**32 - 1)
 
-    set_seed(settings.seed)
+    transformers.set_seed(settings.seed)
 
     print(get_accelerator_info())
 

src/heretic/model.py

@@ -580,11 +580,16 @@ class Model:
                         W = W - W_org
                         # Use a low-rank SVD to get an approximation of the matrix.
                         r = self.peft_config.r
+
                         # svd_lowrank is randomized:
                         # https://github.com/pytorch/pytorch/blob/20919052303c0b5ba87f8bf7e19237dc33ab09d3/torch/_lowrank.py#L108-L109
                         # Reseed immediately before the call so restoring a trial is independent of RNG history.
                         torch.manual_seed(self.settings.seed)
+                        # "It's safe to call this function if CUDA is not available;
+                        # in that case, it is silently ignored."
+                        torch.cuda.manual_seed_all(self.settings.seed)  # ty:ignore[invalid-argument-type]
                         U, S, Vh = torch.svd_lowrank(W, q=2 * r + 4, niter=6)
+
                         # Truncate it to the part we want to store in the LoRA adapter.
                         # Note: svd_lowrank actually returns V, so transpose it to get Vh.
                         U = U[:, :r]

src/heretic/utils.py

@@ -5,7 +5,6 @@ import hashlib
 import json
 import os
 import platform
-import random
 import tempfile
 import traceback
 from dataclasses import dataclass
@@ -15,7 +14,6 @@ from pathlib import Path
 from typing import TypeVar
 
 import huggingface_hub
-import numpy as np
 import tomli_w
 import torch
 from datasets import DatasetDict, ReadInstruction, load_dataset, load_from_disk
@@ -301,14 +299,6 @@ def generate_requirements_txt() -> str:
     return "\n".join(requirements) + "\n"
 
 
-def set_seed(seed: int):
-    """Sets the seed for all RNGs."""
-
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-
 def format_hf_link(
     path: str,
     commit: str | None = None,