05_weightloading/05_part-1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45398736-7e89-4263-89c8-92153baff553",
   "metadata": {},
   "source": [
    "**LLM Workshop 2024 by Sebastian Raschka**\n",
    "\n",
    "This code is based on *Build a Large Language Model (From Scratch)*, [https://github.com/rasbt/LLMs-from-scratch](https://github.com/rasbt/LLMs-from-scratch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66dd524e-864c-4012-b0a2-ccfc56e80024",
   "metadata": {
    "id": "66dd524e-864c-4012-b0a2-ccfc56e80024"
   },
   "source": [
    "<br>\n",
    "<br>\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "# 5) Loading pretrained weights (part 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a57fb9-f69e-44ca-ab29-3537fa5c0157",
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\"matplotlib\", \n",
    "        \"numpy\", \n",
    "        \"tiktoken\", \n",
    "        \"torch\",\n",
    "       ]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83eb6c38-7278-40e0-bd9f-8a2b1feac3ec",
   "metadata": {},
   "source": [
    "- Previously, we only trained a small GPT-2 model using a very small short-story book for educational purposes\n",
    "- Fortunately, we don't have to spend tens to hundreds of thousands of dollars to pretrain the model on a large pretraining corpus but can load pretrained weights (we start with the GPT-2 weights provided by OpenAI)\n",
    "\n",
    "<img src=\"figures/01.png\" width=1000px>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75cab892-a165-4f43-9601-f517bc212ab6",
   "metadata": {},
   "source": [
    "- First, some boilerplate code to download the files from OpenAI and load the weights into Python\n",
    "- Since OpenAI used [TensorFlow](https://www.tensorflow.org/), we will have to install and use TensorFlow for loading the weights; [tqdm](https://github.com/tqdm/tqdm) is a progress bar library\n",
    "- Uncomment and run the next cell to install the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb9fdf02-972a-444e-bf65-8ffcaaf30ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install tensorflow tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0747edc-559c-44ef-a93f-079d60227e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TensorFlow version:\", version(\"tensorflow\"))\n",
    "print(\"tqdm version:\", version(\"tqdm\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5bc89eb-4d39-4287-9b0c-e459ebe7f5ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relative import from the gpt_download.py contained in this folder\n",
    "from gpt_download import download_and_load_gpt2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff76a736-6f9f-4328-872e-f89a7b70a2cc",
   "metadata": {},
   "source": [
    "- We can then download the model weights for the 124 million parameter model as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76271dd7-108d-4f5b-9c01-6ae0aac4b395",
   "metadata": {},
   "outputs": [],
   "source": [
    "settings, params = download_and_load_gpt2(model_size=\"124M\", models_dir=\"gpt2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1a31951-d971-4a6e-9c43-11ee1168ec6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Settings:\", settings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "857c8331-130e-46ba-921d-fa35d7a73cfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Parameter dictionary keys:\", params.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48dac94-8562-4a66-84ef-46c613cdc4cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(params[\"wte\"])\n",
    "print(\"Token embedding weight tensor dimensions:\", params[\"wte\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "466e100c-294e-4afc-a70a-2f398ac4c104",
   "metadata": {},
   "source": [
    "- Alternatively, \"355M\", \"774M\", and \"1558M\" are also supported `model_size` arguments\n",
    "- The difference between these differently sized models is summarized in the figure below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f19d32-5aae-4176-9f86-f391672c8f0d",
   "metadata": {},
   "source": [
    "<img src=\"figures/02.png\" width=800px>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea6e5076-f08d-41fc-bd8b-1cfe53538f41",
   "metadata": {},
   "source": [
    "- Above, we loaded the 124M GPT-2 model weights into Python, however we still need to transfer them into our `GPTModel` instance\n",
    "- First, we initialize a new GPTModel instance\n",
    "- Note that the original GPT model initialized the linear layers for the query, key, and value matrices in the multi-head attention module with bias vectors, which is not required or recommended; however, to be able to load the weights correctly, we have to enable these too by setting `qkv_bias` to `True` in our implementation, too\n",
    "- We are also using the `1024` token context length that was used by the original GPT-2 model(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fef90dd-0654-4667-844f-08e28339ef7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "GPT_CONFIG_124M = {\n",
    "    \"vocab_size\": 50257,   # Vocabulary size\n",
    "    \"context_length\": 256, # Shortened context length (orig: 1024)\n",
    "    \"emb_dim\": 768,        # Embedding dimension\n",
    "    \"n_heads\": 12,         # Number of attention heads\n",
    "    \"n_layers\": 12,        # Number of layers\n",
    "    \"drop_rate\": 0.1,      # Dropout rate\n",
    "    \"qkv_bias\": False      # Query-key-value bias\n",
    "}\n",
    "\n",
    "\n",
    "# Define model configurations in a dictionary for compactness\n",
    "model_configs = {\n",
    "    \"gpt2-small (124M)\": {\"emb_dim\": 768, \"n_layers\": 12, \"n_heads\": 12},\n",
    "    \"gpt2-medium (355M)\": {\"emb_dim\": 1024, \"n_layers\": 24, \"n_heads\": 16},\n",
    "    \"gpt2-large (774M)\": {\"emb_dim\": 1280, \"n_layers\": 36, \"n_heads\": 20},\n",
    "    \"gpt2-xl (1558M)\": {\"emb_dim\": 1600, \"n_layers\": 48, \"n_heads\": 25},\n",
    "}\n",
    "\n",
    "# Copy the base configuration and update with specific model settings\n",
    "model_name = \"gpt2-small (124M)\"  # Example model name\n",
    "NEW_CONFIG = GPT_CONFIG_124M.copy()\n",
    "NEW_CONFIG.update(model_configs[model_name])\n",
    "NEW_CONFIG.update({\"context_length\": 1024, \"qkv_bias\": True})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c20cdc14-684d-4f77-a786-6646f66bae81",
   "metadata": {},
   "outputs": [],
   "source": [
    "from supplementary import GPTModel\n",
    "\n",
    "gpt = GPTModel(NEW_CONFIG)\n",
    "gpt.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "272f29ac-8342-4b3d-a57d-9b0166ced314",
   "metadata": {},
   "source": [
    "- The next task is to assign the OpenAI weights to the corresponding weight tensors in our `GPTModel` instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a92229-c002-49a6-8cfb-248297ad8296",
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign(left, right):\n",
    "    if left.shape != right.shape:\n",
    "        raise ValueError(f\"Shape mismatch. Left: {left.shape}, Right: {right.shape}\")\n",
    "    return torch.nn.Parameter(torch.tensor(right))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f22d5d95-ca5a-425c-a9ec-fc432a12d4e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "def load_weights_into_gpt(gpt, params):\n",
    "    gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])\n",
    "    gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])\n",
    "    \n",
    "    for b in range(len(params[\"blocks\"])):\n",
    "        q_w, k_w, v_w = np.split(\n",
    "            (params[\"blocks\"][b][\"attn\"][\"c_attn\"])[\"w\"], 3, axis=-1)\n",
    "        gpt.trf_blocks[b].att.W_query.weight = assign(\n",
    "            gpt.trf_blocks[b].att.W_query.weight, q_w.T)\n",
    "        gpt.trf_blocks[b].att.W_key.weight = assign(\n",
    "            gpt.trf_blocks[b].att.W_key.weight, k_w.T)\n",
    "        gpt.trf_blocks[b].att.W_value.weight = assign(\n",
    "            gpt.trf_blocks[b].att.W_value.weight, v_w.T)\n",
    "\n",
    "        q_b, k_b, v_b = np.split(\n",
    "            (params[\"blocks\"][b][\"attn\"][\"c_attn\"])[\"b\"], 3, axis=-1)\n",
    "        gpt.trf_blocks[b].att.W_query.bias = assign(\n",
    "            gpt.trf_blocks[b].att.W_query.bias, q_b)\n",
    "        gpt.trf_blocks[b].att.W_key.bias = assign(\n",
    "            gpt.trf_blocks[b].att.W_key.bias, k_b)\n",
    "        gpt.trf_blocks[b].att.W_value.bias = assign(\n",
    "            gpt.trf_blocks[b].att.W_value.bias, v_b)\n",
    "\n",
    "        gpt.trf_blocks[b].att.out_proj.weight = assign(\n",
    "            gpt.trf_blocks[b].att.out_proj.weight, \n",
    "            params[\"blocks\"][b][\"attn\"][\"c_proj\"][\"w\"].T)\n",
    "        gpt.trf_blocks[b].att.out_proj.bias = assign(\n",
    "            gpt.trf_blocks[b].att.out_proj.bias, \n",
    "            params[\"blocks\"][b][\"attn\"][\"c_proj\"][\"b\"])\n",
    "\n",
    "        gpt.trf_blocks[b].ff.layers[0].weight = assign(\n",
    "            gpt.trf_blocks[b].ff.layers[0].weight, \n",
    "            params[\"blocks\"][b][\"mlp\"][\"c_fc\"][\"w\"].T)\n",
    "        gpt.trf_blocks[b].ff.layers[0].bias = assign(\n",
    "            gpt.trf_blocks[b].ff.layers[0].bias, \n",
    "            params[\"blocks\"][b][\"mlp\"][\"c_fc\"][\"b\"])\n",
    "        gpt.trf_blocks[b].ff.layers[2].weight = assign(\n",
    "            gpt.trf_blocks[b].ff.layers[2].weight, \n",
    "            params[\"blocks\"][b][\"mlp\"][\"c_proj\"][\"w\"].T)\n",
    "        gpt.trf_blocks[b].ff.layers[2].bias = assign(\n",
    "            gpt.trf_blocks[b].ff.layers[2].bias, \n",
    "            params[\"blocks\"][b][\"mlp\"][\"c_proj\"][\"b\"])\n",
    "\n",
    "        gpt.trf_blocks[b].norm1.scale = assign(\n",
    "            gpt.trf_blocks[b].norm1.scale, \n",
    "            params[\"blocks\"][b][\"ln_1\"][\"g\"])\n",
    "        gpt.trf_blocks[b].norm1.shift = assign(\n",
    "            gpt.trf_blocks[b].norm1.shift, \n",
    "            params[\"blocks\"][b][\"ln_1\"][\"b\"])\n",
    "        gpt.trf_blocks[b].norm2.scale = assign(\n",
    "            gpt.trf_blocks[b].norm2.scale, \n",
    "            params[\"blocks\"][b][\"ln_2\"][\"g\"])\n",
    "        gpt.trf_blocks[b].norm2.shift = assign(\n",
    "            gpt.trf_blocks[b].norm2.shift, \n",
    "            params[\"blocks\"][b][\"ln_2\"][\"b\"])\n",
    "\n",
    "    gpt.final_norm.scale = assign(gpt.final_norm.scale, params[\"g\"])\n",
    "    gpt.final_norm.shift = assign(gpt.final_norm.shift, params[\"b\"])\n",
    "    gpt.out_head.weight = assign(gpt.out_head.weight, params[\"wte\"])\n",
    "    \n",
    "    \n",
    "load_weights_into_gpt(gpt, params)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "gpt.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f7472cb-54dc-4311-96d8-b2694f885cee",
   "metadata": {},
   "source": [
    "- If the model is loaded correctly, we can use it to generate new text using our previous `generate` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f690253-f845-4347-b7b6-43fabbd2affa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tiktoken\n",
    "from supplementary import (\n",
    "    generate_text_simple,\n",
    "    text_to_token_ids,\n",
    "    token_ids_to_text\n",
    ")\n",
    "\n",
    "\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "token_ids = generate_text_simple(\n",
    "    model=gpt,\n",
    "    idx=text_to_token_ids(\"Every effort moves you\", tokenizer).to(device),\n",
    "    max_new_tokens=10,\n",
    "    context_size=GPT_CONFIG_124M[\"context_length\"]\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d079f98-a7c4-462e-8416-5a64f670861c",
   "metadata": {},
   "source": [
    "- We know that we loaded the model weights correctly because the model can generate coherent text; if we made even a small mistake, the mode would not be able to do that"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a30d071-30cd-43df-ba83-a6b162593b19",
   "metadata": {},
   "source": [
    "<br>\n",
    "<br>\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "# Exercise 1: Trying larger LLMs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a459f5a-578e-4145-8bac-aad43915de0d",
   "metadata": {},
   "source": [
    "- Load one of the larger LLMs and see how the output quality compares\n",
    "- Ask it to answer specific instructions, for example to summarize text or correct the spelling of a sentence"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}