FIX PVeRA forward implementation for bitsandbytes (#3189)

In the forward path of PVeRA, a sampling step is required. For the bnb
layers, this was missing. The PR now adds the sampling step.

This fixes failing tests in the nightly CI.

Author: BenjaminBossan

src/peft/tuners/pvera/bnb.py

@@ -18,6 +18,7 @@
 
 import bitsandbytes as bnb
 import torch
+import torch.nn.functional as F
 
 from peft.import_utils import is_bnb_4bit_available, is_bnb_available
 from peft.tuners.tuners_utils import check_adapters_to_merge
@@ -44,6 +45,7 @@ def __init__(
             super().__init__()
             PveraLayer.__init__(self, base_layer)
             self.fan_in_fan_out = config.fan_in_fan_out
+            self.sample_at_inference = config.sample_at_inference
 
             self._active_adapter = adapter_name
             self.update_layer(
@@ -223,9 +225,9 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
                     sliced_B = pvera_B[: self.out_features, :].to(x.device)
 
                     x_temp = dropout(x.to(lambda_d.dtype))
-
-                    adapter_output = lambda_b * torch.nn.functional.linear(
-                        lambda_d * torch.nn.functional.linear(x_temp, sliced_A), sliced_B
+                    mu, logvar = (lambda_d * F.linear(x_temp, sliced_A)).chunk(2, dim=-1)
+                    adapter_output = lambda_b * F.linear(
+                        self._reparametrize(mu, logvar, self.sample_at_inference), sliced_B
                     )
 
                     if requires_conversion:
@@ -257,6 +259,7 @@ def __init__(
             super().__init__()
             PveraLayer.__init__(self, base_layer)
             self.fan_in_fan_out = config.fan_in_fan_out
+            self.sample_at_inference = config.sample_at_inference
 
             self._active_adapter = adapter_name
             self.update_layer(
@@ -392,9 +395,9 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
                     sliced_B = pvera_B[: self.out_features, :].to(x.device)
 
                     x_temp = dropout(x.to(lambda_d.dtype))
-
-                    adapter_output = lambda_b * torch.nn.functional.linear(
-                        lambda_d * torch.nn.functional.linear(x_temp, sliced_A), sliced_B
+                    mu, logvar = (lambda_d * F.linear(x_temp, sliced_A)).chunk(2, dim=-1)
+                    adapter_output = lambda_b * F.linear(
+                        self._reparametrize(mu, logvar, self.sample_at_inference), sliced_B
                     )
 
                     if requires_conversion:

src/peft/tuners/pvera/layer.py

@@ -143,6 +143,15 @@ def reset_pvera_parameters(self, adapter_name, d_initial: float = 0.1):
                 nn.init.zeros_(self.pvera_lambda_d[adapter_name]).fill_(d_initial)
                 nn.init.zeros_(self.pvera_lambda_b[adapter_name])
 
+    def _reparametrize(self, mu, logvar, sample_at_inference):
+        if self.training or (not self.training and sample_at_inference):
+            std = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(std)
+            z = mu + eps * std
+        else:
+            z = mu
+        return z
+
 
 class Linear(nn.Linear, PveraLayer):
     # PVeRA implemented in a dense layer
@@ -259,15 +268,6 @@ def get_delta_weight(self, adapter) -> torch.Tensor:
 
         return output_tensor
 
-    def _reparametrize(self, mu, logvar, sample_at_inference):
-        if self.training or (not self.training and sample_at_inference):
-            std = torch.exp(0.5 * logvar)
-            eps = torch.randn_like(std)
-            z = mu + eps * std
-        else:
-            z = mu
-        return z
-
     def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         previous_dtype = x.dtype