mirror of
https://github.com/p-e-w/heretic.git
synced 2026-10-15 15:10:08 -07: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:
+12
-3
@@ -221,6 +221,15 @@ class Model:
|
|||||||
# so the result is a PeftModel rather than a PeftMixedModel.
|
# so the result is a PeftModel rather than a PeftMixedModel.
|
||||||
self.model = cast(PeftModel, get_peft_model(self.model, self.peft_config))
|
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:
|
def _get_quantization_config(self, dtype: str) -> BitsAndBytesConfig | None:
|
||||||
"""
|
"""
|
||||||
Creates quantization config based on settings.
|
Creates quantization config based on settings.
|
||||||
@@ -315,9 +324,9 @@ class Model:
|
|||||||
|
|
||||||
if current_model == self.settings.model and not self.needs_reload:
|
if current_model == self.settings.model and not self.needs_reload:
|
||||||
# Reset LoRA adapters to zero (identity transformation).
|
# Reset LoRA adapters to zero (identity transformation).
|
||||||
for name, module in self.model.named_modules():
|
# Zeroing A as well ensures that adapter weights of modules not modified
|
||||||
if "lora_B" in name and hasattr(module, "weight"):
|
# by the next trial don't depend on what previous trials did.
|
||||||
torch.nn.init.zeros_(module.weight)
|
self._zero_lora_weights()
|
||||||
return True
|
return True
|
||||||
|
|
||||||
# Purge existing model object from memory to make space.
|
# Purge existing model object from memory to make space.
|
||||||
|
|||||||
@@ -4,6 +4,7 @@
|
|||||||
# Arbitrary-Rank Ablation (ARA) (Weidmann 2026)
|
# Arbitrary-Rank Ablation (ARA) (Weidmann 2026)
|
||||||
# See https://github.com/p-e-w/heretic/pull/211 for more information.
|
# See https://github.com/p-e-w/heretic/pull/211 for more information.
|
||||||
|
|
||||||
|
import math
|
||||||
from dataclasses import asdict, dataclass
|
from dataclasses import asdict, dataclass
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
@@ -247,6 +248,11 @@ class ARA(Modifier[Parameters]):
|
|||||||
def modify_model(self, ctx: Context, parameters: Parameters) -> None:
|
def modify_model(self, ctx: Context, parameters: Parameters) -> None:
|
||||||
model = ctx.get_model()
|
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(
|
for layer_index in range(
|
||||||
parameters.start_layer_index,
|
parameters.start_layer_index,
|
||||||
parameters.end_layer_index,
|
parameters.end_layer_index,
|
||||||
@@ -280,6 +286,17 @@ class ARA(Modifier[Parameters]):
|
|||||||
lora_A = cast(Tensor, module.lora_A["default"].weight)
|
lora_A = cast(Tensor, module.lora_A["default"].weight)
|
||||||
lora_B = cast(Tensor, module.lora_B["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.
|
# Move I/O tensors to the device of the adapter weights.
|
||||||
good_input, good_output = self.good_module_io[layer_index][
|
good_input, good_output = self.good_module_io[layer_index][
|
||||||
component
|
component
|
||||||
|
|||||||
Reference in New Issue
Block a user