{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNTvbmGZHYOR+skFmiV5gvt",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/soumyadip1995/BabyGPT/blob/main/Notebook/lora.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "words = open('/content/ALL_eminem.txt', 'r').read().split()\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "chars = sorted(list(set(words)))\n",
        "string2integer = {ch: i for i, ch in enumerate(chars)}\n",
        "# print(string2integer)\n",
        "\n",
        "integer2string = {i:ch for ch,i in string2integer.items()}\n",
        "encode = lambda s: [string2integer[c] for c in s]\n",
        "# print(encode)\n",
        "\n",
        "decode = lambda l: ''.join([integer2string[i] for i in l])\n",
        "# print(decode)\n",
        "\n",
        "data = torch.tensor(encode(words), dtype = torch.long)\n",
        "# print(data)\n",
        "# data.size()\n",
        "\n",
        "vocab_size = len(chars)"
      ],
      "metadata": {
        "id": "veIfNs0y3Ztw"
      },
      "execution_count": 97,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "from math import sqrt\n",
        "\n",
        "\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "class Attention(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads):\n",
        "    super(Attention, self).__init__()\n",
        "\n",
        "    self.atten = nn.Linear(embedded_dim, 3 * embedded_dim)\n",
        "    self.projection = nn.Linear(embedded_dim, embedded_dim)\n",
        "    self.num_heads = num_heads\n",
        "    self.embedded_dim = embedded_dim\n",
        "    self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    B,T,C = x.size()\n",
        "    q, k ,v  = self.atten(x).split(self.embedded_dim, dim=2)\n",
        "    q = q.view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)\n",
        "    k = k.view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)\n",
        "    v = v.view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)\n",
        "\n",
        "\n",
        "\n",
        "    # manual implementation of attention\n",
        "    # from karpathy\n",
        "    att = (q @ k.transpose(-2, -1)) * ((1.0 / math.sqrt(k.size(-1))))\n",
        "    att = att.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
        "    att = F.softmax(att, dim=-1)\n",
        "    y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
        "\n",
        "    y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
        "\n",
        "    # output projection\n",
        "    y = self.projection(y)\n",
        "    return y"
      ],
      "metadata": {
        "id": "_PF4pa1OUkIQ"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dropout = 0.2\n",
        "class FeedForward(nn.Module):\n",
        "  def __init__(self, embedded_dim):\n",
        "    super(FeedForward, self).__init__()\n",
        "    self.net = nn.Sequential(nn.Linear(embedded_dim, 4 * embedded_dim),\n",
        "    nn.Linear(4 * embedded_dim, embedded_dim),\n",
        "    nn.GELU(),\n",
        "    nn.Dropout(dropout))\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.net(x)"
      ],
      "metadata": {
        "id": "srO-4qIOUlSd"
      },
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "### A simple Transformer Block\n",
        "class Transformer(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads):\n",
        "    super(Transformer, self).__init__()\n",
        "    self.attention = Attention(embedded_dim, num_heads)\n",
        "    self.feed_forward = FeedForward(embedded_dim)\n",
        "    self.layer_norm_1 = nn.LayerNorm(embedded_dim)\n",
        "    self.layer_norm_2 = nn.LayerNorm(embedded_dim)\n",
        "\n",
        "  def forward(self, x):\n",
        "\n",
        "    x = x + self.attention(self.layer_norm_1(x))\n",
        "    x = x + self.feed_forward(self.layer_norm_2(x))\n",
        "    return x"
      ],
      "metadata": {
        "id": "mKR1Y7p2Us5F"
      },
      "execution_count": 53,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class BabyGPTmodel(nn.Module):\n",
        "  def __init__(self, vocab_size, block_size, num_layers, embedded_dim, num_heads):\n",
        "    super(BabyGPTmodel, self).__init__()\n",
        "    self.token = nn.Embedding(vocab_size, embedded_dim)\n",
        "    self.positional_embeddings = nn.Embedding(block_size, embedded_dim)\n",
        "    self.layers1 = nn.ModuleList([Transformer(embedded_dim, num_heads) for _ in range(num_layers)])\n",
        "    self.ln_f = nn.LayerNorm(embedded_dim, eps = 1e-12) # final layer\n",
        "    self.ln_head = nn.Linear(embedded_dim, vocab_size)\n",
        "\n",
        "\n",
        "    # init all weights\n",
        "    ## from karpathy\n",
        "    self.apply(self._init_weights)\n",
        "    # apply special scaled init to the residual projections, per GPT-2 paper\n",
        "    for pn, p in self.named_parameters():\n",
        "      if pn.endswith('projection.weight'):\n",
        "        torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * num_layers))\n",
        "\n",
        "        # report number of parameters\n",
        "        print(\"number of parameters: %d\" % (sum(p.nelement() for p in self.parameters()),))\n",
        "\n",
        "  def _init_weights(self, module):\n",
        "      if isinstance(module, nn.Linear):\n",
        "          torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "          if module.bias is not None:\n",
        "              torch.nn.init.zeros_(module.bias)\n",
        "      elif isinstance(module, nn.Embedding):\n",
        "          torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "  def forward(self, idx):\n",
        "    device = idx.device\n",
        "    b, t = idx.size()\n",
        "    tok_emb = self.token(idx)\n",
        "    position_ids = torch.arange(0, t, dtype = torch.long).unsqueeze(0)\n",
        "    pos_emb = self.positional_embeddings(position_ids)\n",
        "    x = tok_emb + pos_emb\n",
        "    for layers1 in self.layers1:\n",
        "      x = layers1(x)\n",
        "    x = self.ln_f(x)\n",
        "    logits = self.ln_head(x[:, -1, :])\n",
        "    return logits\n"
      ],
      "metadata": {
        "id": "f3vDjpPuUtsu"
      },
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vocab_size =   len(chars)\n",
        "block_size = 64\n",
        "embedded_dim = 256\n",
        "num_heads = 4\n",
        "num_layers = 4\n",
        "\n",
        "gpt = BabyGPTmodel(vocab_size, block_size, num_layers, embedded_dim, num_heads)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zWQ8_ioEmwkf",
        "outputId": "6edc6e19-e381-4712-8bd4-153c06bfa549"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 14527087\n",
            "number of parameters: 14527087\n",
            "number of parameters: 14527087\n",
            "number of parameters: 14527087\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "input_dim = 16\n",
        "rank = 4\n",
        "W_A = nn.Parameter(torch.empty(input_dim, rank))\n",
        "W_A.shape"
      ],
      "metadata": {
        "id": "N3uQZD2AoTeG",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5be40c72-0391-419a-95f3-ecbca2b1a894"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([16, 4])"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class LowRankAttention(nn.Module):\n",
        "    def __init__(self, dim, rank):\n",
        "        super(LowRankAttention, self).__init__()\n",
        "        self.rank = rank\n",
        "        self.Wq = nn.Linear(dim, rank, bias=False)\n",
        "        self.Wk = nn.Linear(dim, rank, bias=False)\n",
        "        self.Wv = nn.Linear(dim, rank, bias=False)\n",
        "        self.Wo = nn.Linear(rank, dim, bias=False)\n",
        "\n",
        "    def forward(self, q, k, v):\n",
        "        Q = self.Wq(q)\n",
        "        K = self.Wk(k)\n",
        "        V = self.Wv(v)\n",
        "\n",
        "        # Compute the attention scores using low-rank approximation\n",
        "        A = torch.bmm(Q, K.transpose(-2, -1)) / (self.rank ** 0.5)\n",
        "\n",
        "        # Softmax along the key dimension\n",
        "        A = torch.softmax(A, dim=-1)\n",
        "\n",
        "        # Compute the attention-weighted values using low-rank approximation\n",
        "        AV = torch.bmm(A, V)\n",
        "\n",
        "        # Apply the output layer to the attention-weighted values\n",
        "        out = self.Wo(AV)\n",
        "\n",
        "        return out\n",
        "\n",
        "class LowRankTransformerLayer(nn.Module):\n",
        "    def __init__(self, dim, rank, dropout=0.2):\n",
        "        super(LowRankTransformerLayer, self).__init__()\n",
        "        self.attention = LowRankAttention(dim, rank)\n",
        "        self.norm1 = nn.LayerNorm(dim)\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "        self.feedforward = nn.Sequential(\n",
        "            nn.Linear(dim, dim * 3),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(dim * 3, dim)\n",
        "        )\n",
        "        self.norm2 = nn.LayerNorm(dim)\n",
        "        self.dropout2 = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Compute the self-attention layer\n",
        "        attention_out = self.attention(x, x, x)\n",
        "\n",
        "        # Add residual connection and normalize\n",
        "        x = self.norm1(x + self.dropout1(attention_out))\n",
        "\n",
        "        # Feed-forward layer\n",
        "        ff_out = self.feedforward(x)\n",
        "\n",
        "        # Add residual connection and normalize\n",
        "        x = self.norm2(x + self.dropout2(ff_out))\n",
        "\n",
        "        return x\n",
        "\n",
        "class LowRankTransformer(nn.Module):\n",
        "    def __init__(self, vocab_size, num_layers, dim, rank, num_heads, dropout= 0.2):\n",
        "        super(LowRankTransformer, self).__init__()\n",
        "        self.layers = nn.ModuleList([LowRankTransformerLayer(dim, rank, dropout) for _ in range(num_layers)])\n",
        "        self.num_layers = num_layers\n",
        "        self.dim = dim\n",
        "        self.rank = rank\n",
        "        self.num_heads = num_heads\n",
        "        self.pos_embedding = nn.Embedding(vocab_size, dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "        # init all weights\n",
        "        ## from karpathy\n",
        "        self.apply(self._init_weights)\n",
        "        # apply special scaled init to the residual projections, per GPT-2 paper\n",
        "        for pn, p in self.named_parameters():\n",
        "          if pn.endswith('Wo.weight'):\n",
        "            torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * num_layers))\n",
        "\n",
        "            # report number of parameters\n",
        "            print(\"number of parameters: %d\" % (sum(p.nelement() for p in self.parameters()),))\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "      if isinstance(module, nn.Linear):\n",
        "        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "        if module.bias is not None:\n",
        "          torch.nn.init.zeros_(module.bias)\n",
        "        elif isinstance(module, nn.Embedding):\n",
        "          torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Add positional embeddings\n",
        "        x = x + self.pos_embedding[:, :x.size(1)]\n",
        "\n",
        "        # Apply dropout\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        # Apply the transformer layers\n",
        "        for layer in self.layers:\n",
        "            x = layer(x)\n",
        "\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "qD35ejNYzi59"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# words = open(r\"/content/ALL_eminem.txt\", 'r', encoding='utf-8').read().split()\n",
        "\n",
        "# chars = sorted(list(set(words)))\n",
        "# vocab_size = len(chars)\n",
        "\n",
        "dim = 256\n",
        "num_heads = 4\n",
        "num_layers = 4\n",
        "rank = 4\n",
        "\n",
        "\n",
        "lrt = LowRankTransformer(vocab_size, num_layers, dim, rank,  num_heads)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "47TM6blM71y3",
        "outputId": "880e7ddf-e830-4432-baba-cdf4437b007e"
      },
      "execution_count": 69,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 7261952\n",
            "number of parameters: 7261952\n",
            "number of parameters: 7261952\n",
            "number of parameters: 7261952\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A comparison between BabyGPT and Low rank adaptation. BabyGPT is b/w  14. 5 M  parametres as it stands . Low rank adaptation improves the parametre efficiency around 7.26 M parametres. (Almost half)\n",
        "\n",
        "Note: Parametre size is also directly linked to context length."
      ],
      "metadata": {
        "id": "KPKToZKsqPHA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### scaling laws\n",
        "\n",
        "We check for the FLOPs for BabyGPT model"
      ],
      "metadata": {
        "id": "-jkjnPRWQALv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "words = open('/content/ALL_eminem.txt', 'r').read().split()\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "chars = sorted(list(set(words)))\n",
        "string2integer = {ch: i for i, ch in enumerate(chars)}\n",
        "# print(string2integer)\n",
        "\n",
        "integer2string = {i:ch for ch,i in string2integer.items()}\n",
        "encode = lambda s: [string2integer[c] for c in s]\n",
        "# print(encode)\n",
        "\n",
        "decode = lambda l: ''.join([integer2string[i] for i in l])\n",
        "# print(decode)\n",
        "\n",
        "data = torch.tensor(encode(words), dtype = torch.long)\n",
        "# print(data)\n",
        "# data.size()"
      ],
      "metadata": {
        "id": "O-p730EcKRHo"
      },
      "execution_count": 98,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# block_size = 32\n",
        "# batch_size = 128\n",
        "# ix = torch.randint(len(data) - block_size, (batch_size,))\n",
        "# x = torch.stack([data[i:i + block_size] for i in ix])\n",
        "# x.shape\n",
        "\n",
        "data.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7lyu96OpK49q",
        "outputId": "c043dc22-1bfe-45c6-a9ee-0200ed617175"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([180194])"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# def gpt_params(seq_len, vocab_size, embedded_dim, num_heads, num_layers):\n",
        "#     \"\"\" Given GPT config calculate total number of parameters \"\"\"\n",
        "#     ffw_size = 4*embedded_dim # in GPT the number of intermediate features is always 4*embedded_dim\n",
        "#     # token and position embeddings\n",
        "#     embeddings = embedded_dim * vocab_size + embedded_dim * seq_len\n",
        "#     # transformer blocks\n",
        "#     attention = 3*embedded_dim**2 + 3*embedded_dim # weights and biases\n",
        "#     attproj = embedded_dim**2 + embedded_dim\n",
        "#     ffw = embedded_dim*(ffw_size) + ffw_size\n",
        "#     ffwproj = ffw_size*embedded_dim + embedded_dim\n",
        "#     layernorms = 2*2*embedded_dim\n",
        "#     # dense\n",
        "#     ln_f = 2*embedded_dim\n",
        "#     dense = embedded_dim*vocab_size # note: no bias here\n",
        "#     # note: embeddings are not included in the param count!\n",
        "#     total_params = num_layers*(attention + attproj + ffw + ffwproj + layernorms) + ln_f + dense\n",
        "#     return total_params\n",
        "\n",
        "# gpt2 = dict(seq_len = 1346, vocab_size = 478, embedded_dim = 16, num_heads = 4, num_layers = 4)\n",
        "# gpt_params(**gpt2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qq9WH0cpJ2P8",
        "outputId": "421f85a9-d17f-4a5f-ef71-1026bac3c336"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "20800"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def count_flops(seq_len, vocab_size, embedded_dim, num_heads, num_layers, ffw_size):\n",
        "  key_size = embedded_dim // num_heads\n",
        "  embeddings = 2 * seq_len * vocab_size * embedded_dim\n",
        "  # attention\n",
        "\n",
        "  attention = 2 * 3 * seq_len * embedded_dim * (key_size * num_heads)\n",
        "  # key @ query logits\n",
        "  attlogits = 2 * seq_len * seq_len * (key_size * num_heads)\n",
        "  # softmax\n",
        "  attsoftmax = 3 * num_heads * seq_len * seq_len # 3* is for subtract (max), exp, divide (?)\n",
        "  # softmax @ value reductions\n",
        "  attvalue = 2 * seq_len * seq_len * (key_size * num_heads)\n",
        "  # final linear\n",
        "  attlinear = 2 * seq_len * (key_size * num_heads) * embedded_dim\n",
        "  att = attention + attlogits + attsoftmax + attvalue + attlinear\n",
        "  # feed forward\n",
        "  dense = 2 * seq_len * (embedded_dim * ffw_size + embedded_dim * ffw_size)\n",
        "\n",
        "  # logits\n",
        "  logits = 2 * seq_len * embedded_dim * vocab_size\n",
        "\n",
        "  # this is what you'd expect:\n",
        "  # forward_flops = embeddings + num_layers * (att + dense) + logits\n",
        "  # but:\n",
        "  # per author correspondence apparently there is typo in the paper,\n",
        "  # they do not count embeddings and logits to repro table 4. So instead:\n",
        "  forward_flops = num_layers * (att + dense)\n",
        "  backward_flops = 2 * forward_flops # as in Kaplan et al. 2020\n",
        "  total_flops = forward_flops + backward_flops\n",
        "\n",
        "  return total_flops\n",
        "\n",
        "\n",
        "class BabyGPTmodel(nn.Module):\n",
        "  def __init__(self, vocab_size, block_size, num_layers, embedded_dim, num_heads):\n",
        "    super(BabyGPTmodel, self).__init__()\n",
        "    self.token = nn.Embedding(vocab_size, embedded_dim)\n",
        "    self.positional_embeddings = nn.Embedding(block_size, embedded_dim)\n",
        "    self.layers1 = nn.ModuleList([Transformer(embedded_dim, num_heads) for _ in range(num_layers)])\n",
        "    self.ln_f = nn.LayerNorm(embedded_dim, eps = 1e-12) # final layer\n",
        "    self.ln_head = nn.Linear(embedded_dim, vocab_size)\n",
        "\n",
        "\n",
        "    # init all weights\n",
        "    ## from karpathy\n",
        "    self.apply(self._init_weights)\n",
        "    # apply special scaled init to the residual projections, per GPT-2 paper\n",
        "    for pn, p in self.named_parameters():\n",
        "      if pn.endswith('projection.weight'):\n",
        "        torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * num_layers))\n",
        "\n",
        "        # report number of parameters\n",
        "        print(\"number of parameters: %d\" % (sum(p.nelement() for p in self.parameters()),))\n",
        "\n",
        "  def _init_weights(self, module):\n",
        "      if isinstance(module, nn.Linear):\n",
        "          torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "          if module.bias is not None:\n",
        "              torch.nn.init.zeros_(module.bias)\n",
        "      elif isinstance(module, nn.Embedding):\n",
        "          torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "  def forward(self, idx):\n",
        "    device = idx.device\n",
        "    b, t = idx.size()\n",
        "    tok_emb = self.token(idx)\n",
        "    position_ids = torch.arange(0, t, dtype = torch.long, device = device).unsqueeze(0)\n",
        "    pos_emb = self.positional_embeddings(position_ids)\n",
        "    x = tok_emb + pos_emb\n",
        "    for layers1 in self.layers1:\n",
        "      x = layers1(x)\n",
        "    x = self.ln_f(x)\n",
        "    logits = self.ln_head(x[:, -1, :])\n",
        "    return logits"
      ],
      "metadata": {
        "id": "ry0RBUYB2v22"
      },
      "execution_count": 70,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vocab_size =  len(chars)\n",
        "block_size = 4\n",
        "embedded_dim =256\n",
        "num_heads = 4\n",
        "num_layers = 4\n",
        "\n",
        "gpt = BabyGPTmodel(vocab_size, block_size, num_layers, embedded_dim, num_heads)\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JKmIAmaL4F1Q",
        "outputId": "09c08c18-61aa-4a6b-de5f-461f9a02d4c7"
      },
      "execution_count": 87,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 14511727\n",
            "number of parameters: 14511727\n",
            "number of parameters: 14511727\n",
            "number of parameters: 14511727\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "seq_len = data.size()\n",
        "seq_len"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A5ANyPsp40Ae",
        "outputId": "414d7357-4d0d-493e-eb0f-7b68cb2743ac"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([180194])"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "vocab_size"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fhXVY5GH5xzq",
        "outputId": "b2b39552-1ff4-4097-f1e6-a7f7a2203dfd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "22127"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "vocab_size = len(chars)\n",
        "ffw_size = embedded_dim * 4\n",
        "flops = count_flops(180194, 22127, 256, 4, 4, 1024)\n",
        "print( flops/1e15, \"PFLOPS\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TB6DnDDyE3VA",
        "outputId": "a06640b0-c587-4e64-f16f-cd01935cbcb2"
      },
      "execution_count": 88,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.407066566638144 PFLOPS\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.save(lrt.state_dict(), '/content/lrt.pth')\n",
        "\n",
        "\n",
        "lrt.load_state_dict(torch.load('/content/lrt.pth'))\n",
        "print(lrt.eval())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oa_Z2cwl-wGn",
        "outputId": "099aba18-fa2c-4873-8d72-91f07a41b22e"
      },
      "execution_count": 79,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "LowRankTransformer(\n",
            "  (layers): ModuleList(\n",
            "    (0-3): 4 x LowRankTransformerLayer(\n",
            "      (attention): LowRankAttention(\n",
            "        (Wq): Linear(in_features=256, out_features=4, bias=False)\n",
            "        (Wk): Linear(in_features=256, out_features=4, bias=False)\n",
            "        (Wv): Linear(in_features=256, out_features=4, bias=False)\n",
            "        (Wo): Linear(in_features=4, out_features=256, bias=False)\n",
            "      )\n",
            "      (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
            "      (dropout1): Dropout(p=0.2, inplace=False)\n",
            "      (feedforward): Sequential(\n",
            "        (0): Linear(in_features=256, out_features=768, bias=True)\n",
            "        (1): GELU(approximate='none')\n",
            "        (2): Linear(in_features=768, out_features=256, bias=True)\n",
            "      )\n",
            "      (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
            "      (dropout2): Dropout(p=0.2, inplace=False)\n",
            "    )\n",
            "  )\n",
            "  (pos_embedding): Embedding(22127, 256)\n",
            "  (dropout): Dropout(p=0.2, inplace=False)\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "quantize the LORA model."
      ],
      "metadata": {
        "id": "3m7U-sriKo_R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import torch.quantization\n",
        "import torch\n",
        "\n",
        "import pandas\n",
        "\n",
        "quantized_lrt = torch.quantization.quantize_dynamic(lrt, {torch.nn.Linear}, dtype=torch.qint8)\n",
        "quantized_lrt.eval()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RMDADjey_eft",
        "outputId": "553b0606-8293-4b2e-ff9a-a0305a1720c4"
      },
      "execution_count": 92,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "LowRankTransformer(\n",
              "  (layers): ModuleList(\n",
              "    (0-3): 4 x LowRankTransformerLayer(\n",
              "      (attention): LowRankAttention(\n",
              "        (Wq): DynamicQuantizedLinear(in_features=256, out_features=4, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
              "        (Wk): DynamicQuantizedLinear(in_features=256, out_features=4, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
              "        (Wv): DynamicQuantizedLinear(in_features=256, out_features=4, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
              "        (Wo): DynamicQuantizedLinear(in_features=4, out_features=256, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
              "      )\n",
              "      (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "      (dropout1): Dropout(p=0.2, inplace=False)\n",
              "      (feedforward): Sequential(\n",
              "        (0): DynamicQuantizedLinear(in_features=256, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
              "        (1): GELU(approximate='none')\n",
              "        (2): DynamicQuantizedLinear(in_features=768, out_features=256, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
              "      )\n",
              "      (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "      (dropout2): Dropout(p=0.2, inplace=False)\n",
              "    )\n",
              "  )\n",
              "  (pos_embedding): Embedding(22127, 256)\n",
              "  (dropout): Dropout(p=0.2, inplace=False)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 92
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let us estimate size."
      ],
      "metadata": {
        "id": "yMJJBrSlKjQa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "lrt.eval()\n",
        "lrt_int8 = torch.ao.quantization.quantize_dynamic(\n",
        "    lrt,  # the original model\n",
        "    {torch.nn.Linear},  # a set of layers to dynamically quantize\n",
        "    dtype=torch.qint8)\n",
        "\n",
        "def print_model_size(lrt):\n",
        "    torch.save(lrt.state_dict(), \"/content/lrt.pth\")\n",
        "    print(\"%.4f MB\" %(os.path.getsize(\"/content/lrt.pth\")/1e6))\n",
        "\n",
        "\n",
        "print_model_size(lrt)\n",
        "print_model_size(lrt_int8)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S4TNB3ilJq0p",
        "outputId": "f81870fc-7ee9-49ce-da70-cddb0403cb23"
      },
      "execution_count": 96,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "29.0614 MB\n",
            "24.3112 MB\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A reduction of about 1.3 has been observed. Not good, not terrible..!!!"
      ],
      "metadata": {
        "id": "svygpWmzKXHf"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "0BZSIcNmKWQM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}