Skip to content

Commit

Permalink
FIX TST Small regression in BNB LoRA output (#2238)
Browse files Browse the repository at this point in the history
Our regression tests reveal that the 8bit LoRA BNB regression test is
failing. To reproduce, run:

pytest tests/regression/test_regression.py -s --regression -k
test_lora_8bit

The regression was introduced in #2122. We didn't notice this earlier
because of other failing tests in the nightly CI.

The cause of the error is subtle. In the original code, we would
calculate the LoRA output, convert the dtype if necessary, then add it
to the base output. After the mentioned PR, we calculate the LoRA
output, add it to the base output, then convert the dtype if necessary.
The difference is very small on a per layer basis, but it can accumulate
over the layers, leading to a significant difference in outputs, as
witnessed by the regression test.

This PR rolls back this specific part of the PR (both for 8bit and 4bit)
while leaving the main change of that PR intact.
  • Loading branch information
BenjaminBossan authored Nov 28, 2024
1 parent 943daf1 commit 131efba
Showing 1 changed file with 8 additions and 6 deletions.
14 changes: 8 additions & 6 deletions src/peft/tuners/lora/bnb.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,15 +248,15 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
x = x.to(compute_dtype)

if not self.use_dora[active_adapter]:
result = result + lora_B(lora_A(dropout(x))) * scaling
output = lora_B(lora_A(dropout(x))) * scaling
else:
if isinstance(dropout, torch.nn.Identity) or not self.training:
base_result = result
else:
x = dropout(x)
base_result = None

result = result + self.lora_magnitude_vector[active_adapter](
output = self.lora_magnitude_vector[active_adapter](
x,
lora_A=lora_A,
lora_B=lora_B,
Expand All @@ -265,7 +265,8 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
base_result=base_result,
)
if requires_conversion:
result = result.to(expected_dtype)
output = output.to(expected_dtype)
result = result + output

return result

Expand Down Expand Up @@ -515,15 +516,15 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
x = x.to(lora_A.weight.dtype)

if not self.use_dora[active_adapter]:
result = result + lora_B(lora_A(dropout(x))) * scaling
output = lora_B(lora_A(dropout(x))) * scaling
else:
if isinstance(dropout, torch.nn.Identity) or not self.training:
base_result = result
else:
x = dropout(x)
base_result = None

result = result + self.lora_magnitude_vector[active_adapter](
output = self.lora_magnitude_vector[active_adapter](
x,
lora_A=lora_A,
lora_B=lora_B,
Expand All @@ -532,7 +533,8 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
base_result=base_result,
)
if requires_conversion:
result = result.to(expected_dtype)
output = output.to(expected_dtype)
result = result + output

return result

Expand Down

0 comments on commit 131efba

Please sign in to comment.