{ "cells": [ { "cell_type": "markdown", "id": "6d6bc54f-2b16-4b0f-be69-957eed5d112f", "metadata": {}, "source": [ "\n", "Supplementary code for \"Build a Large Language Model From Scratch\": https://www.manning.com/books/build-a-large-language-model-from-scratch by Sebastian Raschka
\n", "Code repository: https://github.com/rasbt/LLMs-from-scratch\n", "" ] }, { "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.25.2\n", "torch version: 2.2.1\n", "transformers version: 4.33.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": [ { "name": "stderr", "output_type": "stream", "text": [ "/Users/sebastian/miniforge3/envs/book/lib/python3.11/site-packages/transformers/utils/generic.py:311: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n", " torch.utils._pytree._register_pytree_node(\n", "/Users/sebastian/miniforge3/envs/book/lib/python3.11/site-packages/transformers/utils/generic.py:311: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n", " torch.utils._pytree._register_pytree_node(\n" ] }, { "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\": {\"emb_dim\": 768, \"n_layers\": 12, \"n_heads\": 12},\n", " \"gpt2-medium\": {\"emb_dim\": 1024, \"n_layers\": 24, \"n_heads\": 16},\n", " \"gpt2-large\": {\"emb_dim\": 1280, \"n_layers\": 36, \"n_heads\": 20},\n", " \"gpt2-xl\": {\"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(torch.tensor(right))" ] }, { "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": [ { "name": "stderr", "output_type": "stream", "text": [ "/var/folders/jg/tpqyh1fd5js5wsr1d138k3n40000gn/T/ipykernel_32618/3877979348.py:4: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n", " return torch.nn.Parameter(torch.tensor(right))\n" ] } ], "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)\n", "gpt.to(device);" ] }, { "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,\n", " idx=text_to_token_ids(\"Every effort moves\", tokenizer),\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.6" } }, "nbformat": 4, "nbformat_minor": 5 }