Add industrial embedding training practices section to notebook 03

Covers weight tying, parameter grouping for AdamW, initialization,
and a simulated training step with gradient clipping.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: sanbuphy

notebooks/part1-foundation/03-embedding-position.ipynb

@@ -429,12 +429,148 @@
    "source": "## 5. 向量空间的可视化\n\n前面提到训练后相似词的向量会靠近，下面用一个小实验来验证：把高维 embedding 用 t-SNE 压到 2D，对比未训练和训练后的分布。"
   },
   {
-   "cell_type": "code",
-   "source": "# === 构建词表：14 个词，分 5 个语义组 + 1 个特殊 token ===\nwords = [\n    \"cat\", \"dog\", \"bird\",          # 动物\n    \"sat\", \"ran\", \"flew\",          # 动作\n    \"on\", \"under\",                 # 介词\n    \"mat\", \"rug\", \"tree\", \"sky\",   # 物品\n    \"the\", \"a\",                    # 冠词\n    \"[PAD]\",                       # 特殊 token\n]\nviz_vocab_size = len(words)\nviz_d_model = 16\n\nword_groups = {\n    \"cat\": \"动物\", \"dog\": \"动物\", \"bird\": \"动物\",\n    \"sat\": \"动作\", \"ran\": \"动作\", \"flew\": \"动作\",\n    \"on\": \"介词\", \"under\": \"介词\",\n    \"mat\": \"物品\", \"rug\": \"物品\", \"tree\": \"物品\", \"sky\": \"物品\",\n    \"the\": \"冠词\", \"a\": \"冠词\",\n    \"[PAD]\": \"特殊\",\n}\n\n# === 未训练的 Embedding（随机初始化）===\ntorch.manual_seed(42)\nuntrained_emb = nn.Embedding(viz_vocab_size, viz_d_model)\n\n# === 模拟「已训练」的 Embedding ===\n# 同组词共享基础向量 + 小噪声 → 同组词在空间中自然聚在一起\ntorch.manual_seed(999)\ngroup_bases = {g: torch.randn(viz_d_model) * 0.6\n               for g in [\"动物\", \"动作\", \"介词\", \"物品\", \"冠词\"]}\n\ntrained_emb = nn.Embedding(viz_vocab_size, viz_d_model)\nwith torch.no_grad():\n    for i, w in enumerate(words):\n        group = word_groups[w]\n        if group == \"特殊\":\n            trained_emb.weight[i] = torch.tensor([10.0, 10.0] + [0.0] * (viz_d_model - 2))\n        else:\n            trained_emb.weight[i] = group_bases[group] + 0.12 * torch.randn(viz_d_model)\n\n# === t-SNE 降维到 2D ===\nall_emb = torch.cat([untrained_emb.weight.data, trained_emb.weight.data], dim=0)\n\ntry:\n    from sklearn.manifold import TSNE\n    tsne = TSNE(n_components=2, random_state=42, perplexity=5)\n    all_2d = tsne.fit_transform(all_emb.numpy())\nexcept ImportError:\n    from sklearn.decomposition import PCA\n    pca = PCA(n_components=2)\n    all_2d = pca.fit_transform(all_emb.numpy())\n\nn = viz_vocab_size\nuntrained_2d = all_2d[:n]\ntrained_2d = all_2d[n:]\n\n# === 画图：左右对比 ===\ngroup_colors = {\n    \"动物\": \"#e74c3c\", \"动作\": \"#2ecc71\", \"介词\": \"#3498db\",\n    \"物品\": \"#f39c12\", \"冠词\": \"#9b59b6\", \"特殊\": \"#95a5a6\",\n}\ncolors = [group_colors[word_groups[w]] for w in words]\n\nfig, axes = plt.subplots(1, 2, figsize=(16, 7))\n\naxes[0].scatter(untrained_2d[:, 0], untrained_2d[:, 1], c=colors, s=200,\n                edgecolors='black', linewidth=1.2, alpha=0.85, zorder=5)\nfor i, w in enumerate(words):\n    axes[0].annotate(w, (untrained_2d[i, 0], untrained_2d[i, 1]),\n                     fontsize=10, textcoords=\"offset points\", xytext=(6, 8 if w != \"[PAD]\" else -15))\naxes[0].set_title(\"Untrained embeddings (t-SNE)\\nRandom init, no clear structure\", fontsize=14, fontweight='bold')\naxes[0].set_xlabel(\"t-SNE dim 1\"); axes[0].set_ylabel(\"t-SNE dim 2\")\naxes[0].grid(True, alpha=0.15)\n\naxes[1].scatter(trained_2d[:, 0], trained_2d[:, 1], c=colors, s=200,\n                edgecolors='black', linewidth=1.2, alpha=0.85, zorder=5)\nfor i, w in enumerate(words):\n    axes[1].annotate(w, (trained_2d[i, 0], trained_2d[i, 1]),\n                     fontsize=10, textcoords=\"offset points\", xytext=(6, 8 if w != \"[PAD]\" else -15))\naxes[1].set_title(\"Trained embeddings (t-SNE)\\nSimilar groups cluster together\", fontsize=14, fontweight='bold')\naxes[1].set_xlabel(\"t-SNE dim 1\"); axes[1].set_ylabel(\"t-SNE dim 2\")\naxes[1].grid(True, alpha=0.15)\n\nfrom matplotlib.patches import Patch\naxes[1].legend(\n    handles=[Patch(facecolor=c, edgecolor='black', label=g) for g, c in group_colors.items()],\n    loc='lower right', fontsize=10, title=\"Semantic group\", title_fontsize=11)\n\nplt.tight_layout()\nplt.show()\n\nprint(\"关键观察：\")\nprint(\"  左图：cat/dog/bird 散落在不同角落，没有结构\")\nprint(\"  右图：cat/dog/bird（红）聚成一团，sat/ran/flew（绿）聚在另一块\")",
+   "cell_type": "markdown",
+   "source": "## 6. 工业界的 Embedding 训练实践\n\n前面的 Embedding 是一个 `nn.Embedding` 查表，概念是对的。但在训练真实的大语言模型时，工业界有几个关键的工程实践，直接影响参数效率和训练稳定性。\n\n**权重共享（Weight Tying）**\n\n现代 LLM 通常将输入 Embedding 和输出投影层共享同一组权重。Token Embedding 的矩阵形状是 `[vocab_size, d_model]`，模型最后的输出层 `lm_head` 也是把 d_model 维的向量映射回 vocab_size 维——前者把 token ID 变成向量，后者把向量变回 token ID 的概率分布。如果这两个矩阵是同一块内存，就能节省 `vocab_size × d_model` 个参数。\n\n以 LLaMA 7B 为例：vocab_size=32000，d_model=4096，不共享要多约 1.3 亿参数。GPT-2、GPT-3、LLaMA、DeepSeek 等都用了这个技巧。\n\n**Embedding 不做权重衰减**\n\nL2 正则化（weight decay）限制权重的绝对值，防止过拟合。但 Embedding 的每一行代表一个 token 的语义向量——它需要在训练中自由移动到合适的位置。施加 L2 惩罚等于把所有向量往原点拽，干扰语义学习。业界惯例：Embedding、LayerNorm 的 weight/bias、所有 bias 项都不参与 weight decay。\n\n**初始化：更小的标准差**\n\n`nn.Embedding` 默认用 `N(0,1)` 初始化。实践中常用更小的标准差——例如 `N(0, 1/√d_model)`——避免初始嵌入值过大，导致早期训练梯度不稳定。\n\n**混合精度下的 Embedding**\n\n用 FP16/BF16 训练时，Embedding 矩阵通常保留一份 FP32 副本（master weights），前向时转成低精度，反向更新时在 FP32 下累积梯度。这是为了避免 FP16 动态范围不足导致的梯度下溢。\n\n下面用代码把这几条串起来，模拟一个真实训练循环中 Embedding 的配置方式。</cell id=\"cell-21\">\n<｜｜DSML｜｜parameter name=\"edit_mode\" string=\"true\">insert",
    "metadata": {},
    "execution_count": null,
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# === 工业界 Embedding 训练实践：权重共享 + 参数分组 + 初始化 ===\n",
+    "import math\n",
+    "\n",
+    "# ---------- 1. 权重共享（Weight Tying）----------\n",
+    "# 输入 Embedding 和输出 lm_head 共享权重——这是 GPT/LLaMA 的标准做法\n",
+    "vocab_size, d_model = 10000, 512\n",
+    "\n",
+    "class GPTStyleModel(nn.Module):\n",
+    "    def __init__(self, vocab_size, d_model):\n",
+    "        super().__init__()\n",
+    "        self.wte = nn.Embedding(vocab_size, d_model)   # token → vector\n",
+    "        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)  # vector → token probs\n",
+    "        # ★ 关键：让 lm_head 复用 wte 的权重矩阵\n",
+    "        self.lm_head.weight = self.wte.weight\n",
+    "\n",
+    "    def forward(self, input_ids):\n",
+    "        x = self.wte(input_ids)\n",
+    "        return self.lm_head(x)\n",
+    "\n",
+    "model = GPTStyleModel(vocab_size, d_model)\n",
+    "\n",
+    "# 验证：两个 weight 指向同一块内存\n",
+    "print(\"=\" * 50)\n",
+    "print(\"1. 权重共享验证\")\n",
+    "print(f\"   wte.weight 内存地址:       {model.wte.weight.data_ptr()}\")\n",
+    "print(f\"   lm_head.weight 内存地址:   {model.lm_head.weight.data_ptr()}\")\n",
+    "print(f\"   是同一块内存:              {model.wte.weight.data_ptr() == model.lm_head.weight.data_ptr()}\")\n",
+    "print(f\"   节省参数量:                {vocab_size * d_model:,} → {vocab_size * d_model * 4 / 1e6:.1f} MB (fp32)\")\n",
+    "\n",
+    "# 反向传播验证：同一块内存上的梯度正确累加\n",
+    "input_ids = torch.randint(0, vocab_size, (2, 16))\n",
+    "logits = model(input_ids)\n",
+    "logits.mean().backward()\n",
+    "print(f\"   两个 grad 指向同一内存:    {model.wte.weight.grad.data_ptr() == model.lm_head.weight.grad.data_ptr()}\")\n",
+    "\n",
+    "# ---------- 2. 不共享 vs 共享：参数量对比 ----------\n",
+    "class NoTieModel(nn.Module):\n",
+    "    def __init__(self, vocab_size, d_model):\n",
+    "        super().__init__()\n",
+    "        self.wte = nn.Embedding(vocab_size, d_model)\n",
+    "        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)\n",
+    "        # 没有 weight tying，两个矩阵各自独立\n",
+    "\n",
+    "    def forward(self, input_ids):\n",
+    "        return self.lm_head(self.wte(input_ids))\n",
+    "\n",
+    "tie_params = sum(p.numel() for p in GPTStyleModel(vocab_size, d_model).parameters())\n",
+    "no_tie_params = sum(p.numel() for p in NoTieModel(vocab_size, d_model).parameters())\n",
+    "print(f\"\\n2. 参数量对比 (vocab={vocab_size:,}, d_model={d_model}):\")\n",
+    "print(f\"   共享权重: {tie_params:,}\")\n",
+    "print(f\"   独立权重: {no_tie_params:,}\")\n",
+    "print(f\"   节省:     {no_tie_params - tie_params:,} ({100 * (no_tie_params - tie_params) / no_tie_params:.1f}%)\")\n",
+    "\n",
+    "# 类比真实模型尺寸\n",
+    "for name, vs, dm in [(\"GPT-2 small\", 50257, 768), (\"LLaMA 7B\", 32000, 4096)]:\n",
+    "    saved = vs * dm\n",
+    "    print(f\"   {name}: vocab={vs:,}, d_model={dm} → 节省 {saved:,} 参数 ({saved * 4 / 1e6:.1f} MB)\")\n",
+    "\n",
+    "# ---------- 3. Embedding 不做 weight decay ----------\n",
+    "# LLaMA / GPT 训练的标准参数分组：某些参数不做正则化\n",
+    "def make_param_groups(model, weight_decay=0.1):\n",
+    "    \"\"\"把参数分成两组：一组做 decay，一组不做\"\"\"\n",
+    "    decay_params, no_decay_params = [], []\n",
+    "    no_decay_keywords = ['wte', 'norm', 'bias']  # Embedding / Norm / Bias 不做衰减\n",
+    "    \n",
+    "    for name, param in model.named_parameters():\n",
+    "        if not param.requires_grad:\n",
+    "            continue\n",
+    "        if any(kw in name.lower() for kw in no_decay_keywords):\n",
+    "            no_decay_params.append(param)\n",
+    "        else:\n",
+    "            decay_params.append(param)\n",
+    "    \n",
+    "    return [\n",
+    "        {'params': decay_params, 'weight_decay': weight_decay},\n",
+    "        {'params': no_decay_params, 'weight_decay': 0.0},\n",
+    "    ]\n",
+    "\n",
+    "model2 = GPTStyleModel(vocab_size, d_model)\n",
+    "param_groups = make_param_groups(model2, weight_decay=0.1)\n",
+    "\n",
+    "decay_count = sum(p.numel() for p in param_groups[0]['params'])\n",
+    "no_decay_count = sum(p.numel() for p in param_groups[1]['params'])\n",
+    "print(f\"\\n3. 参数分组（weight_decay=0.1）:\")\n",
+    "print(f\"   做衰减的参数:   {decay_count:,}\")\n",
+    "print(f\"   不做衰减的参数: {no_decay_count:,}  ← Embedding 在这里\")\n",
+    "print(f\"   Embedding 在 no_decay 组: {any('wte' in n for n, _ in model2.named_parameters())}\")\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(param_groups, lr=1e-3)\n",
+    "for i, pg in enumerate(optimizer.param_groups):\n",
+    "    wd = pg['weight_decay']\n",
+    "    count = sum(p.numel() for p in pg['params'])\n",
+    "    print(f\"   Group {i}: {count:,} params, weight_decay={wd}\")\n",
+    "\n",
+    "# ---------- 4. Embedding 初始化 ----------\n",
+    "print(f\"\\n4. 初始化对比:\")\n",
+    "torch.manual_seed(42)\n",
+    "default_emb = nn.Embedding(100, 16)\n",
+    "custom_emb = nn.Embedding(100, 16)\n",
+    "nn.init.normal_(custom_emb.weight, mean=0.0, std=1.0 / math.sqrt(d_model))\n",
+    "\n",
+    "print(f\"   nn.Embedding 默认 N(0,1):       std = {default_emb.weight.std().item():.4f}\")\n",
+    "print(f\"   自定义 N(0, 1/√d_model):         std = {custom_emb.weight.std().item():.4f}\")\n",
+    "print(f\"   目标 1/√{d_model} = {1.0 / math.sqrt(d_model):.4f}\")\n",
+    "print(f\"   → 更小的初始值 → 训练早期梯度更稳定\")\n",
+    "\n",
+    "# 模拟一个 mini-step：走一遍完整的训练步骤\n",
+    "print(f\"\\n5. 模拟一次训练 step（含梯度裁剪）:\")\n",
+    "model3 = GPTStyleModel(vocab_size, d_model)\n",
+    "optimizer = torch.optim.AdamW(\n",
+    "    make_param_groups(model3, weight_decay=0.1), lr=1e-3\n",
+    ")\n",
+    "\n",
+    "batch = torch.randint(0, vocab_size, (4, 32))   # 随机 batch\n",
+    "targets = torch.randint(0, vocab_size, (4, 32))  # 随机 targets\n",
+    "\n",
+    "logits = model3(batch)                            # 前向\n",
+    "loss = nn.functional.cross_entropy(\n",
+    "    logits.view(-1, vocab_size), targets.view(-1)\n",
+    ")\n",
+    "loss.backward()                                   # 反向\n",
+    "torch.nn.utils.clip_grad_norm_(model3.parameters(), 1.0)  # 梯度裁剪\n",
+    "optimizer.step()                                  # 更新参数\n",
+    "optimizer.zero_grad()\n",
+    "\n",
+    "print(f\"   loss: {loss.item():.4f}\")\n",
+    "print(f\"   wte.weight 已更新: {model3.wte.weight.grad is None} → 梯度已清零\")\n",
+    "print(f\"   → 完整流程：forward → loss → backward → clip → step → zero_grad\")\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "markdown",
    "id": "summary-section",