Fix test

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Author: jeejeelee

tests/lora/test_layers.py

@@ -164,8 +164,8 @@ def populate_loras(
                         weight=layer_weights,
                         generate_embeddings_tensor=generate_embeddings_tensor,
                     )
-                sublora.lora_b = sublora.lora_b[:, (sublora_len *
-                                                    i):(sublora_len * (i + 1))]
+                sublora.lora_b = sublora.lora_b[(sublora_len *
+                                                 i):(sublora_len * (i + 1)), :]
                 sublora.optimize()
                 subloras.append(sublora)
 
@@ -304,9 +304,9 @@ def create_random_embedding_layer():
             result = embedding(input_)
             after_a = F.embedding(
                 input_,
-                lora.lora_a,
+                lora.lora_a.T,
             )
-            result += (after_a @ lora.lora_b)
+            result += (after_a @ lora.lora_b.T)
             expected_results.append(result)
         expected_result = torch.cat(expected_results)
 
@@ -445,9 +445,9 @@ def create_random_embedding_layer():
             result = expanded_embedding(input_)
             after_a = F.embedding(
                 original_input_,
-                lora.lora_a,
+                lora.lora_a.T,
             )
-            result += (after_a @ lora.lora_b)
+            result += (after_a @ lora.lora_b.T)
             expected_results.append(result)
         expected_result = torch.cat(expected_results)
 
@@ -575,7 +575,7 @@ def _pretest():
                                                   lm_head=linear,
                                                   embedding_bias=None)
             result[:, vocab_size + embeddings_tensor_len:] = float("-inf")
-            result += input_ @ lora.lora_a @ lora.lora_b * lora.scaling
+            result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling
             expected_results.append(result)
         expected_result = torch.cat(expected_results)
         logits_processor.org_vocab_size = vocab_size
@@ -692,9 +692,10 @@ def create_random_linear_replicated_layer():
 
         expected_results: list[torch.Tensor] = []
         for input_, lora_id in zip(inputs, prompt_mapping):
+
             lora = lora_dict[lora_id]
             result = linear(input_)[0]
-            result += input_ @ lora.lora_a @ lora.lora_b * lora.scaling
+            result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling
             expected_results.append(result)
         expected_result = torch.cat(expected_results)
 
@@ -817,7 +818,7 @@ def create_random_linear_parallel_layer():
         for input_, lora_id in zip(inputs, prompt_mapping):
             lora = lora_dict[lora_id]
             result = linear(input_)[0]
-            result += input_ @ lora.lora_a @ lora.lora_b * lora.scaling
+            result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling
             expected_results.append(result)
         expected_result = torch.cat(expected_results)
 
@@ -965,9 +966,10 @@ class FakeConfig:
             result = linear(input_)[0]
             subloras = sublora_dict[lora_id]
             for i, sublora in enumerate(subloras):
-                result[:, sublora.lora_b.shape[1] * i:sublora.lora_b.shape[1] *
-                       (i + 1)] += (input_ @ sublora.lora_a @ sublora.lora_b *
-                                    sublora.scaling)
+                result[:, sublora.lora_b.shape[0] * i:sublora.lora_b.shape[0] *
+                       (i + 1)] += (
+                           input_ @ sublora.lora_a.T @ sublora.lora_b.T *
+                           sublora.scaling)
             expected_results.append(result)
         expected_result = torch.cat(expected_results)
 

tests/lora/utils.py

@@ -36,10 +36,10 @@ def init_random_lora(
             module_name,
             rank=rank,
             lora_alpha=1,
-            lora_a=torch.rand([weight.shape[1], rank],
+            lora_a=torch.rand([rank, weight.shape[1]],
                               dtype=weight.dtype,
                               device=self._device),
-            lora_b=torch.rand([rank, weight.shape[0]],
+            lora_b=torch.rand([weight.shape[0], rank],
                               dtype=weight.dtype,
                               device=self._device),
         )
@@ -67,8 +67,8 @@ def init_lora(
             module_name,
             rank=rank,
             lora_alpha=1,
-            lora_a=torch.rand([input_dim, rank], device="cuda"),
-            lora_b=torch.rand([rank, output_dim], device="cuda"),
+            lora_a=torch.rand([rank, input_dim], device="cuda"),
+            lora_b=torch.rand([output_dim, input_dim], device="cuda"),
             embeddings_tensor=embeddings_tensor,
         )
         self.set_module_lora(module_name, lora)