LLaMA/building_llama_complete.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aff495ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from types import SimpleNamespace"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74e401bc",
   "metadata": {},
   "source": [
    "#### a starting point: https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/07_gpt_to_llama/converting-gpt-to-llama2.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd25909e",
   "metadata": {},
   "source": [
    "# RMSNorm\n",
    "replaces LayerNorm\n",
    "\n",
    "$$\n",
    "y_i = \\frac{x_i}{\\text{RMS}(x)}\\gamma_i\\\\\n",
    "RMS(x) = \\sqrt{\\epsilon + \\frac{1}{n} \\sum x_i^2}\n",
    "$$\n",
    "$\\gamma$ is a learnable parameter\n",
    "- $x$ is input, $x_i$ will be one feature/neuron\n",
    "i.e. if x is of shape (1,128,1024) # bs, seq_len, num_fts then we need to normalize along the last dim of 1024 neurons\n",
    "- we init gamma with 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "277a7f08",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, embed_dim):\n",
    "        super().__init__()\n",
    "        self.embed_dim = embed_dim\n",
    "        self.gamma = nn.Parameter(\n",
    "            torch.ones(self.embed_dim,dtype=torch.float32),\n",
    "            requires_grad=True\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x [B, S, D]\n",
    "        mean = x.pow(2).mean(dim=-1,keepdim=True)\n",
    "        r_sqrt = x * torch.rsqrt(mean + 1e-5) # [B, S, 1]\n",
    "        y = r_sqrt * self.gamma\n",
    "        return y.to(x.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0517e8e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 3, 8]), True)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# testing RMSNorm\n",
    "rms = RMSNorm(embed_dim=8)\n",
    "x = torch.rand(1,3,8)\n",
    "rms(x).shape,torch.allclose(rms(x),nn.RMSNorm(8,eps=1e-5)(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "729ac427",
   "metadata": {},
   "source": [
    "# New activation function: SiLU (Swish)\n",
    "\n",
    "$$\\text{SiLU}(x) = x * \\sigma(x)$$\n",
    "$\\sigma(x)$ is sigmoid function "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3980d7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SiLU(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    def forward(self, x):\n",
    "        # x [B S D]\n",
    "        return x * F.sigmoid(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2673be1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# testing SiLU\n",
    "torch.allclose(SiLU()(x),F.silu(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0dd0d26",
   "metadata": {},
   "source": [
    "# New MLP! SwiGLU\n",
    "\n",
    "$$\n",
    "\\text{SwiGLU}(x) = \\text{SiLU}(\\text{linear}_1(x)) * \\text{linear}_2(x) \\\\\n",
    "\\text{output} = \\text{linear}_3(\\text{SwiGLU}(x))\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "083e3748",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.embed_dim = config.embed_dim\n",
    "        self.intermediate_dim = config.intermediate_dim\n",
    "        self.linear1 = nn.Linear(self.embed_dim, self.intermediate_dim, bias=False, dtype=config.dtype)\n",
    "        self.linear2 = nn.Linear(self.embed_dim, self.intermediate_dim, bias=False, dtype=config.dtype)\n",
    "        self.linear3 = nn.Linear(self.intermediate_dim, self.embed_dim, bias=False, dtype=config.dtype)\n",
    "        self.act = SiLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x [B S D]\n",
    "        x1 = self.linear1(x)\n",
    "        x2 = self.linear2(x)\n",
    "        x = self.act(x1) * x2\n",
    "        x = self.linear3(x)\n",
    "        return x\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8c736c",
   "metadata": {},
   "source": [
    "# RoPE\n",
    "rotary positional embeddings!\n",
    "\n",
    "- applied to q,k at MHA step\n",
    "- precomputed angles, their sine and cosine based on model's context length\n",
    "- current implementation input shape: B, nH, S, H\n",
    "\n",
    "$$\n",
    "\\text{inv_freq} = \\frac{1}{\\theta_b^{k/H}}, \\quad k \\in \\{0,2,\\dots,H/2\\}\\\\ \\\\\n",
    "$$\n",
    "$H$ is head_size, ${\\theta_b}$ is base theta = 10,000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a83658a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def precompute_rope(head_dim, base_theta=10_000, context_length=4096):\n",
    "    k = torch.arange(0,head_dim,2)[:head_dim//2].float()\n",
    "    inv_freq = 1 / (base_theta ** (k/head_dim))\n",
    "    \n",
    "    positions = torch.arange(context_length)\n",
    "    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) # [S, H/2]\n",
    "    angles = torch.cat([angles, angles],dim=-1) # [S, H]\n",
    "    \n",
    "    cos = torch.cos(angles) # [S, H]\n",
    "    sin = torch.sin(angles) # [S, H]\n",
    "    \n",
    "    \n",
    "    return cos, sin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e7f2e1c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_rope(x, cos, sin):\n",
    "    B, nH, S, H = x.shape\n",
    "    x1 = x[...,:H//2] # [B, nH, S, H/2]\n",
    "    x2 = x[...,H//2:] # [B, nH, S, H/2]\n",
    "    cos_values = cos[:S,:].unsqueeze(0).unsqueeze(1) # [1,1,S,H]\n",
    "    sin_values = sin[:S,:].unsqueeze(0).unsqueeze(1) # [1,1,S,H]\n",
    "    rotated = torch.cat([-x2,x1],dim=-1)\n",
    "    x_rope = (x * cos_values) + (rotated * sin_values)\n",
    "    return x_rope.to(x.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9ac6dc04",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 256]) torch.Size([256, 1]) torch.Size([256, 256])\n",
      "torch.Size([4096, 512]) torch.Size([4096, 512])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 128, 64])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "head_dim = 512\n",
    "inv_freq2 = 1.0 / (10000 ** (torch.arange(0, head_dim, 2)[: (head_dim // 2)].float() / head_dim))\n",
    "# plt.plot(inv_freq2.numpy())\n",
    "# plt.show()\n",
    "print(torch.arange(256)[None,:].shape ,inv_freq2[:,None].shape, (torch.arange(256)[None,:] * inv_freq2[:,None]).shape)\n",
    "cos, sin = precompute_rope(head_dim)\n",
    "print(cos.shape, sin.shape)\n",
    "\n",
    "\n",
    "# plot it to make sure\n",
    "# import numpy\n",
    "# import matplotlib.pyplot as plt\n",
    "# plt.figure(figsize=(12,36))\n",
    "# plt.imshow(cos.numpy())\n",
    "# plt.show()\n",
    "# plt.figure(figsize=(12,36))\n",
    "# plt.imshow(sin.numpy())\n",
    "# plt.show()\n",
    "\n",
    "x = torch.rand(2,8,128,64)\n",
    "cos,sin = precompute_rope(64,context_length=128)\n",
    "x_rope = apply_rope(x, cos, sin)\n",
    "x_rope.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1eea3c44",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GroupedQueryAttention(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.embed_dim = config.embed_dim\n",
    "        self.num_kv_heads = config.num_kv_heads\n",
    "        self.num_q_heads = config.num_q_heads\n",
    "        \n",
    "        assert self.embed_dim % self.num_q_heads == 0, 'embed_dim should be div. by num. of query heads'\n",
    "        assert self.num_q_heads % self.num_kv_heads ==0, 'num. query heads should be div. by num. key-value heads'\n",
    "        \n",
    "        self.head_dim = self.embed_dim // self.num_q_heads\n",
    "        \n",
    "        self.q_proj = nn.Linear(self.embed_dim, self.head_dim * self.num_q_heads, bias=False, dtype=config.dtype)\n",
    "        self.k_proj = nn.Linear(self.embed_dim, self.head_dim * self.num_kv_heads, bias=False, dtype=config.dtype)\n",
    "        self.v_proj = nn.Linear(self.embed_dim, self.head_dim * self.num_kv_heads, bias=False, dtype=config.dtype)\n",
    "        \n",
    "        self.drop = nn.Dropout(config.attn_dropout)\n",
    "        \n",
    "        self.o_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False, dtype=config.dtype)\n",
    "        \n",
    "        self.register_buffer('causal_mask',\n",
    "                             torch.triu(torch.ones(\n",
    "                                 config.max_position_embeddings,\n",
    "                                 config.max_position_embeddings\n",
    "                             ),diagonal=1))\n",
    "        \n",
    "        cos, sin = precompute_rope(self.head_dim,base_theta=config.base_theta,context_length=config.max_position_embeddings)\n",
    "        self.register_buffer('rope_cos', cos)\n",
    "        self.register_buffer('rope_sin', sin)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x [B S D]\n",
    "        \n",
    "        B,S,D = x.shape\n",
    "        \n",
    "        q = self.q_proj(x) # [B S H*nQ]\n",
    "        k = self.k_proj(x) # [B S H*nKV]\n",
    "        v = self.v_proj(x) # [B S H*nKV]\n",
    "        \n",
    "        q = q.view(B, S, self.num_q_heads, self.head_dim).transpose(1,2) # [B nQ S H]\n",
    "        k = k.view(B, S, self.num_kv_heads, self.head_dim).transpose(1,2) # [B nKV S H]\n",
    "        v = v.view(B, S, self.num_kv_heads, self.head_dim).transpose(1,2) # [B nKV S H]\n",
    "        \n",
    "        q = apply_rope(q, self.rope_cos, self.rope_sin)\n",
    "        k = apply_rope(k, self.rope_cos, self.rope_sin)\n",
    "        \n",
    "        k = k.repeat_interleave(self.num_q_heads//self.num_kv_heads, dim=1) # [B nQ S H]\n",
    "        v = v.repeat_interleave(self.num_q_heads//self.num_kv_heads, dim=1) # [B nQ S H]\n",
    "        \n",
    "        attn = q @ k.transpose(2,3) # [B nQ S H] @ [B nQ H S] = [B nQ S S]\n",
    "        \n",
    "        # apply mask, mul with v, reshape, return\n",
    "        mask = self.causal_mask[:S,:S].bool()\n",
    "        attn.masked_fill_(mask,-torch.inf)\n",
    "        \n",
    "        attn = F.softmax(attn / (self.head_dim ** 0.5), dim=-1)\n",
    "        \n",
    "        attn = self.drop(attn)\n",
    "        \n",
    "        out = attn @ v # [B nQ S S] @ [B nQ S H] = [B nQ S H]\n",
    "        out = out.transpose(1,2) # [B S nQ H]\n",
    "        out = out.reshape(B, S, D)\n",
    "        \n",
    "        proj = self.o_proj(out)\n",
    "        \n",
    "        return proj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a7bbebed",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.ln1 = RMSNorm(config.embed_dim)\n",
    "        self.self_attention = GroupedQueryAttention(config)\n",
    "\n",
    "        self.ln2 = RMSNorm(config.embed_dim)\n",
    "        self.mlp = FeedForward(config)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x [B S D]\n",
    "        skip = x\n",
    "        x = self.ln1(x)\n",
    "        x = self.self_attention(x)\n",
    "        x = x + skip\n",
    "        \n",
    "        skip = x\n",
    "        x = self.ln2(x)\n",
    "        x = self.mlp(x)\n",
    "        x = x + skip\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c02261a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LLaMA(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.token_embedding = nn.Embedding(self.config.vocab_size, self.config.embed_dim, dtype=self.config.dtype)\n",
    "        self.transformer_layers = nn.Sequential(*[\n",
    "            TransformerBlock(self.config) for _ in range(self.config.num_layers)\n",
    "        ])\n",
    "        self.ln = RMSNorm(self.config.embed_dim)\n",
    "        self.lm_head = nn.Linear(self.config.embed_dim, self.config.vocab_size, bias=False, dtype=self.config.dtype)\n",
    "        \n",
    "    def forward(self, input_ids):\n",
    "        # input_ids [B S]\n",
    "        token_embeddings = self.token_embedding(input_ids)\n",
    "        x = self.transformer_layers(token_embeddings)\n",
    "        x = self.ln(x)\n",
    "        logits = self.lm_head(x)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3882e0f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = SimpleNamespace(\n",
    "    embed_dim = 576,\n",
    "    intermediate_dim = 1536,\n",
    "    max_position_embeddings = 8192,\n",
    "    base_theta = 100000,\n",
    "    num_q_heads = 9,\n",
    "    num_kv_heads = 3,\n",
    "    attn_dropout = 0.,\n",
    "    num_layers = 30,\n",
    "    vocab_size = 49152,\n",
    "    dtype = torch.float32,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "954ed46b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LLaMA(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2e7a2144",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float32 (PyTorch default): 8.83 GB with 162,826,560 parameters\n"
     ]
    }
   ],
   "source": [
    "def model_memory_size(model, input_dtype=torch.float32):\n",
    "    total_params = 0\n",
    "    total_grads = 0\n",
    "    for param in model.parameters():\n",
    "        # Calculate total number of elements per parameter\n",
    "        param_size = param.numel()\n",
    "        total_params += param_size\n",
    "        # Check if gradients are stored for this parameter\n",
    "        if param.requires_grad:\n",
    "            total_grads += param_size\n",
    "\n",
    "    # Calculate buffer size (non-parameters that require memory)\n",
    "    total_buffers = sum(buf.numel() for buf in model.buffers())\n",
    "\n",
    "    # Size in bytes = (Number of elements) * (Size of each element in bytes)\n",
    "    # We assume parameters and gradients are stored in the same type as input dtype\n",
    "    element_size = torch.tensor(0, dtype=input_dtype).element_size()\n",
    "    total_memory_bytes = (total_params + total_grads + total_buffers) * element_size\n",
    "\n",
    "    # Convert bytes to gigabytes\n",
    "    total_memory_gb = total_memory_bytes / (1024**3)\n",
    "\n",
    "    return total_memory_gb, total_params\n",
    "\n",
    "total_mem, total_params = model_memory_size(model, input_dtype=torch.float32)\n",
    "print(f\"float32 (PyTorch default): {total_mem:.2f} GB with {total_params:,} parameters\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d5b6c75b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[28301, 26290, 33026, 39400,  1447,  1477, 28014,  7785, 40707, 13677]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randint(0,config.vocab_size,(1,10)).long()\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "49644e26",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = model(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9413d424",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 10, 49152])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}