ch05/02_alternative_weight_loading/weight-loading-hf-transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6d6bc54f-2b16-4b0f-be69-957eed5d112f",
   "metadata": {},
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72953590-5363-4398-85ce-54bde07f3d8a",
   "metadata": {},
   "source": [
    "# Bonus Code for Chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4ab5ee-e7b9-45d3-a82b-a12bcfc0945a",
   "metadata": {},
   "source": [
    "## Alternative Weight Loading from Hugging Face Model Hub using Transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2feea87-49f0-48b9-b925-b8f0dda4096f",
   "metadata": {},
   "source": [
    "- In the main chapter, we loaded the GPT model weights directly from OpenAI\n",
    "- This notebook provides alternative weight loading code to load the model weights from the [Hugging Face Model Hub](https://huggingface.co/docs/hub/en/models-the-hub) using the `transformers` Python library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "99b77109-5215-4d07-a618-4d10eff1a488",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b0467eff-b43c-4a38-93e8-5ed87a5fc2b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy version: 1.24.3\n",
      "torch version: 2.3.0\n",
      "transformers version: 4.41.2\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\"numpy\", \"torch\", \"transformers\"]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ffc17d7d-bcd8-42ee-82a9-04fd55acf15d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT2Model(\n",
       "  (wte): Embedding(50257, 768)\n",
       "  (wpe): Embedding(1024, 768)\n",
       "  (drop): Dropout(p=0.1, inplace=False)\n",
       "  (h): ModuleList(\n",
       "    (0-11): 12 x GPT2Block(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPT2Attention(\n",
       "        (c_attn): Conv1D()\n",
       "        (c_proj): Conv1D()\n",
       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPT2MLP(\n",
       "        (c_fc): Conv1D()\n",
       "        (c_proj): Conv1D()\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import GPT2Model\n",
    "\n",
    "\n",
    "# allowed model names\n",
    "model_names = {\n",
    "    \"gpt2-small (124M)\": \"openai-community/gpt2\",\n",
    "    \"gpt2-medium (355M)\": \"openai-community/gpt2-medium\",\n",
    "    \"gpt2-large (774M)\": \"openai-community/gpt2-large\",\n",
    "    \"gpt2-xl (1558M)\": \"openai-community/gpt2-xl\"\n",
    "}\n",
    "\n",
    "CHOOSE_MODEL = \"gpt2-small (124M)\"\n",
    "\n",
    "gpt_hf = GPT2Model.from_pretrained(model_names[CHOOSE_MODEL], cache_dir=\"checkpoints\")\n",
    "gpt_hf.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9ea9b1bc-7881-46ad-9555-27a9cf23faa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_CONFIG = {\n",
    "    \"vocab_size\": 50257,    # Vocabulary size\n",
    "    \"context_length\": 1024, # Context length\n",
    "    \"drop_rate\": 0.0,       # Dropout rate\n",
    "    \"qkv_bias\": True        # Query-key-value bias\n",
    "}\n",
    "\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",
    "\n",
    "BASE_CONFIG.update(model_configs[CHOOSE_MODEL])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4e2a4cf4-a54e-4307-9141-fb9f288e4dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign_check(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(right.clone().detach())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "75be3077-f141-44bb-af88-62580ffd224c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "def load_weights(gpt, gpt_hf):\n",
    "\n",
    "    d = gpt_hf.state_dict()\n",
    "\n",
    "    gpt.pos_emb.weight = assign_check(gpt.pos_emb.weight, d[\"wpe.weight\"])\n",
    "    gpt.tok_emb.weight = assign_check(gpt.tok_emb.weight, d[\"wte.weight\"])\n",
    "    \n",
    "    for b in range(BASE_CONFIG[\"n_layers\"]):\n",
    "        q_w, k_w, v_w = np.split(d[f\"h.{b}.attn.c_attn.weight\"], 3, axis=-1)\n",
    "        gpt.trf_blocks[b].att.W_query.weight = assign_check(gpt.trf_blocks[b].att.W_query.weight, q_w.T)\n",
    "        gpt.trf_blocks[b].att.W_key.weight = assign_check(gpt.trf_blocks[b].att.W_key.weight, k_w.T)\n",
    "        gpt.trf_blocks[b].att.W_value.weight = assign_check(gpt.trf_blocks[b].att.W_value.weight, v_w.T)\n",
    "    \n",
    "        q_b, k_b, v_b = np.split(d[f\"h.{b}.attn.c_attn.bias\"], 3, axis=-1)\n",
    "        gpt.trf_blocks[b].att.W_query.bias = assign_check(gpt.trf_blocks[b].att.W_query.bias, q_b)\n",
    "        gpt.trf_blocks[b].att.W_key.bias = assign_check(gpt.trf_blocks[b].att.W_key.bias, k_b)\n",
    "        gpt.trf_blocks[b].att.W_value.bias = assign_check(gpt.trf_blocks[b].att.W_value.bias, v_b)\n",
    "    \n",
    "    \n",
    "        gpt.trf_blocks[b].att.out_proj.weight = assign_check(gpt.trf_blocks[b].att.out_proj.weight, d[f\"h.{b}.attn.c_proj.weight\"].T)\n",
    "        gpt.trf_blocks[b].att.out_proj.bias = assign_check(gpt.trf_blocks[b].att.out_proj.bias, d[f\"h.{b}.attn.c_proj.bias\"])\n",
    "    \n",
    "        gpt.trf_blocks[b].ff.layers[0].weight = assign_check(gpt.trf_blocks[b].ff.layers[0].weight, d[f\"h.{b}.mlp.c_fc.weight\"].T)\n",
    "        gpt.trf_blocks[b].ff.layers[0].bias = assign_check(gpt.trf_blocks[b].ff.layers[0].bias, d[f\"h.{b}.mlp.c_fc.bias\"])\n",
    "        gpt.trf_blocks[b].ff.layers[2].weight = assign_check(gpt.trf_blocks[b].ff.layers[2].weight, d[f\"h.{b}.mlp.c_proj.weight\"].T)\n",
    "        gpt.trf_blocks[b].ff.layers[2].bias = assign_check(gpt.trf_blocks[b].ff.layers[2].bias, d[f\"h.{b}.mlp.c_proj.bias\"])\n",
    "    \n",
    "        gpt.trf_blocks[b].norm1.scale = assign_check(gpt.trf_blocks[b].norm1.scale, d[f\"h.{b}.ln_1.weight\"])\n",
    "        gpt.trf_blocks[b].norm1.shift = assign_check(gpt.trf_blocks[b].norm1.shift, d[f\"h.{b}.ln_1.bias\"])\n",
    "        gpt.trf_blocks[b].norm2.scale = assign_check(gpt.trf_blocks[b].norm2.scale, d[f\"h.{b}.ln_2.weight\"])\n",
    "        gpt.trf_blocks[b].norm2.shift = assign_check(gpt.trf_blocks[b].norm2.shift, d[f\"h.{b}.ln_2.bias\"])\n",
    "    \n",
    "        gpt.final_norm.scale = assign_check(gpt.final_norm.scale, d[f\"ln_f.weight\"])\n",
    "        gpt.final_norm.shift = assign_check(gpt.final_norm.shift, d[f\"ln_f.bias\"])\n",
    "        gpt.out_head.weight = assign_check(gpt.out_head.weight, d[\"wte.weight\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cda44d37-92c0-4c19-a70a-15711513afce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from previous_chapters import GPTModel\n",
    "\n",
    "\n",
    "gpt = GPTModel(BASE_CONFIG)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "load_weights(gpt, gpt_hf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4ddd0d51-3ade-4890-9bab-d63f141d095f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves forward, but it's not enough.\n",
      "\n",
      "\"I'm not going to sit here and say, 'I'm not going to do this,'\n"
     ]
    }
   ],
   "source": [
    "import tiktoken\n",
    "from previous_chapters import generate, text_to_token_ids, token_ids_to_text\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "\n",
    "token_ids = generate(\n",
    "    model=gpt.to(device),\n",
    "    idx=text_to_token_ids(\"Every effort moves\", tokenizer).to(device),\n",
    "    max_new_tokens=30,\n",
    "    context_size=BASE_CONFIG[\"context_length\"],\n",
    "    top_k=1,\n",
    "    temperature=1.0\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}