update chap 4

Author: JasonJiaxiangLi

lectures/Chapter4/lr_mup.ipynb

@@ -5,12 +5,33 @@
    "id": "583a9274",
    "metadata": {},
    "source": [
-    "# Practical concerns on training stability\n",
+    "# Practical concerns on large-scale training\n",
     "\n",
     "So far we have witnessed quite some modern optimizers, are we ready to go ahead and apply them directly on training extreme large models? The answer is no. For example, we still don't know how to choose learning rate, how to make it efficient on multiple (sometimes hundred or thousands of) GPUs, etc. In this section we briefly discuss several techniques popular in practice. We will not be able to provide all techincal details but will introduce related concepts and references for further readings.\n",
     "\n",
     "## lr scheduler\n",
     "\n",
+    "Let's begin with the first topic: how to choose the most important hyperparameter, namely the learning rate. Let's take a look at the AdamW algorithm again.\n",
+    "\n",
+    "````{prf:algorithm} AdamW\n",
+    ":label: alg:adamw_again\n",
+    "\n",
+    "- **Input**: initial point $\\vw^0$, stepsize sequence $\\{ \\eta_t \\}_{t \\geq 0}$, max. no. of iterations $T$, momentum parameter $(\\beta_1,\\beta_2)$, weight decay parameter $\\lambda$.\n",
+    "\n",
+    "- For $t=0,1,2,..., T-1$,\n",
+    "```{math}\n",
+    "\\begin{split}\n",
+    "\\xi_{i_{j,t}} &\\sim {\\cal D} \\text{ (draw i.i.d. samples)} \\\\\n",
+    "\\vg^t &\\leftarrow \\frac{1}{B}\\sum_{j=1}^{B}\\nabla \\ell(\\vw^t, \\xi_{i_{j,t}}) \\\\\n",
+    "\\vm^t &\\leftarrow \\beta_1\\vm^{t-1} + (1-\\beta_1)\\vg^t \\\\\n",
+    "\\vv^t &\\leftarrow \\beta_2\\vv^{t-1} + (1-\\beta_2)\\vg^t\\odot \\vg^t \\\\\n",
+    "\\hat{\\vm}^t &\\gets \\frac{\\vm^t}{1-\\beta_1^t},\\ \\hat{\\vv}^t \\gets \\frac{\\vm^t}{1-\\beta_2^t} \\\\\n",
+    "\\vw^{t+1} &\\leftarrow \\vw^{t} - \\eta_t \\frac{\\hat{\\vm}^t}{\\sqrt{\\hat{\\vv}^t}+\\epsilon} - \\eta_t\\lambda \\vw^{t}\n",
+    "\\end{split}\n",
+    "```\n",
+    "- **Output**: last iterate $\\vw^T$, or the solution sequence $\\{ \\vw^t \\}_{t=1}^T$.\n",
+    "````\n",
+    "\n",
     "Constant schedule\n",
     "\n",
     "Linear and Cosine decay\n",