Fixing nits

Signed-off-by: Vladimir Suvorov <[email protected]>

Author: mydatascience

src/MaxText/examples/rl_llama3_demo.ipynb

@@ -87,18 +87,24 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
-   "source": []
+   "outputs": [],
+   "source": [
+    "#Choose the loss algorithm between GSPO or GRPO\n",
+    "LOSS_ALGO=\"grpo\" #  or \"gspo-token\" if you want to use GSPO"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Configuration for GRPO training\n",
     "import os\n",
+    "import sys\n",
+    "from pathlib import Path\n",
     "import MaxText\n",
     "from huggingface_hub import login\n",
     "import jax\n",
@@ -110,7 +116,6 @@
     "MODEL_NAME = \"llama3.1-8b\"\n",
     "HF_REPO_ID = \"meta-llama/Llama-3.1-8B-Instruct\"\n",
     "CHAT_TEMPLATE_PATH = f\"{MAXTEXT_REPO_ROOT}/examples/chat_templates/gsm8k_rl.json\"\n",
-    "LOSS_ALGO=\"grpo\" #  or \"gspo-token\" if you want to use GSPO\n",
     "\n",
     "# Required: Set these before running\n",
     "MODEL_CHECKPOINT_PATH = \"\"  # Update this!\n",
@@ -151,11 +156,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Import required modules\n",
-    "import os\n",
-    "import sys\n",
-    "from pathlib import Path\n",
-    "\n",
     "# Add MaxText to Python path\n",
     "maxtext_path = Path(MAXTEXT_REPO_ROOT) \n",
     "sys.path.insert(0, str(maxtext_path))\n",

src/MaxText/examples/sft_llama3_demo.ipynb

@@ -136,14 +136,16 @@
    "source": [
     "## Set the model, checkpoint path and output directory\n",
     "MODEL_NAME = \"llama3.1-8b\"\n",
-    "# Case 1: Set `MODEL_CHECKPOINT_PATH` to the path (local or gs://) that already has `Llama3.1-8B-Instruct` model checkpoint\n",
-    "# Case 2: If you do not have the checkpoint, then do not update `MODEL_CHECKPOINT_PATH`\n",
-    "# and this colab will download the checkpoint from HF and store at `\"{MAXTEXT_REPO_ROOT}/llama_checkpoint\"`\n",
-    "MODEL_CHECKPOINT_PATH = f\"{MAXTEXT_REPO_ROOT}/llama_checkpoint\"\n",
+    "# set the path to the model checkpoint or leave empty to download from HuggingFace\n",
+    "MODEL_CHECKPOINT_PATH = \"\"\n",
+    "if not MODEL_CHECKPOINT_PATH:\n",
+    "   MODEL_CHECKPOINT_PATH = f\"{MAXTEXT_REPO_ROOT}/llama_checkpoint\"\n",
+    "   print(\"Model checkpoint will be downloaded from HuggingFace at: \",  MODEL_CHECKPOINT_PATH)\n",
+    "   print(\"Set MODEL_CHECKPOINT_PATH if you do not wish to download the checkpoint.\")\n",
     "\n",
-    "# This is the directory where the fine-tuned model will be saved\n",
-    "# You can change it to any path you want (local or gs://)\n",
-    "BASE_OUTPUT_DIRECTORY = \"/tmp/out/maxtext_llama3_8b\"\n"
+    "BASE_OUTPUT_DIRECTORY = \"\"\n",
+    "if not BASE_OUTPUT_DIRECTORY:\n",
+    "   print(\"Please set BASE_OUTPUT_DIRECTORY to store output/logs.\")\n"
    ]
   },
   {