Fix sft_qwen3_demo jupyter notebook

Author: SurbhiJainUSC

.github/workflows/run_jupyter_notebooks.yml

@@ -84,7 +84,7 @@ jobs:
 
           for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl}*.ipynb; do
             filename=$(basename "$notebook")
-            if [[ "$filename" == "sft_qwen3_demo.ipynb" || "$filename" == "sft_llama3_demo_gpu.ipynb" ]]; then
+            if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" ]]; then
               echo "Skipping $filename"
               continue
             fi

src/maxtext/examples/sft_qwen3_demo.ipynb

@@ -95,21 +95,6 @@
     "    IN_COLAB = False"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "try:\n",
-    "    import google.colab\n",
-    "    print(\"Running the notebook in Google Colab\")\n",
-    "    IN_COLAB = True\n",
-    "except ImportError:\n",
-    "    print(\"Running the notebook on JupyterLab\")\n",
-    "    IN_COLAB = False"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -168,6 +153,7 @@
     "from datetime import datetime\n",
     "from flax import nnx\n",
     "from huggingface_hub import login\n",
+    "from jax import numpy as jnp\n",
     "\n",
     "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
    ]
@@ -192,10 +178,9 @@
    },
    "outputs": [],
    "source": [
-    "try:\n",
-    "    from google.colab import userdata\n",
+    "if IN_COLAB:\n",
     "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
-    "except ImportError:\n",
+    "else:\n",
     "    HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
     "\n",
     "# If not found in the environment, prompt the user for input securely\n",
@@ -438,6 +423,7 @@
     "        rollout_vllm_hbm_utilization=0.8,\n",
     "        rollout_vllm_init_with_random_weights=True,\n",
     "        rollout_vllm_tpu_backend_type=\"jax\",\n",
+    "        data_type=jnp.bfloat16,\n",
     "    ),\n",
     ")"
    ]