{ "cells": [ { "cell_type": "markdown", "id": "e1b280ab-b61f-4d1a-bf7e-44e5f9ed3a5c", "metadata": { "id": "e1b280ab-b61f-4d1a-bf7e-44e5f9ed3a5c" }, "source": [ "\n", "\n", "\n", "\n", "\n", "
\n", "\n", "Supplementary code for the Build a Large Language Model From Scratch book by Sebastian Raschka
\n", "
Code repository: https://github.com/rasbt/LLMs-from-scratch\n", "\n", "		\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "efde77f2-6af3-4781-8597-89ecd3f41a52", "metadata": { "id": "efde77f2-6af3-4781-8597-89ecd3f41a52" }, "source": [ "# Qwen3 Mixture-of-Experts From Scratch (A Standalone Notebook)" ] }, { "cell_type": "markdown", "id": "55cdef4d-de59-4a65-89f9-fa2a8ef3471d", "metadata": { "id": "55cdef4d-de59-4a65-89f9-fa2a8ef3471d" }, "source": [ "- This notebook is purposefully minimal and focuses on the code to implement Qwen3-30B-A3B model (with support for **Coder**, **Instruct** and **Thinking** variants); for more information about this model, please see the original blog post, technical report, and model hub pages:\n", " - [Qwen3: Think Deeper, Act Faster](https://qwenlm.github.io/blog/qwen3/)\n", " - [Qwen3 Technical Report](https://arxiv.org/abs/2505.09388)\n", " - https://huggingface.co/Qwen/Qwen3-Coder-30B-A3B-Instruct (Qwen3 Coder Flash)\n", " - https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507 (new thinking model)\n", " - https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507 (new instruct model)\n", " - https://huggingface.co/Qwen/Qwen3-30B-A3B (original Instruct/Thinking hybrid model)\n", "- Many architectural components in Qwen3 are similar to Llama 3; for a step-by-step guide that explains the individual components and the relationship between GPT and the components used here, you may like the GPT-to-Llama conversion notebooks:\n", " - [Converting a From-Scratch GPT Architecture to Llama 2](../07_gpt_to_llama/converting-gpt-to-llama2.ipynb)\n", " - [Converting Llama 2 to Llama 3.2 From Scratch](../07_gpt_to_llama/converting-llama2-to-llama3.ipynb)\n", " \n", "\n", "**By default, this notebook runs Qwen3-Coder-30B-A3B-Instruct (aka Qwen3 Coder Flash) and requires 80 GB of VRAM (e.g., a single A100 or H100). Note that [this related KV-cache notebook](standalone-qwen3-moe-plus-kvcache.ipynb) adds more code complexity but runs about 3x faster.**\n", "\n", "
\n", "\n", "\n", "\n", "
\n", " \n", "- About the code:\n", " - all code is my own code, mapping the Qwen3 architecture onto the model code implemented in my [Build A Large Language Model (From Scratch)](http://mng.bz/orYv) book; the code is released under a permissive open-source Apache 2.0 license (see [LICENSE.txt](https://github.com/rasbt/LLMs-from-scratch/blob/main/LICENSE.txt))" ] }, { "cell_type": "code", "execution_count": 1, "id": "7c201adb-747e-437b-9a62-442802941e01", "metadata": {}, "outputs": [], "source": [ "# pip install -r https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/refs/heads/main/ch05/07_gpt_to_llama/requirements-extra.txt" ] }, { "cell_type": "code", "execution_count": 2, "id": "dd1b65a8-4301-444a-bd7c-a6f2bd1df9df", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "dd1b65a8-4301-444a-bd7c-a6f2bd1df9df", "outputId": "4f762354-e0a3-4cc2-e5d4-e61a227a202c" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "huggingface_hub version: 0.35.0\n", "tokenizers version: 0.22.1\n", "torch version: 2.7.1+cu128\n" ] } ], "source": [ "from importlib.metadata import version\n", "\n", "pkgs = [\n", " \"huggingface_hub\", # to download pretrained weights\n", " \"tokenizers\", # to implement the tokenizer\n", " \"torch\", # to implement the model\n", "]\n", "for p in pkgs:\n", " print(f\"{p} version: {version(p)}\")" ] }, { "cell_type": "markdown", "id": "653410a6-dd2b-4eb2-a722-23d9782e726d", "metadata": { "id": "653410a6-dd2b-4eb2-a722-23d9782e726d" }, "source": [ " \n", "# 1. Architecture code" ] }, { "cell_type": "code", "execution_count": 3, "id": "82076c21-9331-4dcd-b017-42b046cf1a60", "metadata": { "id": "82076c21-9331-4dcd-b017-42b046cf1a60" }, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "\n", "\n", "class FeedForward(nn.Module):\n", " def __init__(self, cfg):\n", " super().__init__()\n", " self.fc1 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n", " self.fc2 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n", " self.fc3 = nn.Linear(cfg[\"hidden_dim\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"], bias=False)\n", "\n", " def forward(self, x):\n", " x_fc1 = self.fc1(x)\n", " x_fc2 = self.fc2(x)\n", " x = nn.functional.silu(x_fc1) * x_fc2\n", " return self.fc3(x)\n", "\n", "\n", "class MoEFeedForward(nn.Module):\n", " def __init__(self, cfg):\n", " super().__init__()\n", " self.num_experts_per_tok = cfg[\"num_experts_per_tok\"]\n", " self.num_experts = cfg[\"num_experts\"]\n", " self.emb_dim = cfg[\"emb_dim\"]\n", " self.gate = nn.Linear(cfg[\"emb_dim\"], cfg[\"num_experts\"], bias=False, dtype=cfg[\"dtype\"])\n", "\n", " self.fc1 = nn.ModuleList([nn.Linear(cfg[\"emb_dim\"], cfg[\"moe_hidden_dim\"], bias=False, dtype=cfg[\"dtype\"])\n", " for _ in range(cfg[\"num_experts\"])])\n", " self.fc2 = nn.ModuleList([nn.Linear(cfg[\"emb_dim\"], cfg[\"moe_hidden_dim\"], bias=False, dtype=cfg[\"dtype\"])\n", " for _ in range(cfg[\"num_experts\"])])\n", " self.fc3 = nn.ModuleList([nn.Linear(cfg[\"moe_hidden_dim\"], cfg[\"emb_dim\"], bias=False, dtype=cfg[\"dtype\"])\n", " for _ in range(cfg[\"num_experts\"])])\n", "\n", " def forward(self, x):\n", " scores = self.gate(x) # (b, seq_len, num_experts)\n", " topk_scores, topk_indices = torch.topk(scores, self.num_experts_per_tok, dim=-1)\n", " topk_probs = torch.softmax(topk_scores, dim=-1)\n", "\n", " batch, seq_len, _ = x.shape\n", " x_flat = x.reshape(batch * seq_len, -1)\n", " out_flat = torch.zeros(batch * seq_len, self.emb_dim, device=x.device, dtype=x.dtype)\n", "\n", " topk_indices_flat = topk_indices.reshape(-1, self.num_experts_per_tok)\n", " topk_probs_flat = topk_probs.reshape(-1, self.num_experts_per_tok)\n", "\n", " unique_experts = torch.unique(topk_indices_flat)\n", "\n", " for expert_id_tensor in unique_experts:\n", " expert_id = int(expert_id_tensor.item())\n", " mask = topk_indices_flat == expert_id\n", " if not mask.any():\n", " continue\n", "\n", " token_mask = mask.any(dim=-1)\n", " selected_idx = token_mask.nonzero(as_tuple=False).squeeze(-1)\n", " if selected_idx.numel() == 0:\n", " continue\n", "\n", " expert_input = x_flat.index_select(0, selected_idx)\n", " hidden = torch.nn.functional.silu(self.fc1[expert_id](expert_input)) * self.fc2[expert_id](expert_input)\n", " expert_out = self.fc3[expert_id](hidden)\n", "\n", " mask_selected = mask[selected_idx]\n", " slot_indices = mask_selected.int().argmax(dim=-1, keepdim=True)\n", " selected_probs = torch.gather(topk_probs_flat.index_select(0, selected_idx), dim=-1, index=slot_indices).squeeze(-1)\n", "\n", " out_flat.index_add_(0, selected_idx, expert_out * selected_probs.unsqueeze(-1))\n", "\n", " return out_flat.reshape(batch, seq_len, self.emb_dim)" ] }, { "cell_type": "code", "execution_count": 4, "id": "56715760-37e1-433e-89da-04864c139a9e", "metadata": {}, "outputs": [], "source": [ "class RMSNorm(nn.Module):\n", " def __init__(self, emb_dim, eps=1e-6, bias=False, qwen3_compatible=True):\n", " super().__init__()\n", " self.eps = eps\n", " self.qwen3_compatible = qwen3_compatible\n", " self.scale = nn.Parameter(torch.ones(emb_dim))\n", " self.shift = nn.Parameter(torch.zeros(emb_dim)) if bias else None\n", "\n", " def forward(self, x):\n", " input_dtype = x.dtype\n", "\n", " if self.qwen3_compatible:\n", " x = x.to(torch.float32)\n", "\n", " variance = x.pow(2).mean(dim=-1, keepdim=True)\n", " norm_x = x * torch.rsqrt(variance + self.eps)\n", " norm_x = norm_x * self.scale\n", "\n", " if self.shift is not None:\n", " norm_x = norm_x + self.shift\n", "\n", " return norm_x.to(input_dtype)" ] }, { "cell_type": "code", "execution_count": 5, "id": "4b9a346f-5826-4083-9162-abd56afc03f0", "metadata": { "id": "4b9a346f-5826-4083-9162-abd56afc03f0" }, "outputs": [], "source": [ "def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, dtype=torch.float32):\n", " assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n", "\n", " # Compute the inverse frequencies\n", " inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype)[: (head_dim // 2)].float() / head_dim))\n", "\n", " # Generate position indices\n", " positions = torch.arange(context_length, dtype=dtype)\n", "\n", " # Compute the angles\n", " angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) # Shape: (context_length, head_dim // 2)\n", "\n", " # Expand angles to match the head_dim\n", " angles = torch.cat([angles, angles], dim=1) # Shape: (context_length, head_dim)\n", "\n", " # Precompute sine and cosine\n", " cos = torch.cos(angles)\n", " sin = torch.sin(angles)\n", "\n", " return cos, sin\n", "\n", "\n", "def apply_rope(x, cos, sin):\n", " # x: (batch_size, num_heads, seq_len, head_dim)\n", " batch_size, num_heads, seq_len, head_dim = x.shape\n", " assert head_dim % 2 == 0, \"Head dimension must be even\"\n", "\n", " # Split x into first half and second half\n", " x1 = x[..., : head_dim // 2] # First half\n", " x2 = x[..., head_dim // 2 :] # Second half\n", "\n", " # Adjust sin and cos shapes\n", " cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(0) # Shape: (1, 1, seq_len, head_dim)\n", " sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(0)\n", "\n", " # Apply the rotary transformation\n", " rotated = torch.cat((-x2, x1), dim=-1)\n", " x_rotated = (x * cos) + (rotated * sin)\n", "\n", " # It's ok to use lower-precision after applying cos and sin rotation\n", " return x_rotated.to(dtype=x.dtype)" ] }, { "cell_type": "code", "execution_count": 6, "id": "e8169ab5-f976-4222-a2e1-eb1cabf267cb", "metadata": { "id": "e8169ab5-f976-4222-a2e1-eb1cabf267cb" }, "outputs": [], "source": [ "class GroupedQueryAttention(nn.Module):\n", " def __init__(\n", " self, d_in, num_heads, num_kv_groups, head_dim=None, qk_norm=False, dtype=None\n", " ):\n", " super().__init__()\n", " assert num_heads % num_kv_groups == 0, \"num_heads must be divisible by num_kv_groups\"\n", "\n", " self.num_heads = num_heads\n", " self.num_kv_groups = num_kv_groups\n", " self.group_size = num_heads // num_kv_groups\n", "\n", " if head_dim is None:\n", " assert d_in % num_heads == 0, \"`d_in` must be divisible by `num_heads` if `head_dim` is not set\"\n", " head_dim = d_in // num_heads\n", "\n", " self.head_dim = head_dim\n", " self.d_out = num_heads * head_dim\n", "\n", " self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)\n", " self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)\n", " self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)\n", "\n", " self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)\n", "\n", " if qk_norm:\n", " self.q_norm = RMSNorm(head_dim, eps=1e-6)\n", " self.k_norm = RMSNorm(head_dim, eps=1e-6)\n", " else:\n", " self.q_norm = self.k_norm = None\n", "\n", " def forward(self, x, mask, cos, sin):\n", " b, num_tokens, _ = x.shape\n", "\n", " # Apply projections\n", " queries = self.W_query(x) # (b, num_tokens, num_heads * head_dim)\n", " keys = self.W_key(x) # (b, num_tokens, num_kv_groups * head_dim)\n", " values = self.W_value(x) # (b, num_tokens, num_kv_groups * head_dim)\n", "\n", " # Reshape\n", " queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)\n", " keys = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)\n", " values = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)\n", "\n", " # Optional normalization\n", " if self.q_norm:\n", " queries = self.q_norm(queries)\n", " if self.k_norm:\n", " keys = self.k_norm(keys)\n", "\n", " # Apply RoPE\n", " queries = apply_rope(queries, cos, sin)\n", " keys = apply_rope(keys, cos, sin)\n", "\n", " # Expand K and V to match number of heads\n", " keys = keys.repeat_interleave(self.group_size, dim=1)\n", " values = values.repeat_interleave(self.group_size, dim=1)\n", "\n", " # Attention\n", " attn_scores = queries @ keys.transpose(2, 3)\n", " attn_scores = attn_scores.masked_fill(mask, -torch.inf)\n", " attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)\n", "\n", " context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)\n", " return self.out_proj(context)" ] }, { "cell_type": "code", "execution_count": 7, "id": "457cb2f8-50c1-4045-8a74-f181bfb5fea9", "metadata": { "id": "457cb2f8-50c1-4045-8a74-f181bfb5fea9" }, "outputs": [], "source": [ "class TransformerBlock(nn.Module):\n", " def __init__(self, cfg):\n", " super().__init__()\n", " self.att = GroupedQueryAttention(\n", " d_in=cfg[\"emb_dim\"],\n", " num_heads=cfg[\"n_heads\"],\n", " head_dim=cfg[\"head_dim\"],\n", " num_kv_groups=cfg[\"n_kv_groups\"],\n", " qk_norm=cfg[\"qk_norm\"],\n", " dtype=cfg[\"dtype\"]\n", " )\n", " if cfg[\"num_experts\"] > 0:\n", " self.ff = MoEFeedForward(cfg)\n", " else:\n", " self.ff = FeedForward(cfg)\n", " self.norm1 = RMSNorm(cfg[\"emb_dim\"], eps=1e-6)\n", " self.norm2 = RMSNorm(cfg[\"emb_dim\"], eps=1e-6)\n", "\n", " def forward(self, x, mask, cos, sin):\n", " # Shortcut connection for attention block\n", " shortcut = x\n", " x = self.norm1(x)\n", " x = self.att(x, mask, cos, sin) # Shape [batch_size, num_tokens, emb_size]\n", " x = x + shortcut # Add the original input back\n", "\n", " # Shortcut connection for feed-forward block\n", " shortcut = x\n", " x = self.norm2(x)\n", " x = self.ff(x)\n", " x = x + shortcut # Add the original input back\n", "\n", " return x" ] }, { "cell_type": "code", "execution_count": 8, "id": "e88de3e3-9f07-42cc-816b-28dbd46e96c4", "metadata": { "id": "e88de3e3-9f07-42cc-816b-28dbd46e96c4" }, "outputs": [], "source": [ "class Qwen3Model(nn.Module):\n", " def __init__(self, cfg):\n", " super().__init__()\n", "\n", " # Main model parameters\n", " self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"])\n", "\n", " self.trf_blocks = nn.ModuleList( # ModuleList since Sequential can only accept one input, and we need `x, mask, cos, sin`\n", " [TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])]\n", " )\n", "\n", " self.final_norm = RMSNorm(cfg[\"emb_dim\"])\n", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n", "\n", " # Reusuable utilities\n", " if cfg[\"head_dim\"] is None:\n", " head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n", " else:\n", " head_dim = cfg[\"head_dim\"]\n", " cos, sin = compute_rope_params(\n", " head_dim=head_dim,\n", " theta_base=cfg[\"rope_base\"],\n", " context_length=cfg[\"context_length\"]\n", " )\n", " self.register_buffer(\"cos\", cos, persistent=False)\n", " self.register_buffer(\"sin\", sin, persistent=False)\n", " self.cfg = cfg\n", "\n", "\n", " def forward(self, in_idx):\n", " # Forward pass\n", " tok_embeds = self.tok_emb(in_idx)\n", " x = tok_embeds\n", "\n", " num_tokens = x.shape[1]\n", " mask = torch.triu(torch.ones(num_tokens, num_tokens, device=x.device, dtype=torch.bool), diagonal=1)\n", " \n", " for block in self.trf_blocks:\n", " x = block(x, mask, self.cos, self.sin)\n", " x = self.final_norm(x)\n", " logits = self.out_head(x.to(self.cfg[\"dtype\"]))\n", " return logits" ] }, { "cell_type": "markdown", "id": "be2d201f-74ad-4d63-ab9c-601b00674a48", "metadata": { "id": "be2d201f-74ad-4d63-ab9c-601b00674a48" }, "source": [ " \n", "# 2. Initialize model" ] }, { "cell_type": "code", "execution_count": 9, "id": "caa142fa-b375-4e78-b392-2072ced666f3", "metadata": { "id": "caa142fa-b375-4e78-b392-2072ced666f3" }, "outputs": [], "source": [ "# Same config for\n", "\n", "# https://huggingface.co/Qwen/Qwen3-Coder-30B-A3B-Instruct (Qwen3 Coder Flash)\n", "# https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507\n", "# https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507\n", "# https://huggingface.co/Qwen/Qwen3-30B-A3B (original Instruct/Thinking hybrid model)\n", "\n", "QWEN3_CONFIG = {\n", " \"vocab_size\": 151_936,\n", " \"context_length\": 262_144,\n", " \"emb_dim\": 2048,\n", " \"n_heads\": 32,\n", " \"n_layers\": 48,\n", " \"head_dim\": 128,\n", " \"qk_norm\": True,\n", " \"n_kv_groups\": 4,\n", " \"rope_base\": 10_000_000.0,\n", " \"dtype\": torch.bfloat16,\n", " \"num_experts\": 128,\n", " \"num_experts_per_tok\": 8,\n", " \"moe_hidden_dim\": 768,\n", "}" ] }, { "cell_type": "code", "execution_count": 10, "id": "313effd0", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "cuda\n" ] } ], "source": [ "if torch.cuda.is_available():\n", " device = torch.device(\"cuda\")\n", "elif torch.backends.mps.is_available():\n", " device = torch.device(\"mps\")\n", "else:\n", " device = torch.device(\"cpu\")\n", "\n", "print(device)" ] }, { "cell_type": "code", "execution_count": 11, "id": "156253fe-aacd-4da2-8f13-705f05c4b11e", "metadata": { "id": "156253fe-aacd-4da2-8f13-705f05c4b11e" }, "outputs": [], "source": [ "torch.manual_seed(123)\n", "\n", "with device:\n", " model = Qwen3Model(QWEN3_CONFIG)\n", "\n", "#model.to(device)" ] }, { "cell_type": "markdown", "id": "90aca91d-4bee-45ce-993a-4ec5393abe2b", "metadata": {}, "source": [ "- A quick check that the forward pass works before continuing (nan values are ok for now since we are using a \"meta\" device upon instantiation to save memory):" ] }, { "cell_type": "code", "execution_count": 12, "id": "adf0a6b7-b688-42c9-966e-c223d34db99f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "tensor([[[ 0.3223, -0.0562, 0.2490, ..., 0.4551, -0.0542, 0.8242],\n", " [ 0.0688, 0.0786, -0.0312, ..., 0.6406, -0.9141, 0.8672],\n", " [-0.6172, 0.4121, 0.3750, ..., 0.1699, -0.2500, 0.6953]]],\n", " device='cuda:0', dtype=torch.bfloat16, grad_fn=)" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "model(torch.tensor([1, 2, 3]).unsqueeze(0).to(device))" ] }, { "cell_type": "code", "execution_count": 13, "id": "364e76ca-52f8-4fa5-af37-c4069f9694bc", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "364e76ca-52f8-4fa5-af37-c4069f9694bc", "outputId": "00d7e983-262e-4c65-f322-f4d999311988" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Total number of parameters: 30,532,122,624\n", "\n", "Total number of unique parameters: 30,220,957,696\n" ] } ], "source": [ "total_params = sum(p.numel() for p in model.parameters())\n", "print(f\"Total number of parameters: {total_params:,}\")\n", "\n", "# Account for weight tying\n", "total_params_normalized = total_params - model.tok_emb.weight.numel()\n", "print(f\"\\nTotal number of unique parameters: {total_params_normalized:,}\")" ] }, { "cell_type": "code", "execution_count": 14, "id": "fd5efb03-5a07-46e8-8607-93ed47549d2b", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "fd5efb03-5a07-46e8-8607-93ed47549d2b", "outputId": "65c1a95e-b502-4150-9e2e-da619d9053d5" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "float32 (PyTorch default): 227.73 GB\n", "bfloat16: 113.87 GB\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\n", "\n", "print(f\"float32 (PyTorch default): {model_memory_size(model, input_dtype=torch.float32):.2f} GB\")\n", "print(f\"bfloat16: {model_memory_size(model, input_dtype=torch.bfloat16):.2f} GB\")" ] }, { "cell_type": "markdown", "id": "4686eeb7-281f-4c5c-b37a-ed21d0a10427", "metadata": {}, "source": [ "- Don't be concerned; the model runs fine on an A100 card with 80 GB RAM due to offloading some layers to CPU RAM" ] }, { "cell_type": "markdown", "id": "c172f89f-d301-439f-b809-46169e5f5945", "metadata": { "id": "c172f89f-d301-439f-b809-46169e5f5945" }, "source": [ " \n", "# 4. Load pretrained weights" ] }, { "cell_type": "code", "execution_count": 15, "id": "75166128-5899-4995-9b88-9672e135650e", "metadata": { "id": "75166128-5899-4995-9b88-9672e135650e" }, "outputs": [], "source": [ "def load_weights_into_qwen(model, param_config, params):\n", " def assign(left, right, tensor_name=\"unknown\"):\n", " if left.shape != right.shape:\n", " raise ValueError(f\"Shape mismatch in tensor '{tensor_name}'. Left: {left.shape}, Right: {right.shape}\")\n", " \n", " with torch.no_grad():\n", " if isinstance(right, torch.Tensor):\n", " left.copy_(right)\n", " else:\n", " left.copy_(torch.as_tensor(right, dtype=left.dtype, device=left.device))\n", " \n", " return left \n", "\n", " model.tok_emb.weight = assign(model.tok_emb.weight, params[\"model.embed_tokens.weight\"], \"model.embed_tokens.weight\")\n", "\n", " for l in range(param_config[\"n_layers\"]):\n", " block = model.trf_blocks[l]\n", " att = block.att\n", "\n", " # Q, K, V projections\n", " att.W_query.weight = assign(\n", " att.W_query.weight,\n", " params[f\"model.layers.{l}.self_attn.q_proj.weight\"],\n", " f\"model.layers.{l}.self_attn.q_proj.weight\"\n", " )\n", " att.W_key.weight = assign(\n", " att.W_key.weight,\n", " params[f\"model.layers.{l}.self_attn.k_proj.weight\"],\n", " f\"model.layers.{l}.self_attn.k_proj.weight\"\n", " )\n", " att.W_value.weight = assign(\n", " att.W_value.weight,\n", " params[f\"model.layers.{l}.self_attn.v_proj.weight\"],\n", " f\"model.layers.{l}.self_attn.v_proj.weight\"\n", " )\n", "\n", " # Output projection\n", " att.out_proj.weight = assign(\n", " att.out_proj.weight,\n", " params[f\"model.layers.{l}.self_attn.o_proj.weight\"],\n", " f\"model.layers.{l}.self_attn.o_proj.weight\"\n", " )\n", "\n", " # QK norms\n", " if hasattr(att, \"q_norm\") and att.q_norm is not None:\n", " att.q_norm.scale = assign(\n", " att.q_norm.scale,\n", " params[f\"model.layers.{l}.self_attn.q_norm.weight\"],\n", " f\"model.layers.{l}.self_attn.q_norm.weight\"\n", " )\n", " if hasattr(att, \"k_norm\") and att.k_norm is not None:\n", " att.k_norm.scale = assign(\n", " att.k_norm.scale,\n", " params[f\"model.layers.{l}.self_attn.k_norm.weight\"],\n", " f\"model.layers.{l}.self_attn.k_norm.weight\"\n", " )\n", "\n", " # Attention layernorm\n", " block.norm1.scale = assign(\n", " block.norm1.scale,\n", " params[f\"model.layers.{l}.input_layernorm.weight\"],\n", " f\"model.layers.{l}.input_layernorm.weight\"\n", " )\n", "\n", " # Feedforward weights\n", " if \"num_experts\" in param_config and param_config[\"num_experts\"] > 0:\n", " # Load router (gating) weights\n", " block.ff.gate.weight = assign(\n", " block.ff.gate.weight,\n", " params[f\"model.layers.{l}.mlp.gate.weight\"],\n", " f\"model.layers.{l}.mlp.gate.weight\"\n", " )\n", " # Load expert weights\n", " for e in range(param_config[\"num_experts\"]):\n", " prefix = f\"model.layers.{l}.mlp.experts.{e}\"\n", " block.ff.fc1[e].weight = assign(\n", " block.ff.fc1[e].weight,\n", " params[f\"{prefix}.gate_proj.weight\"],\n", " f\"{prefix}.gate_proj.weight\"\n", " )\n", " block.ff.fc2[e].weight = assign(\n", " block.ff.fc2[e].weight,\n", " params[f\"{prefix}.up_proj.weight\"],\n", " f\"{prefix}.up_proj.weight\"\n", " )\n", " block.ff.fc3[e].weight = assign(\n", " block.ff.fc3[e].weight,\n", " params[f\"{prefix}.down_proj.weight\"],\n", " f\"{prefix}.down_proj.weight\"\n", " )\n", "\n", " else:\n", " block.ff.fc1.weight = assign(\n", " block.ff.fc1.weight,\n", " params[f\"model.layers.{l}.mlp.gate_proj.weight\"],\n", " f\"model.layers.{l}.mlp.gate_proj.weight\"\n", " )\n", " block.ff.fc2.weight = assign(\n", " block.ff.fc2.weight,\n", " params[f\"model.layers.{l}.mlp.up_proj.weight\"],\n", " f\"model.layers.{l}.mlp.up_proj.weight\"\n", " )\n", " block.ff.fc3.weight = assign(\n", " block.ff.fc3.weight,\n", " params[f\"model.layers.{l}.mlp.down_proj.weight\"],\n", " f\"model.layers.{l}.mlp.down_proj.weight\"\n", " )\n", "\n", " block.norm2.scale = assign(\n", " block.norm2.scale,\n", " params[f\"model.layers.{l}.post_attention_layernorm.weight\"],\n", " f\"model.layers.{l}.post_attention_layernorm.weight\"\n", " )\n", "\n", " # Final normalization and output head\n", " model.final_norm.scale = assign(model.final_norm.scale, params[\"model.norm.weight\"], \"model.norm.weight\")\n", "\n", " if \"lm_head.weight\" in params:\n", " model.out_head.weight = assign(model.out_head.weight, params[\"lm_head.weight\"], \"lm_head.weight\")\n", " else:\n", " model.out_head.weight = model.tok_emb.weight\n", " print(\"Model uses weight tying.\")" ] }, { "cell_type": "code", "execution_count": 16, "id": "699cb1b8-a67d-49fb-80a6-0dad9d81f392", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 17, "referenced_widgets": [ "9881b6995c3f49dc89e6992fd9ab660b", "17a3174e65c54476b2e0d1faf8f011ca", "1bbf2e62c0754d1593beb4105a7f1ac1", "b82112e1dec645d98aa1c1ba64abcb61", "271e2bd6a35e4a8b92de8697f7c0be5f", "90a79523187446dfa692723b2e5833a7", "431ffb83b8c14bf182f0430e07ea6154", "a8f1b72a33dd4b548de23fbd95e0da18", "25cc36132d384189acfbecc59483134b", "bfd06423ad544218968648016e731a46", "d029630b63ff44cf807ade428d2eb421" ] }, "id": "699cb1b8-a67d-49fb-80a6-0dad9d81f392", "outputId": "55b2f28c-142f-4698-9d23-d27456d3ed6d" }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "acf19bb84d754884821e1794cedb25a4", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Fetching 28 files: 0%| | 0/28 [00:00\",\n", " \"<|im_start|>\", \"<|im_end|>\",\n", " \"<|object_ref_start|>\", \"<|object_ref_end|>\",\n", " \"<|box_start|>\", \"<|box_end|>\",\n", " \"<|quad_start|>\", \"<|quad_end|>\",\n", " \"<|vision_start|>\", \"<|vision_end|>\",\n", " \"<|vision_pad|>\", \"<|image_pad|>\", \"<|video_pad|>\",\n", " \"\", \"\"\n", " ]\n", " _SPLIT_RE = re.compile(r\"(<\\|[^>]+?\\|>||)\")\n", "\n", " def __init__(self, tokenizer_file_path=\"tokenizer.json\", repo_id=None,\n", " apply_chat_template=True, add_generation_prompt=False, add_thinking=False):\n", "\n", " self.apply_chat_template = apply_chat_template\n", " self.add_generation_prompt = add_generation_prompt\n", " self.add_thinking = add_thinking\n", "\n", " tok_file = Path(tokenizer_file_path)\n", " self._tok = Tokenizer.from_file(str(tok_file))\n", " self._special_to_id = {}\n", " for t in self._SPECIALS:\n", " tid = self._tok.token_to_id(t)\n", " if tid is not None:\n", " self._special_to_id[t] = tid\n", "\n", " self.pad_token_id = self._special_to_id[\"<|endoftext|>\"]\n", " self.eos_token_id = self.pad_token_id\n", "\n", " if repo_id and \"Base\" not in repo_id:\n", " eos_token = \"<|im_end|>\"\n", " else:\n", " eos_token = \"<|endoftext|>\"\n", " if eos_token in self._special_to_id:\n", " self.eos_token_id = self._special_to_id[eos_token]\n", "\n", " def encode(self, text, chat_wrapped=None):\n", " if chat_wrapped is None:\n", " chat_wrapped = self.apply_chat_template\n", "\n", " stripped = text.strip()\n", " if stripped in self._special_to_id and \"\\n\" not in stripped:\n", " return [self._special_to_id[stripped]]\n", "\n", " if chat_wrapped:\n", " text = self._wrap_chat(text)\n", "\n", " ids = []\n", " for part in filter(None, self._SPLIT_RE.split(text)):\n", " if part in self._special_to_id:\n", " ids.append(self._special_to_id[part])\n", " else:\n", " ids.extend(self._tok.encode(part).ids)\n", " return ids\n", "\n", " def decode(self, ids):\n", " return self._tok.decode(ids, skip_special_tokens=False)\n", "\n", " def _wrap_chat(self, user_msg):\n", " s = f\"<|im_start|>user\\n{user_msg}<|im_end|>\\n\"\n", " if self.add_generation_prompt:\n", " s += \"<|im_start|>assistant\"\n", " if self.add_thinking:\n", " s += \"\\n\"\n", " else:\n", " s += \"\\n\\n\\n\\n\\n\"\n", " return s" ] }, { "cell_type": "code", "execution_count": 18, "id": "7b6df8bc-7308-468e-93ce-2d5529ea7866", "metadata": {}, "outputs": [], "source": [ "tokenizer_file_path = f\"{Path(repo_id).parts[-1]}/tokenizer.json\"\n", "\n", "tokenizer = Qwen3Tokenizer(\n", " tokenizer_file_path=tokenizer_file_path,\n", " repo_id=repo_id,\n", " apply_chat_template=True,\n", " add_generation_prompt=True,\n", " add_thinking=True\n", ")" ] }, { "cell_type": "code", "execution_count": 19, "id": "1946b534-e3af-431a-a222-391a60bfa892", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "'<|im_start|>user\\nImplement a binary search function in Python<|im_end|>\\n<|im_start|>assistant\\n'" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# prompt = \"Give me a short introduction to large language models.\"\n", "prompt = \"Implement a binary search function in Python\"\n", "\n", "\n", "input_token_ids = tokenizer.encode(prompt)\n", "text = tokenizer.decode(input_token_ids)\n", "text" ] }, { "cell_type": "markdown", "id": "57d07df1-4401-4792-b549-7c4cc5632323", "metadata": { "id": "57d07df1-4401-4792-b549-7c4cc5632323" }, "source": [ " \n", "# 5. Generate text" ] }, { "cell_type": "code", "execution_count": 20, "id": "60b9fc72", "metadata": {}, "outputs": [], "source": [ "def generate_text_basic_stream(model, token_ids, max_new_tokens, eos_token_id=None):\n", "\n", " model.eval()\n", " with torch.no_grad():\n", " for _ in range(max_new_tokens):\n", " out = model(token_ids)[:, -1]\n", " next_token = torch.argmax(out, dim=-1, keepdim=True)\n", "\n", " if (eos_token_id is not None\n", " and torch.all(next_token == eos_token_id)):\n", " break\n", "\n", " yield next_token\n", " \n", " token_ids = torch.cat([token_ids, next_token], dim=1)" ] }, { "cell_type": "code", "execution_count": 21, "id": "a5b30753", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Here's a comprehensive implementation of binary search in Python with both iterative and recursive approaches:\n", "\n", "## Iterative Binary Search\n", "\n", "```python\n", "def binary_search(arr, target):\n", " \"\"\"\n", " Iterative binary search implementation\n", " \n", " Args:\n", " arr: Sorted list of elements\n", " target: Element to search for\n", " \n", " Returns:\n", " int: Index of target if found, -1 if not found\n", " \"\"\"\n", " left = 0\n", " right = len(arr) - 1\n", " \n", " while left" ] } ], "source": [ "input_token_ids_tensor = torch.tensor(input_token_ids, device=device).unsqueeze(0)\n", "\n", "\n", "for token in generate_text_basic_stream(\n", " model=model,\n", " token_ids=input_token_ids_tensor,\n", " max_new_tokens=100, # Cut-off after 100 tokens because non-kv variant is very slow\n", " # eos_token_id=tokenizer.eos_token_id\n", "):\n", " token_id = token.squeeze(0).tolist()\n", " print(\n", " tokenizer.decode(token_id),\n", " end=\"\",\n", " flush=True\n", " )" ] }, { "cell_type": "markdown", "id": "549324d6-5c71-4147-ae21-2e67675faa3d", "metadata": { "id": "549324d6-5c71-4147-ae21-2e67675faa3d" }, "source": [ " \n", "# What's next?" ] }, { "cell_type": "markdown", "id": "e6edaaae-2de1-406c-8ffa-897cdfa3808c", "metadata": { "id": "e6edaaae-2de1-406c-8ffa-897cdfa3808c" }, "source": [ "- Check out the [README.md](./README.md), to use this model via the `llms_from_scratch` package\n", "- For those interested in a comprehensive guide on building a large language model from scratch and gaining a deeper understanding of its mechanics, you might like my [Build a Large Language Model (From Scratch)](http://mng.bz/orYv)\n", "\n", "" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "A100", "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.13.5" } }, "nbformat": 4, "nbformat_minor": 5 }