Merge pull request #272 from lakshith-403/LoRA

LoRA minor fixes

Author: vpj

labml_nn/lora/__init__.py

@@ -79,7 +79,7 @@ def __init__(self, in_features: int, out_features: int, bias: bool,
         # Matrix $A \in \mathbb{R}^{r \times k}$
         self.lora_a = nn.Parameter(torch.empty((r, in_features)))
         # Matrix $B \in \mathbb{R}^{d \times r}$, we keep $A$ and $B$ transposed
-        self.lora_b = nn.Parameter(torch.empty((outfeatures, r)))
+        self.lora_b = nn.Parameter(torch.empty((out_features, r)))
 
         with torch.no_grad():
             # Initialize $A$ similar to a weight matrix in a normal linear layer

labml_nn/lora/experiment.py

@@ -81,14 +81,14 @@ def _load_pretrained_weights(self):
 
         # Mapping (`hf: ours`) of decoder layers
         for i in range(12):
-            mapping[f'transformer.h.{i}.ln_1.weight'] = f'blocks.{i}.pre_norm.weight'
-            mapping[f'transformer.h.{i}.ln_1.bias'] = f'blocks.{i}.pre_norm.bias'
+            mapping[f'transformer.h.{i}.ln_1.weight'] = f'blocks.{i}.attn_norm.weight'
+            mapping[f'transformer.h.{i}.ln_1.bias'] = f'blocks.{i}.attn_norm.bias'
             mapping[f'transformer.h.{i}.attn.c_attn.weight'] = f'blocks.{i}.attn.qkv_projection.weight'
             mapping[f'transformer.h.{i}.attn.c_attn.bias'] = f'blocks.{i}.attn.qkv_projection.bias'
             mapping[f'transformer.h.{i}.attn.c_proj.weight'] = f'blocks.{i}.attn.output_projection.weight'
             mapping[f'transformer.h.{i}.attn.c_proj.bias'] = f'blocks.{i}.attn.output_projection.bias'
-            mapping[f'transformer.h.{i}.ln_2.weight'] = f'blocks.{i}.post_norm.weight'
-            mapping[f'transformer.h.{i}.ln_2.bias'] = f'blocks.{i}.post_norm.bias'
+            mapping[f'transformer.h.{i}.ln_2.weight'] = f'blocks.{i}.ffn_norm.weight'
+            mapping[f'transformer.h.{i}.ln_2.bias'] = f'blocks.{i}.ffn_norm.bias'
             mapping[f'transformer.h.{i}.mlp.c_fc.weight'] = f'blocks.{i}.ffn.linear_in.weight'
             mapping[f'transformer.h.{i}.mlp.c_fc.bias'] = f'blocks.{i}.ffn.linear_in.bias'
             mapping[f'transformer.h.{i}.mlp.c_proj.weight'] = f'blocks.{i}.ffn.linear_out.weight'
@@ -110,7 +110,11 @@ def _load_pretrained_weights(self):
             new_state_dict[layer] = torch.transpose(new_state_dict[layer], 0, 1)
 
         # Load out model. We use `strict = False` because the state does not have LoRA weights
-        self.model.load_state_dict(new_state_dict, strict=False)
+        missing_keys, unexpected_keys = self.model.load_state_dict(new_state_dict, strict=False)
+
+        # make sure that only lora weights are not loaded
+        assert all('lora' in key for key in missing_keys)
+        assert not unexpected_keys
 
     def initialize(self):
         """