{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyO3uStvnbO22Dv+Bq+SK97r",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "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/llama_implementation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "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",
        "\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",
        "    block_size: int = 2048\n",
        "    vocab_size: int = 32000\n",
        "    num_layers: int = 32\n",
        "    num_heads: int = 32\n",
        "    embedded_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(num_layers=32, num_heads=32, embedded_dim=4096),\n",
        "    \"13B\": dict(num_layers=40, num_heads=40, embedded_dim=5120),\n",
        "    \"30B\": dict(num_layers=60, num_heads=52, embedded_dim=6656),\n",
        "    \"65B\": dict(num_layers=80, num_heads=64, embedded_dim=8192),\n",
        "}\n",
        "\n",
        "\n",
        "class AttentionHead(nn.Module):\n",
        "  def __init__(self, config) -> None:\n",
        "    super().__init__()\n",
        "    assert config.embedded_dim % config.num_heads == 0\n",
        "\n",
        "\n",
        "    self.atten = nn.Linear(config.embedded_dim, 3 * config.embedded_dim, bias=False)\n",
        "    self.c_proj = nn.Linear(config.embedded_dim, config.embedded_dim, bias=False)\n",
        "    self.num_heads = config.num_heads\n",
        "    self.embedded_dim = config.embedded_dim\n",
        "    self.block_size = config.block_size\n",
        "    self.rope_cache: Optional[torch.Tensor] = None\n",
        "    self.register_buffer(\"bias\", torch.tril(torch.ones(config.block_size, config.block_size))\n",
        "                                    .view(1, 1, config.block_size, config.block_size))\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",
        "    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.embedded_dim // self.num_heads, \n",
        "          dtype=x.dtype,\n",
        "          device=x.device,\n",
        "          )\n",
        "      q = apply_rope(q, self.rope_cache)\n",
        "      k = apply_rope(k, self.rope_cache)\n",
        "\n",
        "      # from karpathy\n",
        "      # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
        "      # att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
        "      # att = att.masked_fill(self.bias[:,:,: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.c_proj(y)\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.c_proj(y)\n",
        "\n",
        "      return y\n",
        "\n",
        "\n",
        "class RMSNorm(nn.Module):\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",
        "\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",
        "  # $\\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"
      ],
      "metadata": {
        "id": "68qK3erWlstM"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dropout = 0.2\n",
        "class FeedForward(nn.Module):\n",
        "  def __init__(self, config: LLaMAConfig) -> None:\n",
        "    super().__init__()\n",
        "    hidden_dim = 4 * config.embedded_dim\n",
        "    n_hidden = int(2 * hidden_dim / 3)\n",
        "        \n",
        "\n",
        "    self.c_fc1 = nn.Linear(config.embedded_dim, n_hidden, bias=False)\n",
        "    self.c_fc2 = nn.Linear(config.embedded_dim, n_hidden, bias=False)\n",
        "    self.c_proj = nn.Linear(n_hidden, config.embedded_dim, bias=False)\n",
        "    self.dropout = nn.Dropout(dropout)\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",
        "    x = self.dropout(x)\n",
        "    return x\n"
      ],
      "metadata": {
        "id": "8guVYB3doCfc"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Transformer(nn.Module):\n",
        "  def __init__(self, config: LLaMAConfig) -> None:\n",
        "    super().__init__()\n",
        "    self.attention = AttentionHead(config)\n",
        "    self.rms_norm_1 = RMSNorm(config.embedded_dim)\n",
        "    self.rms_norm_2 = RMSNorm(config.embedded_dim)\n",
        "    self.feedforward = FeedForward(config)\n",
        "\n",
        "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "    x = x + self.attention(self.rms_norm_1(x))\n",
        "    x = x + self.feedforward(self.rms_norm_2(x))\n",
        "    return x"
      ],
      "metadata": {
        "id": "PaEXpANDoDVg"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class BabyGPTmodel(nn.Module):\n",
        "    def __init__(self, config: LLaMAConfig) -> None:\n",
        "        super().__init__()\n",
        "        \n",
        "        self.config = config     \n",
        "        self.token = nn.Embedding(config.vocab_size, config.embedded_dim)\n",
        "        self.positional_embeddings = nn.Embedding(config.block_size, config.embedded_dim)\n",
        "        self.layers1 = nn.ModuleList([Transformer(config) for _ in range(config.num_layers)])\n",
        "        self.ln_f = RMSNorm(config.embedded_dim, eps = 1e-12) # final layer \n",
        "        self.lnum_heads = nn.Linear(config.embedded_dim, config.vocab_size)\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('c_proj.weight'):\n",
        "            torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.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: nn.Module) -> None:\n",
        "        if isinstance(module, nn.Linear):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(2 * self.config.num_layers))\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 * self.config.num_layers))\n",
        "\n",
        "\n",
        "    def forward(self, idx: torch.Tensor) -> torch.Tensor:\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.lnum_heads(x[:, -1, :])\n",
        "      return logits\n",
        "\n",
        "    @classmethod\n",
        "    def from_name(cls, name: str) -> Self:\n",
        "        return cls(LLaMAConfig.from_name(name))"
      ],
      "metadata": {
        "id": "DAlB9mG1oL6I"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config =  LLaMAConfig(block_size = 3,\n",
        "    vocab_size = 478, num_layers  = 4,\n",
        "    num_heads = 4,\n",
        "    embedded_dim = 16)\n",
        "\n",
        "llama = BabyGPTmodel(config)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Utrn-7ET2rMb",
        "outputId": "bc799ce8-4f0b-4642-a44b-e3685c357d51"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 28126\n"
          ]
        }
      ]
    }
  ]
}