return singluar values

logix-project · Nov 20, 2023 · 6d5e888 · 6d5e888
1 parent 0c6d13d
commit 6d5e888
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/analog/lora/modules.py b/analog/lora/modules.py
@@ -48,11 +48,15 @@ def init_weight(self, init_strategy: str = "random", hessian=None):
             nn.init.kaiming_uniform_(self.analog_lora_A.weight, a=math.sqrt(5))
             nn.init.kaiming_uniform_(self.analog_lora_C.weight, a=math.sqrt(5))
         elif init_strategy == "pca":
-            top_r_singular_vector_forward = compute_top_k_singular_vectors(
-                hessian[FORWARD], self.rank
-            )
-            top_r_singular_vector_backward = compute_top_k_singular_vectors(
-                hessian[BACKWARD], self.rank
-            )
+            (
+                top_r_singular_vector_forward,
+                top_r_singular_value_forward,
+            ) = compute_top_k_singular_vectors(hessian[FORWARD], self.rank)
+            (
+                top_r_singular_vector_backward,
+                top_r_singular_value_backward,
+            ) = compute_top_k_singular_vectors(hessian[BACKWARD], self.rank)
+            # top_r_singular_vector_forward /= top_r_singular_value_forward.unsqueeze(0)
+            # top_r_singular_vector_backward /= top_r_singular_value_backward.unsqueeze(0)
             self.analog_lora_A.weight.data.copy_(top_r_singular_vector_forward.T)
             self.analog_lora_C.weight.data.copy_(top_r_singular_vector_backward)
diff --git a/analog/lora/utils.py b/analog/lora/utils.py
@@ -7,4 +7,5 @@ def compute_top_k_singular_vectors(matrix, k):
     """
     U, S, Vh = torch.linalg.svd(matrix)
     top_k_singular_vectors = U[:, :k]
-    return top_k_singular_vectors
+    top_k_singular_values = S[:k]
+    return top_k_singular_vectors, top_k_singular_values