Fix LoRA embedding parameter initialization

loralib/layers.py

@@ -56,8 +56,8 @@ def reset_parameters(self):
         nn.Embedding.reset_parameters(self)
         if hasattr(self, 'lora_A'):
             # initialize A the same way as the default for nn.Linear and B to zero
-            nn.init.zeros_(self.lora_A)
-            nn.init.normal_(self.lora_B)
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
 
     def train(self, mode: bool = True):
         nn.Embedding.train(self, mode)

tests/test_layers.py

@@ -0,0 +1,25 @@
+import torch
+
+import loralib as lora
+
+
+def test_embedding_lora_parameter_initialization():
+    layer = lora.Embedding(
+        num_embeddings=100,
+        embedding_dim=32,
+        r=8,
+    )
+
+    assert not torch.all(layer.lora_A == 0)
+    assert torch.all(layer.lora_B == 0)
+
+
+def test_linear_lora_parameter_initialization():
+    layer = lora.Linear(
+        in_features=100,
+        out_features=32,
+        r=8,
+    )
+
+    assert not torch.all(layer.lora_A == 0)
+    assert torch.all(layer.lora_B == 0)