fix: stop lora_A from carrying over between trials (#476)
Some checks failed
CI / Check and build (Python 3.10) (push) Has been cancelled
CI / Check and build (Python 3.11) (push) Has been cancelled
CI / Check and build (Python 3.12) (push) Has been cancelled
CI / Check and build (Python 3.13) (push) Has been cancelled

* 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.
This commit is contained in:
benthecarman 2026-10-04 08:01:40 -04:00 • committed by GitHub
parent a09b4eae07
commit 208c0ca35b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 29 additions and 3 deletions

View file

@ -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.

View file

@ -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