[Bugfix] Fix JambaForCausalLM LoRA (vllm-project#14370)

Signed-off-by: Jee Jee Li <[email protected]>

Author: jeejeelee

tests/lora/conftest.py

@@ -6,7 +6,6 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-import safetensors
 import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
@@ -191,29 +190,6 @@ def mixtral_lora_files_all_target_modules():
     return snapshot_download(repo_id="dyang415/mixtral-lora-v0")
 
 
-@pytest.fixture(scope="session")
-def jamba_lora_files():
-    #   some of the adapters have unnecessary weights for serving,
-    #   hence we remove them
-    def remove_unnecessary_weights(path):
-        lora_path = f"{adapter_path}/adapter_model.safetensors"
-        tensors = safetensors.torch.load_file(lora_path)
-        nonlora_keys = []
-        for k in list(tensors.keys()):
-            if "lora" not in k:
-                nonlora_keys.append(k)
-        for k in nonlora_keys:
-            del tensors[k]
-        safetensors.torch.save_file(tensors, lora_path)
-
-    adapter_path = snapshot_download(
-        repo_id=
-        "hf-100/Jamba-1.5-mini-Spellbound-StoryWriter-0.1-6583896-ckpt53-lora")
-
-    remove_unnecessary_weights(adapter_path)
-    return adapter_path
-
-
 @pytest.fixture(scope="session")
 def gemma_lora_files():
     return snapshot_download(repo_id="wskwon/gemma-7b-test-lora")

tests/lora/test_jamba.py

@@ -632,6 +632,7 @@ def create_random_linear_replicated_layer():
 
         id_to_index = get_random_id_to_index(num_loras, max_loras)
         linear, lora_linear = create_random_linear_replicated_layer()
+        assert torch.equal(linear.weight, lora_linear.weight)
         lora_linear.set_mapping(punica_wrapper)
         lora_dict, _ = populate_loras(
             id_to_index,
@@ -757,6 +758,7 @@ def create_random_linear_parallel_layer():
 
         id_to_index = get_random_id_to_index(num_loras, max_loras)
         linear, lora_linear = create_random_linear_parallel_layer()
+        assert torch.equal(linear.weight, lora_linear.weight)
         lora_linear.set_mapping(punica_wrapper)
         lora_dict, _ = populate_loras(
             id_to_index,
@@ -904,6 +906,7 @@ class FakeConfig:
         id_to_index = get_random_id_to_index(num_loras, max_loras)
 
         linear, lora_linear = create_column_parallel_packed_layer()
+        assert torch.equal(linear.weight, lora_linear.weight)
         lora_linear.set_mapping(punica_wrapper)
         lora_dict, sublora_dict = populate_loras(
             id_to_index,

tests/lora/test_layers.py

@@ -274,6 +274,10 @@ def can_replace_layer(
     ) -> bool:
         return type(source_layer) is VocabParallelEmbedding
 
+    @property
+    def weight(self):
+        return self.base_layer.weight
+
 
 class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
 
@@ -409,6 +413,34 @@ def apply(self,
                                             self.output_slices)
         return output
 
+    @property
+    def weight(self) -> torch.Tensor:
+
+        # unquantizedLinear
+        if hasattr(self.base_layer, "weight"):
+            return self.base_layer.weight
+        # Compressed Tensor
+        elif hasattr(self.base_layer, "weight_packed"):
+            return self.base_layer.weight_packed
+        # GPTQ/AWQ
+        elif hasattr(self.base_layer, "qweight"):
+            return self.base_layer.qweight
+        # marlin
+        elif hasattr(self.base_layer, "B"):
+            return self.base_layer.B
+        # HQQ marlin
+        elif hasattr(self.base_layer, "W_q"):
+            return self.base_layer.W_q
+        else:
+            raise ValueError(f"Unsupported base layer: {self.base_layer}")
+
+    @property
+    def bias(self) -> Optional[torch.Tensor]:
+        if hasattr(self.base_layer, "bias"):
+            return self.base_layer.bias
+        else:
+            return None
+
 
 class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):
 
@@ -902,11 +934,6 @@ def forward(
 
         return output, output_bias
 
-    @property
-    def weight(self):
-        return (self.base_layer.weight if hasattr(self.base_layer, "weight")
-                else self.base_layer.qweight)
-
     @classmethod
     @_not_fully_sharded_can_replace
     def can_replace_layer(