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

* optional weight tying for Qwen3 and Llama3.2

* typo

Author: casinca

ch05/07_gpt_to_llama/standalone-llama32.ipynb

@@ -358,7 +358,7 @@
     "        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",
     "\n",
-    "        # Reusuable utilities\n",
+    "        # Reusable utilities\n",
     "        cos, sin = compute_rope_params(\n",
     "            head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n",
     "            theta_base=cfg[\"rope_base\"],\n",

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

@@ -432,7 +432,7 @@
     "        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",
+    "        # Reusable utilities\n",
     "        if cfg[\"head_dim\"] is None:\n",
     "            head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
     "        else:\n",

ch05/11_qwen3/standalone-qwen3-moe.ipynb

@@ -422,7 +422,7 @@
     "        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",
+    "        # Reusable utilities\n",
     "        if cfg[\"head_dim\"] is None:\n",
     "            head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
     "        else:\n",

ch05/11_qwen3/standalone-qwen3.ipynb

@@ -388,7 +388,7 @@
     "        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",
+    "        # Reusable utilities\n",
     "        if cfg[\"head_dim\"] is None:\n",
     "            head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
     "        else:\n",

ch05/14_ch05_with_other_llms/ch05-llama32.ipynb

@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "86000d74-624a-48f0-86da-f41926cb9e04",
    "metadata": {
     "colab": {
@@ -329,7 +329,11 @@
     "        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",
     "\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",
     "            head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n",
     "            theta_base=cfg[\"rope_base\"],\n",

ch05/14_ch05_with_other_llms/ch05-qwen3.ipynb

@@ -121,7 +121,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "86000d74-624a-48f0-86da-f41926cb9e04",
    "metadata": {
     "colab": {
@@ -332,7 +332,11 @@
     "        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",
+    "        # 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",
     "            head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
     "        else:\n",

pkg/llms_from_scratch/kv_cache/llama3.py

@@ -65,7 +65,7 @@ def __init__(self, cfg):
         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"])
 
-        # Reusuable utilities
+        # Reusable utilities
         cos, sin = compute_rope_params(
             head_dim=cfg["emb_dim"] // cfg["n_heads"],
             theta_base=cfg["rope_base"],