Make lora torch.compile optional for unmerged lora application

This change has caused issues especially with LoRAs that have dynamic rank. Will now be disabled by default, to allow full graph with unmerged LoRAs the option to allow compile is available in the Torch Compile Settings -node

custom_linear.py

@@ -3,15 +3,20 @@ import torch.nn as nn
 from accelerate import init_empty_weights
 
 #based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/quantizers/gguf/utils.py
-def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, scale_weights=None):
+def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, scale_weights=None, compile_args=None):
    
     has_children = list(model.children())
     if not has_children:
         return
+    
+    allow_compile = False
+
     for name, module in model.named_children():
+        if compile_args is not None:
+            allow_compile = compile_args.get("allow_unmerged_lora_compile", False)
         module_prefix = prefix + name + "."
         module_prefix = module_prefix.replace("_orig_mod.", "")
-        _replace_linear(module, compute_dtype, state_dict, module_prefix, patches, scale_weights)
+        _replace_linear(module, compute_dtype, state_dict, module_prefix, patches, scale_weights, compile_args)
 
         if isinstance(module, nn.Linear) and "loras" not in module_prefix:
             in_features = state_dict[module_prefix + "weight"].shape[1]
@@ -25,7 +30,8 @@ def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, s
                     out_features,
                     module.bias is not None,
                     compute_dtype=compute_dtype,
-                    scale_weight=scale_weights.get(scale_key) if scale_weights else None
+                    scale_weight=scale_weights.get(scale_key) if scale_weights else None,
+                    allow_compile=allow_compile
                 )
             model._modules[name].source_cls = type(module)
             model._modules[name].requires_grad_(False)
@@ -77,7 +83,8 @@ class CustomLinear(nn.Linear):
         bias=False,
         compute_dtype=None,
         device=None,
-        scale_weight=None
+        scale_weight=None,
+        allow_compile=False
     ) -> None:
         super().__init__(in_features, out_features, bias, device)
         self.compute_dtype = compute_dtype
@@ -85,6 +92,10 @@ class CustomLinear(nn.Linear):
         self.step = 0
         self.scale_weight = scale_weight
         self.lora_strengths = []
+        self.allow_compile = allow_compile
+
+        if not allow_compile:
+            self._get_weight_with_lora = torch.compiler.disable()(self._get_weight_with_lora)
     
     def set_lora_diffs(self, lora_diffs, device=torch.device("cpu")):
         self.lora_diffs = []
@@ -98,25 +109,11 @@ class CustomLinear(nn.Linear):
                 self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
                 self.lora_diffs.append(f"lora_diff_{i}_0")
 
-    def forward(self, input):
-        if self.bias is not None:
-            bias = self.bias.to(input)
-        else:
-            bias = None
-        weight = self.weight.to(input)
-
-        if self.scale_weight is not None:
-            if weight.numel() < input.numel():
-                weight = weight * self.scale_weight
-            else:
-                input = input * self.scale_weight
-
-        if hasattr(self, f"lora_diff_0_0"):
-            weight = self.apply_lora(weight).to(self.compute_dtype)
-
-        return torch.nn.functional.linear(input, weight, bias)
-
-    def apply_lora(self, weight):
+    def _get_weight_with_lora(self, weight):
+        """Apply LoRA outside compiled region"""
+        if not hasattr(self, "lora_diff_0_0"):
+            return weight
+        
         for lora_diff_names, lora_strength in zip(self.lora_diffs, self.lora_strengths):
             if isinstance(lora_strength, list):
                 lora_strength = lora_strength[self.step]
@@ -131,7 +128,7 @@ class CustomLinear(nn.Linear):
                 patch_diff = torch.mm(
                     lora_diff_0.flatten(start_dim=1),
                     lora_diff_1.flatten(start_dim=1)
-                ).reshape(weight.shape) + 0
+                ).reshape(weight.shape)
                 alpha = lora_diff_2 / lora_diff_1.shape[0] if lora_diff_2 is not None else 1.0
                 scale = lora_strength * alpha
                 weight = weight.add(patch_diff, alpha=scale)
@@ -139,6 +136,26 @@ class CustomLinear(nn.Linear):
                 lora_diff = getattr(self, lora_diff_names)
                 weight = weight.add(lora_diff, alpha=lora_strength)
         return weight
+
+    def forward(self, input):
+        if self.bias is not None:
+            bias = self.bias.to(input)
+        else:
+            bias = None
+        weight = self.weight.to(input)
+
+        if self.scale_weight is not None:
+            if weight.numel() < input.numel():
+                weight = weight * self.scale_weight
+            else:
+                input = input * self.scale_weight
+
+        weight = self._get_weight_with_lora(weight)
+        
+        if self.compute_dtype is not None:
+            weight = weight.to(self.compute_dtype)
+
+        return torch.nn.functional.linear(input, weight, bias)
     
 def remove_lora_from_module(module):
     for name, submodule in module.named_modules():