FIX Use correct argument name in MHA forward (#2510)

BenjaminBossan · web-flow · commit eb5e9bcbdf9a · 2025-05-05T15:40:12.000+02:00
The arguments of the forward method of MultiheadAttention are called query etc. PEFT used x. Therefore, if a caller uses keywords only, the argument is not assigned, resulting in an error. This was initially reported here: #761 (comment) Note: Other layers' forward method (like Linear) also uses incorrect names, like x instead of input, but so far no issues were reported, so I'll leave it as is for now.
diff --git a/src/peft/tuners/lora/layer.py b/src/peft/tuners/lora/layer.py
@@ -1648,16 +1648,16 @@ def _check_forward_args(self, x, *args, **kwargs):
             raise TypeError(f"lora.{self.__class__.__name__} does not support mixed adapter batches.")
         super()._check_forward_args(x, *args, **kwargs)
 
-    def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
-        previous_dtype = x.dtype
-        self._check_forward_args(x, *args, **kwargs)
+    def forward(self, query: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
+        previous_dtype = query.dtype
+        self._check_forward_args(query, *args, **kwargs)
 
         if self.disable_adapters:
             if self.merged:
                 self.unmerge()
-            result = self.base_layer(x, *args, **kwargs)
+            result = self.base_layer(query, *args, **kwargs)
         elif self.merged:
-            result = self.base_layer(x, *args, **kwargs)
+            result = self.base_layer(query, *args, **kwargs)
         else:
             out_proj = self.get_base_layer().out_proj
             if out_proj.active_adapters != self.active_adapters:
@@ -1680,7 +1680,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
             active_adapters = [a for a in self.active_adapters if a in self.lora_A]
             try:
                 self.merge(adapter_names=active_adapters)
-                result = self.base_layer(x, *args, **kwargs)
+                result = self.base_layer(query, *args, **kwargs)
             finally:
                 # it's safe to call unmerge(), which unmerges all adapters, because we checked that not self.merged,
                 # i.e. there is was no merged layer before
