mirror of
https://github.com/p-e-w/heretic.git
synced 2026-10-14 20:33:44 +00:00
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.
This commit is contained in:
parent
a09b4eae07
commit
208c0ca35b
2 changed files with 29 additions and 3 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue