books/LLMs-from-scratch
1
0
Fork 0
mirror of https://github.com/rasbt/LLMs-from-scratch.git synced 2026-04-10 12:33:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

Optional weight tying for Qwen3 and Llama3.2 pretraining (#949)

Browse Source 
* optional weight tying for Qwen3 and Llama3.2

* typo
This commit is contained in:
casinca
2026-01-14 16:07:04 +01:00
committed by GitHub
parent e0dbec3331
commit 9c4be478f8
7 changed files with 17 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ch05/07_gpt_to_llama/standalone-llama32.ipynb
View File

@@ -358,7 +358,7 @@
" self.final_norm = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n", " self.final_norm = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n",
" self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
"\n", "\n",
" # Reusuable utilities\n", " # Reusable utilities\n",
" cos, sin = compute_rope_params(\n", " cos, sin = compute_rope_params(\n",
" head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n", " head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n",
" theta_base=cfg[\"rope_base\"],\n", " theta_base=cfg[\"rope_base\"],\n",

2
ch05/11_qwen3/standalone-qwen3-moe-plus-kvcache.ipynb
View File

@@ -432,7 +432,7 @@
" self.final_norm = RMSNorm(cfg[\"emb_dim\"])\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", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
"\n", "\n",
" # Reusuable utilities\n", " # Reusable utilities\n",
" if cfg[\"head_dim\"] is None:\n", " if cfg[\"head_dim\"] is None:\n",
" head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n", " head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
" else:\n", " else:\n",

2
ch05/11_qwen3/standalone-qwen3-moe.ipynb
View File

@@ -422,7 +422,7 @@
" self.final_norm = RMSNorm(cfg[\"emb_dim\"])\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", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
"\n", "\n",
" # Reusuable utilities\n", " # Reusable utilities\n",
" if cfg[\"head_dim\"] is None:\n", " if cfg[\"head_dim\"] is None:\n",
" head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n", " head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
" else:\n", " else:\n",

2
ch05/11_qwen3/standalone-qwen3.ipynb
View File

@@ -388,7 +388,7 @@
" self.final_norm = RMSNorm(cfg[\"emb_dim\"])\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", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
"\n", "\n",
" # Reusuable utilities\n", " # Reusable utilities\n",
" if cfg[\"head_dim\"] is None:\n", " if cfg[\"head_dim\"] is None:\n",
" head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n", " head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
" else:\n", " else:\n",

8
ch05/14_ch05_with_other_llms/ch05-llama32.ipynb
View File

@@ -113,7 +113,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 3, "execution_count": null,
"id": "86000d74-624a-48f0-86da-f41926cb9e04", "id": "86000d74-624a-48f0-86da-f41926cb9e04",
"metadata": { "metadata": {
"colab": { "colab": {
@@ -329,7 +329,11 @@
" self.final_norm = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n", " self.final_norm = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n",
" self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
"\n", "\n",
" # Reusuable utilities\n", " # Uncomment the following code to tie weights\n",
" # self.out_head.weight = self.tok_emb.weight\n",
" # torch.nn.init.normal_(self.out_head.weight, mean=0.0, std=0.02)\n",
"\n",
" # Reusable utilities\n",
" cos, sin = compute_rope_params(\n", " cos, sin = compute_rope_params(\n",
" head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n", " head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n",
" theta_base=cfg[\"rope_base\"],\n", " theta_base=cfg[\"rope_base\"],\n",

8
ch05/14_ch05_with_other_llms/ch05-qwen3.ipynb
View File

@@ -121,7 +121,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 3, "execution_count": null,
"id": "86000d74-624a-48f0-86da-f41926cb9e04", "id": "86000d74-624a-48f0-86da-f41926cb9e04",
"metadata": { "metadata": {
"colab": { "colab": {
@@ -332,7 +332,11 @@
" self.final_norm = RMSNorm(cfg[\"emb_dim\"])\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", " self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
"\n", "\n",
" # Reusuable utilities\n", " # Uncomment the following code to tie weights\n",
" # self.out_head.weight = self.tok_emb.weight\n",
" # torch.nn.init.normal_(self.out_head.weight, mean=0.0, std=0.02)\n",
"\n",
" # Reusable utilities\n",
" if cfg[\"head_dim\"] is None:\n", " if cfg[\"head_dim\"] is None:\n",
" head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n", " head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
" else:\n", " else:\n",

2
pkg/llms_from_scratch/kv_cache/llama3.py
View File

@@ -65,7 +65,7 @@ class Llama3Model(nn.Module):
self.final_norm = nn.RMSNorm(cfg["emb_dim"], eps=1e-5, dtype=cfg["dtype"]) self.final_norm = nn.RMSNorm(cfg["emb_dim"], eps=1e-5, dtype=cfg["dtype"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"]) self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"])
# Reusuable utilities # Reusable utilities
cos, sin = compute_rope_params( cos, sin = compute_rope_params(
head_dim=cfg["emb_dim"] // cfg["n_heads"], head_dim=cfg["emb_dim"] // cfg["n_heads"],
theta_base=cfg["rope_base"], theta_base=cfg["rope_base"],