chore: import upstream snapshot with attribution

2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
@@ -0,0 +1,151 @@
+"""
+---
+title: Low-Rank Adaptation (LoRA)
+summary: >
+  Annotated implementation of RoRA from paper
+  LoRA: Low-Rank Adaptation of Large Language Models
+---
+
+# Low-Rank Adaptation (LoRA)
+
+This is an implementation of
+[Low-Rank Adaptation (LoRA)](https://arxiv.org/abs/2106.09685)
+in [PyTorch](https://pytorch.org).
+
+Low-Rank Adaptation (LoRA) freezes pre-trained model weights and injects
+ trainable rank decomposition matrices into each layer of the transformer.
+ This makes it possible to efficiently fine-tune large language models by
+ reducing trainable parameters by a large factor.
+
+Here's [the training code](experiment.html) for training a GPT2 model with LoRA
+ on Tiny Shakespeare dataset.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class Linear(nn.Module):
+    """
+    ## LoRA Linear Layer
+
+    LoRA linear layer adds a low-rank decomposition to the pre-trained
+    weight matrix ($W_0 \in \mathbb{R}^{d \times k}$)
+    of the linear layer.
+
+    $$W_0 + \Delta W = W_0 + BA$$
+
+    , where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$,
+     and the rank $r \ll min(d, k)$.
+
+    All parameters are frozen except $A$ and $B$.
+
+    $\Delta W$ is initialized to be zero at the beginning of the training.
+
+    They multiple $x \Delta W^T$ by $\frac{\alpha}{r}$ where $\alpha$ is a hyper-parameter.
+    Once $\alpha$ is tuned it can be kept the same when varying $r$.
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool,
+                 r: int, alpha: int = None):
+        """
+        :param in_features: is the number of input features of the linear layer
+        :param out_features: is the number of output features of the linear layer
+        :param bias: is a flag indicating if there is a bias parameter
+        :param r: is the rank of the decomposition $r$
+        :param alpha: is the scaling factor $\alpha$
+        """
+        super().__init__()
+
+        # Set $\alpha = r$ is not provided. i.e. make the scaling factor $\frac{\alpha}{r} = 1$.
+        if alpha is None:
+            alpha = r
+
+        # The pre-trained weight $W_0$
+        self.weight = nn.Parameter(torch.empty((out_features, in_features)))
+        # Freeze it
+        self.weight.requires_grad = False
+
+        if bias:
+            # Bias parameter $b_0$ (also frozen)
+            self.bias = nn.Parameter(torch.empty(out_features))
+            self.bias.requires_grad = False
+        else:
+            # No bias parameter
+            self.bias = None
+
+        # scaling factor $\frac{\alpha}{r}$
+        self.scaling = alpha / r
+        # 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((out_features, r)))
+
+        with torch.no_grad():
+            # Initialize $A$ similar to a weight matrix in a normal linear layer
+            nn.init.kaiming_uniform_(self.lora_a, a=5 ** 0.5)
+            # Initialize $B$ to $0$ so that $\Delta W = BA$ is $0$ at initialization
+            nn.init.zeros_(self.lora_b)
+
+    def forward(self, x: torch.Tensor):
+        # Compute $x W_0^T + b_0$
+        result = nn.functional.linear(x, self.weight, bias=self.bias)
+
+        # Add $\frac{\alpha}{r} x \Delta W^T = \frac{\alpha}{r} x {(BA)}^T = \frac{\alpha}{r} x A^T B^T$
+        result += (x @ self.lora_a.T @ self.lora_b.T) * self.scaling
+
+        #
+        return result
+
+
+class Embedding(nn.Module):
+    """
+    ## LoRA Embedding Layer
+
+    Similar to LoRA linear layer this adds a low-rank decomposition to the pre-trained
+    embedding weights matrix ($W_0 \in \mathbb{R}^{d \times k}$).
+
+    $$W_0 + \Delta W = W_0 + BA$$
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int,
+                 r: int, alpha: int = None):
+        """
+
+        :param num_embeddings: is the number of embeddings
+        :param embedding_dim: is the number embedding dimensions
+        :param r: is the rank of the decomposition $r$
+        :param alpha: is the scaling factor $\alpha$
+        """
+        super().__init__()
+
+        # Set $\alpha = r$ is not provided. i.e. make the scaling factor $\frac{\alpha}{r} = 1$.
+        if alpha is None:
+            alpha = r
+
+        # The pre-trained embedding weights $W_0^T$ (frozen)
+        self.weight = nn.Parameter(torch.empty((num_embeddings, embedding_dim)))
+        self.weight.requires_grad = False
+
+        # scaling factor $\frac{\alpha}{r}$
+        self.scaling = alpha / r
+        # Matrix $A \in \mathbb{R}^{r \times k}$
+        self.lora_a = nn.Parameter(torch.empty((r, num_embeddings)))
+        # Matrix $B \in \mathbb{R}^{d \times r}$
+        self.lora_b = nn.Parameter(torch.empty((embedding_dim, r)))
+
+        with torch.no_grad():
+            # Initialize $A$ with a normal distribution
+            nn.init.normal_(self.lora_a)
+            # Initialize $B$ to $0$ so that $\Delta W = BA$ is $0$ at initialization
+            nn.init.zeros_(self.lora_b)
+
+    def forward(self, x: torch.Tensor):
+        # Compute the embeddings $\text{onehot}(x) W_0$
+        result = nn.functional.embedding(x, self.weight)
+
+        # Add $\frac{\alpha}{r} \text{onehot}(x) \Delta W^T = \frac{\alpha}{r} \text{onehot}(x) A^T B^T$
+        result += (nn.functional.embedding(x, self.lora_a.T) @ self.lora_b.T) * self.scaling
+
+        #
+        return result
@@ -0,0 +1,97 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "!pip install labml-nn",
+   "id": "c5ed37230628ee76"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "from labml_nn.lora.experiment import Trainer\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": "trainer = Trainer()",
+   "id": "31c9bc08eca2592",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "experiment.configs(trainer)",
+   "id": "fb6ce74326558948",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "trainer.initialize()",
+   "id": "1456cfab47dee3b",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "with experiment.start():\n",
+    "    trainer.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
+}
@@ -0,0 +1,189 @@
+"""
+---
+title: Finetune GPT-2 with LoRA
+summary: This is training code with notes for fine-tuning pre-trained GPT-2 model with LoRA.
+---
+
+# Finetune [GPT-2](gpt2.html) with [LoRA](index.html)
+
+Here's a Colab notebook for training a feedback transformer on Tiny Shakespeare dataset.
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/lora/experiment.ipynb)
+"""
+
+import torch
+from torch.optim import Adam
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from labml import lab, monit, tracker
+from labml.configs import BaseConfigs, option
+from labml.utils.download import download_file
+from labml_nn.helpers.device import DeviceConfigs
+from labml_nn.lora.gpt2 import GPTModel
+
+
+class Trainer(BaseConfigs):
+    """
+    ## Trainer configurations and the training loop
+
+    The default configs can and will be over-ridden when we start the experiment
+    """
+    device: torch.device = DeviceConfigs()
+
+    # GPT-2 configs
+    layer_norm_epsilon: float = 1e-05
+    d_model: int = 768
+    n_layers: int = 12
+    n_heads: int = 12
+    n_positions: int = 1024
+    vocab_size: int = 50257
+
+    # Training configs
+    epochs: int = 10
+    batch_size: int = 32
+    learning_rate: float = 1e-4
+    context_len: int = 512
+
+    # LoRA rank
+    lora_r: int = 32
+
+    # Dataset
+    text: TensorDataset = "tiny_shakespeare"
+    # Huggingface tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    # [GPT2 model](gpt2.html)
+    model: GPTModel
+    # Optimizer
+    optimizer: torch.optim.Adam
+    # Cross entropy loss
+    loss_func = torch.nn.CrossEntropyLoss()
+    # Dataloader
+    data_loader: DataLoader
+
+    def _load_pretrained_weights(self):
+        """
+        ### Load pre-trained [GPT-2 from huggingface](https://huggingface.co/openai-community/gpt2)
+        """
+
+        # Load the huggingface model and get the parameters
+        hf_model = AutoModelForCausalLM.from_pretrained("gpt2")
+        state_dict = hf_model.state_dict()
+
+        # Transformer embedding and prediction layer parameter mapping (`hf: ours`)
+        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'
+        }
+
+        # Mapping (`hf: ours`) of decoder layers
+        for i in range(12):
+            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}.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'
+            mapping[f'transformer.h.{i}.mlp.c_proj.bias'] = f'blocks.{i}.ffn.linear_out.bias'
+
+        # Move the parameters based on mapping
+        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]
+
+        # GPT-2 hugging face uses 1D Convolution layers. We need to transpose those weights since we use linear layers
+        convo_layers = ([f'blocks.{i}.ffn.linear_in.weight' for i in range(12)] +
+                        [f'blocks.{i}.ffn.linear_out.weight' for i in range(12)] +
+                        [f'blocks.{i}.attn.qkv_projection.weight' for i in range(12)] +
+                        [f'blocks.{i}.attn.output_projection.weight' for i in range(12)])
+
+        for layer in convo_layers:
+            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
+        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):
+        """
+        ### Initialize the model, optimizer and dataloader
+        """
+        # Initialize the [GPT2 model](gpt2.html)
+        self.model = GPTModel(
+            layer_norm_epsilon=self.layer_norm_epsilon,
+            d_model=self.d_model,
+            n_layers=self.n_layers,
+            n_heads=self.n_heads,
+            n_positions=self.n_positions,
+            vocab_size=self.vocab_size,
+            r=self.lora_r,
+        )
+        self.model.to(self.device)
+        # Load pre-trained model weights
+        self._load_pretrained_weights()
+
+        # Initialize the optimizer
+        self.optimizer = Adam(self.model.parameters(), lr=self.learning_rate)
+
+        # Initialize the data loader
+        self.data_loader = DataLoader(self.text, batch_size=self.batch_size, shuffle=True)
+
+    def run(self):
+        """
+        ### Training loop
+        """
+
+        for _ in monit.loop(self.epochs):
+            # `inputs` has shape `[batch_size, seq_len]`
+            for (inputs,) in monit.iterate('Train', self.data_loader):
+                # Move `inputs` to device
+                inputs = inputs.to(self.device)
+                # Call the model, with the all but the last token
+                logits = self.model(inputs[:, :-1])
+                # Get cross entropy loss
+                loss = self.loss_func(logits.reshape(-1, logits.shape[-1]), inputs[:, 1:].reshape(-1))
+
+                # Make gradients 0
+                self.optimizer.zero_grad()
+                # Compute gradients
+                loss.backward()
+                # Optimize
+                self.optimizer.step()
+
+                # Log the loss
+                tracker.save({'loss': loss})
+                tracker.add_global_step()
+            #
+            tracker.new_line()
+
+
+@option(Trainer.text)
+def tiny_shakespeare(c: Trainer):
+    """
+    ### 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)
@@ -0,0 +1,192 @@
+"""
+---
+title: GPT-2 with LoRA
+summary: GPT-2 implementation with LoRA modules
+---
+
+# GPT-2 with [LoRA modules](index.html)
+
+Here's [the training code](experiment.html) for training a GPT2 model with LoRA
+ on Tiny Shakespeare dataset.
+"""
+
+import torch
+import torch.nn as nn
+
+from labml_nn.lora import Linear, Embedding
+
+
+class FFN(nn.Module):
+    """
+    ### Feedforward Network
+    """
+
+    def __init__(self, d_model: int, d_ff: int, r: int):
+        """
+        :param d_model: is the number of dimensions
+        :param d_ff: is the size of the hidden dimension
+        :param r: is the lora rank
+        """
+        super().__init__()
+
+        # The linear layers and the activation
+        self.linear_in = Linear(d_model, d_ff, r=r, bias=True)
+        self.linear_out = Linear(d_ff, d_model, r=r, bias=True)
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        :param x: is the embeddings tensor with shape `[batch_size, seq_len, d_model]`
+        """
+        x = self.linear_in(x)
+        x = self.act(x)
+        x = self.linear_out(x)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    ### Multi-Head Attention
+    """
+
+    def __init__(self, d_model: int, n_heads: int, r: int):
+        """
+        :param d_model: is the number of dimensions in the embeddings
+        :param n_heads: is the number of heads
+        :param r: is the lora rank
+        """
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+
+        # Linear transformation for QKV
+        self.qkv_projection = Linear(d_model, d_model * 3, r=r, bias=True)
+        # Output projection
+        self.output_projection = Linear(d_model, d_model, r=r, bias=True)
+
+    def _split_heads(self, x: torch.Tensor):
+        """
+        :param x: is the tensor with shape `[batch_size, seq_len, d_model]`
+        """
+        # Split last dimension to `[n_heads, d_head]`
+        x = x.view(x.shape[:-1] + (self.n_heads, self.d_head))
+        # Reorder to `[batch_size, head, seq_length, d_head]`
+        return x.permute(0, 2, 1, 3)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        :param x: is the embeddings tensor with shape `[batch_size, seq_len, d_model]`
+        """
+        batch_size, seq_length, _ = x.shape
+
+        # Get query, key and value
+        q, k, v = self.qkv_projection(x).split(self.d_model, dim=-1)
+
+        # Transform them from shape  `[batch_size, seq_len, d_model]` to `[batch_size, head, seq_length, d_head]`
+        q = self._split_heads(q)
+        k = self._split_heads(k)
+        v = self._split_heads(v)
+
+        # Apply causal attention
+        attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
+
+        # Transform them from shape  `[batch_size, head, seq_length, d_head]` to `[batch_size, seq_len, d_model]`
+        attn_output = attn_output.permute(0, 2, 1, 3).reshape(batch_size, seq_length, self.d_model)
+
+        # Final project
+        return self.output_projection(attn_output)
+
+
+class Block(nn.Module):
+    """
+    ### Decoder block
+    """
+
+    def __init__(self, d_model: int, n_heads: int, layer_norm_epsilon: float, r: int):
+        """
+        :param d_model: is the number of dimensions in the embeddings
+        :param n_heads: is the number of heads
+        :param layer_norm_epsilon: is the layer norm epsilon
+        :param r: is the lora rank
+        """
+        super().__init__()
+        # Attention pre-normalization layer
+        self.attn_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
+        # Attention layer
+        self.attn = MultiHeadAttention(d_model, n_heads, r)
+        # FFN pre-normalization layer
+        self.ffn_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
+        # Feed-forward network
+        self.ffn = FFN(d_model, d_model * 4, r)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        :param x: is the embeddings tensor with shape `[batch_size, seq_len, d_model]`
+        """
+        # Attention
+        x = x + self.attn(self.attn_norm(x))
+        # FFN
+        x = x + self.ffn(self.ffn_norm(x))
+
+        return x
+
+
+class GPTModel(nn.Module):
+    """
+    ## GPT2 Model
+    """
+
+    def __init__(self, *, d_model: int,
+                 n_heads: int, n_layers: int,
+                 n_positions: int,
+                 layer_norm_epsilon: float,
+                 vocab_size: int, r: int):
+        """
+        :param d_model: is the number of dimensions in the embeddings
+        :param n_heads: is the number of attention heads
+        :param n_layers: is the number of decoder layers
+        :param n_positions: is the number of positional embeddings
+        :param layer_norm_epsilon: is the layer norm epsilon
+        :param vocab_size: is the vocabulary size
+        :param r: is the lora rank
+        """
+        super().__init__()
+
+        # Token and absolute positional embeddings
+        self.token_embedding = Embedding(vocab_size, d_model, r=r)
+        self.position_embedding = Embedding(n_positions, d_model, r=r)
+
+        # Decoder blocks
+        self.blocks = nn.ModuleList([Block(d_model, n_heads, layer_norm_epsilon, r=r)
+                                     for _ in range(n_layers)])
+
+        # Final layer norm
+        self.final_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
+        # Projection layer to logit space
+        self.lm_head = Linear(d_model, vocab_size, r=r, bias=False)
+
+    def forward(self, input_ids: torch.Tensor):
+        """
+        :param input_ids: has shape `[batch_size, seq_len]`
+        """
+        batch_size, seq_len = input_ids.shape
+
+        # Get token embeddings
+        token_embeddings = self.token_embedding(input_ids)
+        # Get position ids
+        position_ids = torch.arange(seq_len, device=input_ids.device)[None, :]
+        # Get position embeddings
+        position_embeddings = self.position_embedding(position_ids)
+
+        # Add position embeddings
+        x = token_embeddings + position_embeddings
+
+        # Run through transformer blocks
+        for block in self.blocks:
+            x = block(x)
+
+        # Final normalization
+        x = self.final_norm(x)
+        # Get logits from projection layer
+        return self.lm_head(x)