From 208c0ca35b3feade91dc74755cb80b529affa325 Mon Sep 17 00:00:00 2001
From: benthecarman <15256660+benthecarman@users.noreply.github.com>
Date: Sun, 4 Oct 2026 08:01:40 -0400
Subject: [PATCH] fix: stop `lora_A` from carrying over between trials (#476)

* Re-initialize ARA's lora_A deterministically

ARA optimizes both LoRA matrices in place, but the model reset between
trials only zeroes lora_B. Each trial therefore started its
optimization from whatever lora_A the previous trial left behind.
Restoring a trial after the study re-optimized it from a different
starting point, so the exported model was not the one that earned the
trial's scores, and reproduction could not match hashes.

Re-initialize lora_A for every optimized module from a generator
seeded with the run seed. This makes the starting point independent of
trial history, the device and the global RNG state.

* Zero lora_A when resetting the model

The fast reset path only zeroed lora_B, and a fresh adapter gets lora_A
from PEFT's random init. Modules a trial doesn't modify therefore kept
whatever lora_A an earlier trial left, or a random one. The model is
unaffected since B is zero, but saved adapters differed byte-for-byte
depending on trial history, breaking hash verification for adapter
exports.

Zero lora_A as well, both on reset and right after applying LoRA.
---
 src/heretic/model.py         | 15 ++++++++++++---
 src/heretic/modifiers/ara.py | 17 +++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/heretic/model.py b/src/heretic/model.py
index 951aa89..aaac268 100644
--- a/src/heretic/model.py
+++ b/src/heretic/model.py
@@ -221,6 +221,15 @@ class Model:
         # so the result is a PeftModel rather than a PeftMixedModel.
         self.model = cast(PeftModel, get_peft_model(self.model, self.peft_config))
 
+        # PEFT initializes A randomly from the global RNG. B is zero, so this doesn't
+        # change the model, but zeroing A too keeps saved adapters reproducible.
+        self._zero_lora_weights()
+
+    def _zero_lora_weights(self):
+        for name, module in self.model.named_modules():
+            if ("lora_A" in name or "lora_B" in name) and hasattr(module, "weight"):
+                torch.nn.init.zeros_(module.weight)
+
     def _get_quantization_config(self, dtype: str) -> BitsAndBytesConfig | None:
         """
         Creates quantization config based on settings.
@@ -315,9 +324,9 @@ class Model:
 
         if current_model == self.settings.model and not self.needs_reload:
             # Reset LoRA adapters to zero (identity transformation).
-            for name, module in self.model.named_modules():
-                if "lora_B" in name and hasattr(module, "weight"):
-                    torch.nn.init.zeros_(module.weight)
+            # Zeroing A as well ensures that adapter weights of modules not modified
+            # by the next trial don't depend on what previous trials did.
+            self._zero_lora_weights()
             return True
 
         # Purge existing model object from memory to make space.
diff --git a/src/heretic/modifiers/ara.py b/src/heretic/modifiers/ara.py
index 428c4ab..a3b7aaf 100644
--- a/src/heretic/modifiers/ara.py
+++ b/src/heretic/modifiers/ara.py
@@ -4,6 +4,7 @@
 # Arbitrary-Rank Ablation (ARA) (Weidmann 2026)
 # See https://github.com/p-e-w/heretic/pull/211 for more information.
 
+import math
 from dataclasses import asdict, dataclass
 from typing import Any, cast
 
@@ -247,6 +248,11 @@ class ARA(Modifier[Parameters]):
     def modify_model(self, ctx: Context, parameters: Parameters) -> None:
         model = ctx.get_model()
 
+        # The model's reset zeroes both A and B, where all gradients vanish,
+        # so A must be initialized below for the optimization to make progress.
+        # Generating on the CPU keeps A independent of the device and the global RNG state.
+        generator = torch.Generator().manual_seed(cast(int, self.heretic_settings.seed))
+
         for layer_index in range(
             parameters.start_layer_index,
             parameters.end_layer_index,
@@ -280,6 +286,17 @@ class ARA(Modifier[Parameters]):
                     lora_A = cast(Tensor, module.lora_A["default"].weight)
                     lora_B = cast(Tensor, module.lora_B["default"].weight)
 
+                    # Initialize A the same way PEFT does, but deterministically. See
+                    # https://github.com/huggingface/peft/blob/v0.21.2/src/peft/tuners/lora/layer.py#L338
+                    initial_A = torch.empty(lora_A.shape)
+                    torch.nn.init.kaiming_uniform_(
+                        initial_A,
+                        a=math.sqrt(5),
+                        generator=generator,
+                    )
+                    with torch.no_grad():
+                        lora_A.copy_(initial_A)
+
                     # Move I/O tensors to the device of the adapter weights.
                     good_input, good_output = self.good_module_io[layer_index][
                         component