{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyM9nJQX/i8Ii0RaLeUQzlak",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "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/trainer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### BabyGPT trainer."
      ],
      "metadata": {
        "id": "kG2T8jLsgekR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "import math\n",
        "from dataclasses import dataclass\n",
        "from typing import Optional\n",
        "\n",
        "# hyperparameters\n",
        "\n",
        "@dataclass\n",
        "class GPTConfig:\n",
        "    # these are default GPT-2 hyperparameters\n",
        "    block_size: int = 1024\n",
        "    vocab_size: int = 50304\n",
        "    n_layer: int = 12\n",
        "    n_head: int = 12\n",
        "    n_embd: int = 768\n",
        "    bias :bool = False\n",
        "\n",
        "\n",
        "### other hyperparametres\n",
        "batch_size = 64\n",
        "max_iters = 11000\n",
        "learning_rate = 3e-4\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "eval_interval = 500\n",
        "eval_iters = 200\n",
        "dropout = 0.2\n",
        "\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "words = open(r\"/content/ALL_eminem.txt\", 'r', encoding='utf-8').read()\n",
        "\n",
        "chars = sorted(list(set(words)))\n",
        "vocab_size = len(chars)\n",
        "\n",
        "\n",
        "string2integer = {ch: i for i, ch in enumerate(chars)}\n",
        "integer2string = {i:ch for ch,i in string2integer.items()}\n",
        "encode = lambda s: [string2integer[c] for c in s]\n",
        "decode = lambda l: ''.join([integer2string[i] for i in l])\n",
        "data = torch.tensor(encode(words), dtype = torch.long)\n",
        "\n",
        "\n",
        "## train and split the data\n",
        "n = int(0.8*len(data))\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]\n",
        "\n",
        "# data loading\n",
        "def get_batch(split):\n",
        "    # generate a small batch of data of inputs x and targets y\n",
        "    data = train_data if split == 'train' else val_data\n",
        "    ix = torch.randint(len(data) - config.block_size, (batch_size,))\n",
        "    x = torch.stack([data[i:i+ config.block_size] for i in ix])\n",
        "    y = torch.stack([data[i+1:i+ config.block_size+1] for i in ix])\n",
        "    x, y = x.to(device), y.to(device)\n",
        "    return x, y\n",
        "\n",
        "\n",
        "### from pytorch GPT tutorial\n",
        "@torch.no_grad()\n",
        "def estimate_loss():\n",
        "    out = {}\n",
        "    model.eval()\n",
        "    for split in ['train', 'val']:\n",
        "        losses = torch.zeros(eval_iters)\n",
        "        for k in range(eval_iters):\n",
        "            X, Y = get_batch(split)\n",
        "            logits, loss = model(X, Y)\n",
        "            losses[k] = loss.item()\n",
        "        out[split] = losses.mean()\n",
        "    model.train()\n",
        "    return out\n",
        "\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "class Attention(nn.Module):\n",
        "  def __init__(self, config):\n",
        "    super(Attention, self).__init__()\n",
        "\n",
        "    assert config.n_embd % config.n_head == 0\n",
        "\n",
        "    self.atten = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)\n",
        "    self.projection = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
        "    self.n_head = config.n_head\n",
        "    self.n_embd = config.n_embd\n",
        "    self.register_buffer('tril', torch.tril(torch.ones(config.block_size, config.block_size)))\n",
        "\n",
        "  def forward(self, x):\n",
        "    B,T,C = x.size()\n",
        "    q, k ,v  = self.atten(x).split(self.n_embd, dim=2)\n",
        "    q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
        "    k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
        "    v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\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",
        "    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\n",
        "\n",
        "dropout = 0.2\n",
        "class FeedForward(nn.Module):\n",
        "  def __init__(self,config):\n",
        "    super(FeedForward, self).__init__()\n",
        "    self.net = nn.Sequential(nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias),\n",
        "    nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias),\n",
        "    nn.GELU(),\n",
        "    nn.Dropout(dropout))\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.net(x)\n",
        "\n",
        "### A simple Transformer Block\n",
        "class Transformer(nn.Module):\n",
        "  def __init__(self,config):\n",
        "    super(Transformer, self).__init__()\n",
        "    self.attention = Attention(config)\n",
        "    self.feed_forward = FeedForward(config)\n",
        "    self.layer_norm_1 = nn.LayerNorm(config.n_embd)\n",
        "    self.layer_norm_2 = nn.LayerNorm(config.n_embd)\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\n",
        "\n",
        "\n",
        "class BabyGPTmodel(nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super(BabyGPTmodel, self).__init__()\n",
        "\n",
        "        assert config.vocab_size is not None\n",
        "        assert config.block_size is not None\n",
        "\n",
        "        self.config = config\n",
        "        self.token = nn.Embedding(config.vocab_size, config.n_embd)\n",
        "        self.positional_embeddings = nn.Embedding(config.block_size, config.n_embd)\n",
        "        self.blocks = nn.Sequential(*[Transformer(config) for _ in range(config.n_layer)])\n",
        "        self.ln_f = nn.LayerNorm(config.n_embd, eps = 1e-12) # final layer norm\n",
        "        self.lnum_heads = nn.Linear(config.n_embd, config.vocab_size)\n",
        "\n",
        "        ## init all weights\n",
        "        ## from karpathy\n",
        "        self.apply(self._init_weights)\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 * config.n_layer))\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 / math.sqrt(2 * config.n_layer))\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 / math.sqrt(2 * config.n_layer))\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        device = idx.device\n",
        "        B, T = idx.shape\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 block in self.blocks:\n",
        "          x = self.blocks(x)\n",
        "        x = self.ln_f(x)\n",
        "        logits = self.lnum_heads(x)\n",
        "\n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            B, T, C = logits.shape\n",
        "            logits = logits.view(B*T, C)\n",
        "            targets = targets.view(B*T)\n",
        "            loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "\n",
        "    ## from karpathy's youtube videos.\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        # idx is (B, T) array of indices in the current context\n",
        "        for _ in range(max_new_tokens):\n",
        "            # crop idx to the last block_size tokens\n",
        "            idx_cond = idx[:, -config.block_size:]\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx_cond)\n",
        "            # focus only on the last time step\n",
        "            logits = logits[:, -1, :] # becomes (B, C)\n",
        "            # apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
        "            # sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
        "            # append sampled index to the running sequence\n",
        "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "\n",
        "config = GPTConfig(\n",
        "    block_size = 4,\n",
        "    vocab_size = len(chars),\n",
        "    n_head = 4,\n",
        "    n_layer = 4,\n",
        "    n_embd = 16)\n",
        "\n",
        "model = BabyGPTmodel(config)\n",
        "\n",
        "m = model.to(device)\n",
        "\n",
        "# create a PyTorch optimizer\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "for iter in range(max_iters):\n",
        "\n",
        "    # every once in a while evaluate the loss on train and val sets\n",
        "    if iter % eval_interval == 0 or iter == max_iters - 1:\n",
        "        losses = estimate_loss()\n",
        "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
        "\n",
        "    # sample a batch of data\n",
        "    xb, yb = get_batch('train')\n",
        "\n",
        "    # evaluate the loss\n",
        "    logits, loss = model(xb, yb)\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "# generate from the model\n",
        "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
        "print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kflnFouSHdM1",
        "outputId": "46e23193-7148-487f-fc9d-1dd6c3496aaa"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 16237\n",
            "step 0: train loss 4.6938, val loss 4.6940\n",
            "step 500: train loss 3.1403, val loss 3.1382\n",
            "step 1000: train loss 2.7314, val loss 2.7349\n",
            "step 1500: train loss 2.5625, val loss 2.5675\n",
            "step 2000: train loss 2.4877, val loss 2.5002\n",
            "step 2500: train loss 2.4302, val loss 2.4225\n",
            "step 3000: train loss 2.4015, val loss 2.4092\n",
            "step 3500: train loss 2.3725, val loss 2.3798\n",
            "step 4000: train loss 2.3553, val loss 2.3565\n",
            "step 4500: train loss 2.3320, val loss 2.3338\n",
            "step 5000: train loss 2.3159, val loss 2.3222\n",
            "step 5500: train loss 2.3110, val loss 2.3110\n",
            "step 6000: train loss 2.2861, val loss 2.2992\n",
            "step 6500: train loss 2.2905, val loss 2.2927\n",
            "step 7000: train loss 2.2843, val loss 2.2974\n",
            "step 7500: train loss 2.2793, val loss 2.2889\n",
            "step 8000: train loss 2.2767, val loss 2.2738\n",
            "step 8500: train loss 2.2616, val loss 2.2740\n",
            "step 9000: train loss 2.2576, val loss 2.2935\n",
            "step 9500: train loss 2.2596, val loss 2.2714\n",
            "step 10000: train loss 2.2609, val loss 2.2769\n",
            "step 10500: train loss 2.2558, val loss 2.2645\n",
            "step 10999: train loss 2.2538, val loss 2.2680\n",
            "\n",
            "Bit with ick as can lin' a deed Kam\n",
            "Caroull I'lld you arint allol 'es to there oof out a raries\n",
            "I finct on an my nover to mad to conghecte, way Nevis you meakinin' a get no?)\n",
            "\"quies\n",
            "An're me, hyoin' my\n",
            "\n",
            "Frog Dover, Ome be oll kare (Dlyer a be ir]\n",
            "And tike a be\n",
            "Trick gat a tirkffockn evul in to you that dim that acough leike wo, get oof dade foundd you sain my a your gins bit I's So shit\n",
            "Saixchere the witg It bit ?, is's the you toine\n",
            "\"'St luak, Davin' aup my at\n",
            "Litsty wam wamined\n",
            "Wame,\n",
            "So chum f\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### llama trainer with MFU"
      ],
      "metadata": {
        "id": "VCsAmIIPgVO9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## from lit-llama repo (partially).\n",
        "\n",
        "\n",
        "import math\n",
        "from dataclasses import dataclass\n",
        "from typing import Optional\n",
        "import time\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "from typing_extensions import Self\n",
        "\n",
        "@dataclass\n",
        "class LLaMAConfig:\n",
        "    ## LLaMa parametres\n",
        "    block_size: int = 2048\n",
        "    vocab_size: int = 32000\n",
        "    n_layer: int = 32\n",
        "    n_head: int = 32\n",
        "    dim: int = 4096\n",
        "    dim :int = 4096\n",
        "\n",
        "\n",
        "    @classmethod\n",
        "    def from_name(cls, name: str) -> Self:\n",
        "        return cls(**llama_configs[name])\n",
        "\n",
        "\n",
        "llama_configs = {\n",
        "    \"7B\": dict(n_layer=32, n_head = 32, n_embd=4096),\n",
        "    \"13B\": dict(n_layer=40, n_head =40, n_embd=5120),\n",
        "    \"30B\": dict(n_layer=60, n_head=52, n_embd=6656),\n",
        "    \"65B\": dict(n_layer=80, n_head =64, n_embd=8192),\n",
        "}\n",
        "\n",
        "\n",
        "### other hyperparametres\n",
        "\n",
        "batch_size = 64\n",
        "max_iters = 3000\n",
        "learning_rate = 3e-4\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "eval_interval = 500\n",
        "eval_iters = 200\n",
        "dropout = 0.2\n",
        "\n",
        "words = open(r\"/content/ALL_eminem.txt\", 'r', encoding='utf-8').read()\n",
        "\n",
        "\n",
        "chars = sorted(list(set(words)))\n",
        "vocab_size = len(chars)\n",
        "\n",
        "\n",
        "string2integer = {ch: i for i, ch in enumerate(chars)}\n",
        "integer2string = {i:ch for ch,i in string2integer.items()}\n",
        "encode = lambda s: [string2integer[c] for c in s]\n",
        "decode = lambda l: ''.join([integer2string[i] for i in l])\n",
        "data = torch.tensor(encode(words), dtype = torch.long)\n",
        "\n",
        "\n",
        "## train and split the data\n",
        "n = int(0.8*len(data))\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]\n",
        "\n",
        "# data loading\n",
        "def get_batch(split):\n",
        "    # generate a small batch of data of inputs x and targets y\n",
        "    data = train_data if split == 'train' else val_data\n",
        "    ix = torch.randint(len(data) - config.block_size, (batch_size,))\n",
        "    x = torch.stack([data[i:i+ config.block_size] for i in ix])\n",
        "    y = torch.stack([data[i+1:i+ config.block_size+1] for i in ix])\n",
        "    x, y = x.to(device), y.to(device)\n",
        "    return x, y\n",
        "\n",
        "\n",
        "### from pytorch GPT tutorial\n",
        "@torch.no_grad()\n",
        "def estimate_loss():\n",
        "    out = {}\n",
        "    model.eval()\n",
        "    for split in ['train', 'val']:\n",
        "        losses = torch.zeros(eval_iters)\n",
        "        for k in range(eval_iters):\n",
        "            X, Y = get_batch(split)\n",
        "            logits, loss = model(X, Y)\n",
        "            losses[k] = loss.item()\n",
        "        out[split] = losses.mean()\n",
        "    model.train()\n",
        "    return out\n",
        "\n",
        "\n",
        "class Attention(nn.Module):\n",
        "  def __init__(self, config : LLaMAConfig):\n",
        "    super(Attention, self).__init__()\n",
        "\n",
        "    self.config = config\n",
        "    self.atten = nn.Linear(config.dim, 3 * config.dim)\n",
        "    self.projection = nn.Linear(config.dim, config.dim)\n",
        "    self.n_head = config.n_head\n",
        "    self.dim = config.dim\n",
        "    self.block_size = config.block_size\n",
        "    self.rope_cache: Optional[torch.Tensor] = None\n",
        "    self.register_buffer('tril', torch.tril(torch.ones(config.block_size, config.block_size)))\n",
        "\n",
        "  def forward(self, x):\n",
        "    B,T,C = x.size()\n",
        "    q, k ,v  = self.atten(x).split(self.dim, dim=2)\n",
        "    q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
        "    k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
        "    v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
        "\n",
        "\n",
        "    if self.rope_cache is None:\n",
        "      # cache for future forward calls\n",
        "      self.rope_cache = build_rope_cache(\n",
        "      seq_len=self.block_size,\n",
        "      n_elem=self.dim // self.n_head,\n",
        "      dtype=x.dtype,\n",
        "      device=x.device,\n",
        "            )\n",
        "\n",
        "      q = apply_rope(q, self.rope_cache)\n",
        "      k = apply_rope(k, self.rope_cache)\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",
        "    # y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
        "\n",
        "    # efficient attention using Flash Attention CUDA kernels\n",
        "    y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)\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",
        "\n",
        "    return y\n",
        "\n",
        "class FeedForward(nn.Module):\n",
        "  def __init__(self, config: LLaMAConfig) -> None:\n",
        "    super().__init__()\n",
        "    hidden_dim = 4 * config.dim\n",
        "    n_hidden = int(2 * hidden_dim / 3)\n",
        "\n",
        "\n",
        "    self.c_fc1 = nn.Linear(config.dim, n_hidden, bias=False)\n",
        "    self.c_fc2 = nn.Linear(config.dim, n_hidden, bias=False)\n",
        "    self.c_proj = nn.Linear(n_hidden,  config.dim, bias=False)\n",
        "\n",
        "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "    x = F.silu(self.c_fc1(x)) * self.c_fc2(x)\n",
        "    x = self.c_proj(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "### A simple Transformer Block\n",
        "class Transformer(nn.Module):\n",
        "  def __init__(self, config : LLaMAConfig) -> None:\n",
        "    super(Transformer, self).__init__()\n",
        "    self.attention = Attention(config)\n",
        "    self.feed_forward = FeedForward(config)\n",
        "    self.layer_norm_1 = RMSNorm(config.dim)\n",
        "    self.layer_norm_2 = RMSNorm(config.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\n",
        "\n",
        "class RMSNorm(nn.Module):\n",
        "    \"\"\"Root Mean Square Layer Normalization.\n",
        "\n",
        "    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:\n",
        "    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:\n",
        "        super().__init__()\n",
        "        self.scale = nn.Parameter(torch.ones(size))\n",
        "        self.eps = eps\n",
        "        self.dim = dim\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        # NOTE: the original RMSNorm paper implementation is not equivalent\n",
        "        # norm_x = x.norm(2, dim=self.dim, keepdim=True)\n",
        "        # rms_x = norm_x * d_x ** (-1. / 2)\n",
        "        # x_normed = x / (rms_x + self.eps)\n",
        "        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)\n",
        "        x_normed = x * torch.rsqrt(norm_x + self.eps)\n",
        "        return self.scale * x_normed\n",
        "\n",
        "\n",
        "def build_rope_cache(seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000) -> torch.Tensor:\n",
        "\n",
        "    \"\"\"Enhanced Transformer with Rotary Position Embedding.\n",
        "\n",
        "    Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/\n",
        "    transformers/rope/__init__.py. MIT License:\n",
        "    https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.\n",
        "    \"\"\"\n",
        "    # $\\Theta = {\\theta_i = 10000^{\\frac{2(i-1)}{d}}, i \\in [1, 2, ..., \\frac{d}{2}]}$\n",
        "    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))\n",
        "\n",
        "    # Create position indexes `[0, 1, ..., seq_len - 1]`\n",
        "    seq_idx = torch.arange(seq_len, dtype=dtype, device=device)\n",
        "\n",
        "    # Calculate the product of position index and $\\theta_i$\n",
        "    idx_theta = torch.outer(seq_idx, theta).float()\n",
        "\n",
        "    cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)\n",
        "\n",
        "    # this is to mimic the behaviour of complex32, else we will get different results\n",
        "    if dtype in (torch.float16, torch.bfloat16, torch.int8):\n",
        "        cache = cache.half()\n",
        "    return cache\n",
        "\n",
        "\n",
        "def apply_rope(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:\n",
        "    x = x.transpose(1, 2)\n",
        "\n",
        "    # truncate to support variable sizes\n",
        "    T = x.size(1)\n",
        "    rope_cache = rope_cache[:T]\n",
        "\n",
        "    # cast because the reference does\n",
        "    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)\n",
        "    rope_cache = rope_cache.view(1, xshaped.size(1), 1, xshaped.size(3), 2)\n",
        "    x_out2 = torch.stack(\n",
        "        [xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],\n",
        "         xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],\n",
        "        ], -1)\n",
        "\n",
        "    x_out2 = x_out2.flatten(3)\n",
        "    return x_out2.transpose(1, 2).type_as(x)\n",
        "\n",
        "class BabyGPTmodel(nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super(BabyGPTmodel, self).__init__()\n",
        "\n",
        "        assert config.vocab_size is not None\n",
        "        assert config.block_size is not None\n",
        "\n",
        "        self.config = config\n",
        "        self.token = nn.Embedding(config.vocab_size, config.dim)\n",
        "        self.positional_embeddings = nn.Embedding(config.block_size, config.dim)\n",
        "        self.blocks = nn.Sequential(*[Transformer(config) for _ in range(config.n_layer)])\n",
        "        self.ln_f = RMSNorm(config.dim, eps = 1e-12) # final layer norm\n",
        "        self.lnum_heads = nn.Linear(config.dim, config.vocab_size)\n",
        "\n",
        "        ## init all weights\n",
        "        ## from karpathy\n",
        "        self.apply(self._init_weights)\n",
        "        for pn, p in self.named_parameters():\n",
        "          if pn.endswith('attention.weight'):\n",
        "            torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))\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 / math.sqrt(2 * config.n_layer))\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 / math.sqrt(2 * config.n_layer))\n",
        "\n",
        "    def num_params(self):\n",
        "      n_params = sum(p.numel() for p in self.parameters())\n",
        "      return n_params\n",
        "\n",
        "\n",
        "    def model_flops(self, for_back, dt):\n",
        "      # from https://arxiv.org/pdf/2204.02311.pdf section B\n",
        "      cfg = self.config\n",
        "      N = self.num_params()\n",
        "\n",
        "      H = cfg.n_head\n",
        "      Q = cfg.dim// config.n_head\n",
        "      T = cfg.block_size\n",
        "      L = cfg.n_layer\n",
        "      flops = 6*N + 12*L*H*Q*T\n",
        "      flops_per_for_back = flops * T\n",
        "      flops_per_iteration =  for_back * flops_per_for_back\n",
        "      flops_received= flops_per_iteration * (1.0/dt) # per second\n",
        "      theoretical_flops = 8e12  # tesla t4 has about 8.1 TFLOPS\n",
        "      mfu = flops_received / theoretical_flops\n",
        "\n",
        "\n",
        "      return mfu\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        device = idx.device\n",
        "        B, T = idx.shape\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 block in self.blocks:\n",
        "          x = self.blocks(x)\n",
        "        x = self.ln_f(x)\n",
        "        logits = self.lnum_heads(x)\n",
        "\n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            B, T, C = logits.shape\n",
        "            logits = logits.view(B*T, C)\n",
        "            targets = targets.view(B*T)\n",
        "            loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "\n",
        "    ## from karpathy's youtube videos.\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        # idx is (B, T) array of indices in the current context\n",
        "        for _ in range(max_new_tokens):\n",
        "            # crop idx to the last block_size tokens\n",
        "            idx_cond = idx[:, -config.block_size:]\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx_cond)\n",
        "            # focus only on the last time step\n",
        "            logits = logits[:, -1, :] # becomes (B, C)\n",
        "            # apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
        "            # sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
        "            # append sampled index to the running sequence\n",
        "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "\n",
        "config = LLaMAConfig(\n",
        "    block_size = 64,\n",
        "    vocab_size = len(chars),\n",
        "    n_head = 4,\n",
        "    n_layer = 4,\n",
        "    dim = 256)\n",
        "\n",
        "model = BabyGPTmodel(config)\n",
        "\n",
        "m = model.to(device)\n",
        "\n",
        "# create a PyTorch optimizer\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "for iter in range(max_iters):\n",
        "\n",
        "    t0 = time.time()\n",
        "    # every once in a while evaluate the loss on train and val sets\n",
        "    if iter % eval_interval == 0 or iter == max_iters - 1:\n",
        "        losses = estimate_loss()\n",
        "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
        "\n",
        "    # sample a batch of data\n",
        "    xb, yb = get_batch('train')\n",
        "\n",
        "    # evaluate the loss\n",
        "    logits, loss = model(xb, yb)\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    t1 = time.time()\n",
        "    dt = t1 - t0\n",
        "\n",
        "# generate from the model\n",
        "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
        "print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))\n",
        "mfu = model.model_flops(batch_size * 1, dt)\n",
        "print(mfu)\n",
        "print(f\" Model Flop Utilization: {mfu*100:.10f}%\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TE_NRBZXbLW7",
        "outputId": "829c22ba-de3d-4a4f-89ec-05652e3b7514"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 3222381\n",
            "step 0: train loss 4.7187, val loss 4.7180\n",
            "step 500: train loss 1.8811, val loss 1.9138\n",
            "step 1000: train loss 1.6014, val loss 1.6814\n",
            "step 1500: train loss 1.4488, val loss 1.6036\n",
            "step 2000: train loss 1.3455, val loss 1.5789\n",
            "step 2500: train loss 1.2653, val loss 1.5880\n",
            "step 2999: train loss 1.1818, val loss 1.6149\n",
            "\n",
            "The fuckin' Outside three, I'm 'bout to tig for\n",
            "So they geteratures but I had a crew feet \n",
            "I can head if I was gone\n",
            "'Cause we don't know my world's always the greatest then the same rape slugs\n",
            "It's till I said her baby\n",
            "I wanna impologire\n",
            "\n",
            "[Joyner Royce's \"L-Tony Manside while he beeps must\n",
            "Call the fore-elliever\n",
            "This is just shut shit down passed the copathy games hill on\n",
            "Entil platinums we're tungry albouncers\n",
            "Man a calls we can't say it\n",
            "Call  on fire exact then it'll breathing her retard\n",
            "Now I\n",
            "0.0005369315281254354\n",
            " Model Flop Utilization: 0.0536931528%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Both The models starts to converge towards the end, which can be fixed. (We need the loss value to reach local minima)\n",
        "\n",
        "\n",
        "Due to the use of Cuda efficient attention, the llama tokens seem more readable."
      ],
      "metadata": {
        "id": "-je2R6-HgkQK"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "```\n",
        "``` number of parameters: 3222381 ```\n",
        "\n",
        "step 0: train loss 4.6894, val loss 4.6895\n",
        "step 500: train loss 2.1731, val loss 2.1832\n",
        "step 1000: train loss 1.7580, val loss 1.8032\n",
        "step 1500: train loss 1.5790, val loss 1.6645\n",
        "step 2000: train loss 1.4482, val loss 1.5992\n",
        "step 2500: train loss 1.3538, val loss 1.5874\n",
        "step 3000: train loss 1.2574, val loss 1.5971\n",
        "step 3500: train loss 1.1835, val loss 1.6254\n",
        "step 4000: train loss 1.1090, val loss 1.6647\n",
        "step 4500: train loss 1.0183, val loss 1.7308\n",
        "step 5000: train loss 0.9455, val loss 1.7937\n",
        "step 5500: train loss 0.8657, val loss 1.8646\n",
        "step 6000: train loss 0.8049, val loss 1.9732\n",
        "step 6500: train loss 0.7404, val loss 2.0515\n",
        "step 7000: train loss 0.6722, val loss 2.1346\n",
        "step 7500: train loss 0.6294, val loss 2.2333\n",
        "step 8000: train loss 0.5899, val loss 2.3033\n",
        "step 8500: train loss 0.5504, val loss 2.4246\n",
        "step 9000: train loss 0.5236, val loss 2.4614\n",
        "step 9500: train loss 0.4916, val loss 2.5494\n",
        "step 10000: train loss 0.4680, val loss 2.6631\n",
        "step 10500: train loss 0.4448, val loss 2.6970\n",
        "step 10999: train loss 0.4341, val loss 2.7462\n",
        "\n",
        "Detroit, revior myself 'til I confused to get the big clead Mastles\n",
        "Slaughterhouse on the blue, that's when he pine I'm hop with the cowprinton\n",
        "robaly I want to a lox on my tempt\n",
        "\n",
        "But now we can't never find a gift killed broke\n",
        "Big before anyone could ever hear the first as I was cooped chill\n",
        "But i this o for a big star\n",
        "I said get chased up!\n",
        "(Hello darkness, my old friend)[Eminem:]\n",
        "If my legacy I acged buving in the tub (might what?)\n",
        "I would know one [*Barrns, worried :]\n",
        "Yeah, so kon bitch, it's\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "J93uBlLWiEm8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "```\n",
        "! git clone https://github.com/soumyadip1995/BabyGPT.git\n",
        "\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "kTx88t_RcAaF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Model FLOP Utilization\n",
        "\n",
        "MFU is the ratio of\n",
        "the observed throughput (tokens-per-second), relative to the theoretical maximum throughput of a system operating\n",
        "at peak FLOPs.\n",
        "\n",
        " L, H, Q, and T are the number of layers, the number of heads, the head dimension, and the\n",
        "sequence length respectively\n",
        "\n",
        "If a given group of accelerators have a total theoretical peak matmul throughput of P\n",
        "FLOPs per second, the model FLOPs utilization is the ratio of the achieved throughput in tokens per second to the\n",
        "theoretical peak throughput R =\n",
        "P/\n",
        "(6N + 12LHQT)\n",
        ".\n",
        "\n",
        "Tesla t4 has a theoritical peak matmul of 8.1 TFLOPS.\n",
        "This can be calculated with a few lines.\n"
      ],
      "metadata": {
        "id": "5y6YsjqOtwlL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "    def model_flops(self, for_back, dt):\n",
        "      # from https://arxiv.org/pdf/2204.02311.pdf section B\n",
        "      cfg = self.config\n",
        "      N = self.num_params()\n",
        "      \n",
        "      H = cfg.n_head\n",
        "      Q = cfg.n_embd // config.n_head\n",
        "      T = cfg.block_size\n",
        "      L = cfg.n_layer\n",
        "      flops = 6*N + 12*L*H*Q*T\n",
        "      flops_per_for_back = flops * T\n",
        "      flops_per_iteration =  for_back * flops_per_for_back\n",
        "      flops_received= flops_per_iteration * (1.0/dt) # per second\n",
        "      theoretical_flops = 8e12  # tesla t4 has about 8.1 TFLOPS\n",
        "      mfu = flops_received / theoretical_flops\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "HAruCYcOm9Fm"
      }
    }
  ]
}