Merge pull request #268 from lakshith-403/LoRA

LoRA experiment

labml_nn/lora/experiment.ipynb

@@ -0,0 +1,89 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "from labml_nn.lora.experiment import Configs\n",
+    "from labml import experiment"
+   ],
+   "id": "1b9da2e59ffce5d5",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true
+   },
+   "source": "experiment.create(name=\"lora_gpt2\")",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "conf = Configs()",
+   "id": "31c9bc08eca2592",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "experiment.configs(conf)",
+   "id": "fb6ce74326558948",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "conf.initialize()",
+   "id": "1456cfab47dee3b",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "with experiment.start():\n",
+    "    conf.run()"
+   ],
+   "id": "3fe4068fd2df9094",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "",
+   "id": "d3c3c723ebbe854a",
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ml)",
+   "language": "python",
+   "name": "ml"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

labml_nn/lora/experiment.py

@@ -0,0 +1,137 @@
+import torch
+from labml import lab, monit, tracker
+from labml.configs import BaseConfigs, option
+from labml.utils.download import download_file
+from labml_helpers.device import DeviceConfigs
+from torch.optim import Adam
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from labml_nn.lora.gpt2 import GPTModel
+
+
+class Configs(BaseConfigs):
+    device: torch.device = DeviceConfigs()
+    layer_norm_epsilon: float = 1e-05
+    n_embed: int = 768
+    n_layer: int = 12
+    n_positions: int = 1024
+    vocab_size: int = 50257
+    epochs: int = 10
+    batch_size: int = 32
+    learning_rate: float = 1e-4
+    context_len: int = 512
+    r: int = 32
+
+    text: TensorDataset = "tiny_shakespeare"
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    model: GPTModel
+    optimizer: torch.optim.Adam
+    criterion = torch.nn.CrossEntropyLoss()
+    data_loader: DataLoader
+
+    def _load_pretrained_weights(self):
+        hf_model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+        state_dict = hf_model.state_dict()
+
+        mapping = {
+            'transformer.wte.weight': 'token_embedding.weight',
+            'transformer.wpe.weight': 'position_embedding.weight',
+            'transformer.ln_f.weight': 'final_norm.weight',
+            'transformer.ln_f.bias': 'final_norm.bias',
+            'lm_head.weight': 'lm_head.weight'
+        }
+
+        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}.attn.c_attn.weight'] = f'blocks.{i}.attn.c_att.weight'
+            mapping[f'transformer.h.{i}.attn.c_attn.bias'] = f'blocks.{i}.attn.c_att.bias'
+            mapping[f'transformer.h.{i}.attn.c_proj.weight'] = f'blocks.{i}.attn.c_proj.weight'
+            mapping[f'transformer.h.{i}.attn.c_proj.bias'] = f'blocks.{i}.attn.c_proj.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}.mlp.c_fc.weight'] = f'blocks.{i}.ffn.c_fc.weight'
+            mapping[f'transformer.h.{i}.mlp.c_fc.bias'] = f'blocks.{i}.ffn.c_fc.bias'
+            mapping[f'transformer.h.{i}.mlp.c_proj.weight'] = f'blocks.{i}.ffn.c_proj.weight'
+            mapping[f'transformer.h.{i}.mlp.c_proj.bias'] = f'blocks.{i}.ffn.c_proj.bias'
+
+        new_state_dict = {}
+        for old_key, new_key in mapping.items():
+            if old_key in state_dict:
+                new_state_dict[new_key] = state_dict[old_key]
+
+        # transpose weight matrices of convo 1d layers to use linear layers instead
+        convo_layers = ([f'blocks.{i}.ffn.c_fc.weight' for i in range(12)] +
+                        [f'blocks.{i}.ffn.c_proj.weight' for i in range(12)] +
+                        [f'blocks.{i}.attn.c_att.weight' for i in range(12)] +
+                        [f'blocks.{i}.attn.c_proj.weight' for i in range(12)])
+
+        for layer in convo_layers:
+            new_state_dict[layer] = torch.transpose(new_state_dict[layer], 0, 1)
+
+        self.model.load_state_dict(new_state_dict, strict=False)  # state dict does not have lora weights
+
+        del hf_model
+        del state_dict
+        del new_state_dict
+
+    def initialize(self):
+        self.model = GPTModel(
+            layer_norm_epsilon=self.layer_norm_epsilon,
+            n_embd=self.n_embed,
+            n_layer=self.n_layer,
+            n_positions=self.n_positions,
+            vocab_size=self.vocab_size,
+            r=self.r,
+            device=self.device
+        ).to(self.device)
+        self._load_pretrained_weights()
+
+        self.optimizer = Adam(self.model.parameters(), lr=self.learning_rate)
+
+        self.data_loader = DataLoader(self.text, batch_size=self.batch_size, shuffle=True)
+
+    def run(self):
+        for _ in monit.loop(self.epochs):
+            for i, batch in monit.enum('Train', self.data_loader):
+                inputs = batch[0]
+                inputs = inputs.to(self.device)
+                labels = inputs.clone()
+
+                outputs = self.model(inputs)
+
+                shift_logits = outputs[..., :-1, :]
+                shift_labels = labels[..., 1:]
+
+                loss = self.criterion(shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1))
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                self.optimizer.step()
+
+                tracker.add({'loss': loss})
+
+                tracker.save()
+                tracker.add_global_step()
+            tracker.new_line()
+
+
+@option(Configs.text)
+def tiny_shakespeare(c: Configs):
+    """
+    ### Tiny Shakespeare dataset
+
+    It will download from the url if not present
+    """
+    path = lab.get_data_path() / 'tiny_shakespeare.txt'
+    if not path.exists():
+        download_file("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt", path)
+    with open(path, 'r', encoding='utf-8') as f:
+        text = f.read()
+
+    tokens = c.tokenizer.encode(text)
+    num_batches = len(tokens) // (c.batch_size * c.context_len)
+    tokens = tokens[:num_batches * c.batch_size * c.context_len]
+    input_ids = torch.tensor(tokens).view(-1, c.context_len)
+    return TensorDataset(input_ids)

labml_nn/lora/gpt2.py

@@ -1,26 +1,13 @@
 import torch
 import torch.nn as nn
-from transformers import AutoTokenizer
 from labml_nn.lora import Linear, Embedding
 
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-
-config = {
-    "layer_norm_epsilon": 1e-05,
-    "n_embd": 768,
-    "n_head": 12,
-    "n_layer": 12,
-    "n_positions": 1024,
-    "vocab_size": 50257,
-    "device": "cuda"
-}
-
 
 class FFN(nn.Module):
-    def __init__(self, dim):
+    def __init__(self, dim: int, n_embed: int, r: int):
         super().__init__()
-        self.c_fc = Linear(config['n_embd'], dim, r=32, bias=True)
-        self.c_proj = Linear(dim, config['n_embd'], r=32, bias=True)
+        self.c_fc = Linear(n_embed, dim, r=r, bias=True)
+        self.c_proj = Linear(dim, n_embed, r=r, bias=True)
         self.act = nn.functional.gelu
 
     def forward(self, hidden_states):
@@ -31,15 +18,15 @@ def forward(self, hidden_states):
 
 
 class MultiHeadAttention(nn.Module):
-    def __init__(self):
+    def __init__(self, n_embed: int, r: int):
         super().__init__()
-        self.embed_dim = config['n_embd']
-        self.num_heads = config['n_head']
+        self.embed_dim = n_embed
+        self.num_heads = n_embed
         self.head_dim = self.embed_dim // self.num_heads
         self.split_size = self.embed_dim
 
-        self.c_att = Linear(config['n_embd'], config['n_embd'] * 3, r=32, bias=True)
-        self.c_proj = Linear(config['n_embd'], config['n_embd'], r=32, bias=True)
+        self.c_att = Linear(n_embed, n_embed * 3, r=r, bias=True)
+        self.c_proj = Linear(n_embed, n_embed, r=r, bias=True)
 
     def _split_heads(self, tensor, num_heads, attn_head_size):
         """
@@ -76,12 +63,12 @@ def forward(self, hidden_states):
 
 
 class Block(nn.Module):
-    def __init__(self):
+    def __init__(self, n_embed: int, layer_norm_epsilon: float, r: int):
         super().__init__()
-        self.pre_norm = nn.LayerNorm(config['n_embd'], eps=config['layer_norm_epsilon'])
-        self.attn = MultiHeadAttention()
-        self.post_norm = nn.LayerNorm(config['n_embd'], eps=config['layer_norm_epsilon'])
-        self.ffn = FFN(config['n_embd'] * 4)
+        self.pre_norm = nn.LayerNorm(n_embed, eps=layer_norm_epsilon)
+        self.attn = MultiHeadAttention(n_embed, r)
+        self.post_norm = nn.LayerNorm(n_embed, eps=layer_norm_epsilon)
+        self.ffn = FFN(n_embed * 4, n_embed, r)
 
     def forward(self, hidden_states):
         residual = hidden_states
@@ -99,23 +86,27 @@ def forward(self, hidden_states):
 
 
 class GPTModel(nn.Module):
-    def __init__(self):
+    def __init__(self, layer_norm_epsilon: float, n_embd: int, n_layer: int, n_positions: int,
+                 vocab_size: int, r: int, device: torch.device):
         super().__init__()
 
-        self.token_embedding = Embedding(config['vocab_size'], config['n_embd'], r=32)
-        self.position_embedding = Embedding(config['n_positions'], config['n_embd'], r=32)
+        self.token_embedding = Embedding(vocab_size, n_embd, r=r)
+        self.position_embedding = Embedding(n_positions, n_embd, r=r)
+
+        self.blocks = nn.ModuleList([Block(n_embd, layer_norm_epsilon, r=r)
+                                     for _ in range(n_layer)])
 
-        self.blocks = nn.ModuleList([Block() for _ in range(config['n_layer'])])
+        self.final_norm = nn.LayerNorm(n_embd, eps=layer_norm_epsilon)
 
-        self.final_norm = nn.LayerNorm(config['n_embd'], eps=config['layer_norm_epsilon'])
+        self.lm_head = Linear(n_embd, vocab_size, r=r, bias=False)
 
-        self.lm_head = Linear(config['n_embd'], config['vocab_size'], r=32, bias=False)
+        self.device = device
 
     def forward(self, input_ids):
         batch_size, input_shape = input_ids.size()
 
         token_embeddings = self.token_embedding(input_ids)  # B T C
-        position_ids = torch.arange(input_shape, device=config['device'])  # T C
+        position_ids = torch.arange(input_shape, device=self.device)  # T C
         position_embeddings = self.position_embedding(position_ids)  # B T C
 
         hidden_states = token_embeddings + position_embeddings