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