Merge pull request AI-Hypercomputer#4154 from AI-Hypercomputer:multimodal_notebook

PiperOrigin-RevId: 938717576

Author: Google-ML-Automation

.github/workflows/run_jupyter_notebooks.yml

@@ -109,10 +109,9 @@ jobs:
           # Run Hugging Face authentication
           hf auth login --token "$HF_TOKEN"
 
-          for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl,lora}*.ipynb; do
+          for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/*.ipynb; do
             filename=$(basename "$notebook")
-            # TODO: Update runnner to v6e-8 as RL with LLama3.1-8b doesn't fit on v6e-4
-            if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" || "$filename" == "maxtext_with_gepa.ipynb" ]]; then
+            if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" || "$filename" == "maxtext_with_gepa.ipynb" || "$filename" == "demo_decoding.ipynb" ]]; then
               echo "Skipping $filename"
               continue
             fi
@@ -123,6 +122,11 @@ jobs:
             echo "------------------------------------------------------"
 
             $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv
+
+            # Clean up any checkpoint directories created by the notebook to avoid filling up disk space
+            echo "Post-notebook disk cleanup for $filename ..."
+            rm -rf "$MAXTEXT_PKG_DIR"/*_output
+            rm -rf "$HOME/.cache/huggingface/hub"
           done
       - name: Upload Outputs
         if: always()

docs/tutorials/posttraining/multimodal.md

@@ -6,7 +6,7 @@ This document provides a guide to use the multimodal functionalities in MaxText
 - **Multimodal Decode**: Inference with text+images as input.
 - **Supervised Fine-Tuning (SFT)**: Apply SFT to the model using a visual-question-answering dataset.
 
-We also provide a [colab](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/examples/multimodal_gemma3_demo.ipynb) for multimodal features demonstration. The following table provides a list of models and modalities we currently support:
+We also provide a [colab](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb) for multimodal features demonstration. The following table provides a list of models and modalities we currently support:
 
 | Models                      | Input: Text | Input: Image | Input: Video | Input: Audio | Output |
 | :-------------------------- | :---------: | :----------: | :----------: | :----------: | :----: |

src/maxtext/examples/lora_llama3_demo.ipynb

@@ -215,16 +215,6 @@
    "outputs": [],
    "source": [
     "if not epath.Path(MODEL_CHECKPOINT_PATH).exists():\n",
-    "    # Install torch for the conversion script\n",
-    "    print(\"Installing torch...\")\n",
-    "    subprocess.run(\n",
-    "        [\n",
-    "            sys.executable, \"-m\", \"pip\", \"install\",\n",
-    "            \"torch\", \"--index-url\", \"https://download.pytorch.org/whl/cpu\"\n",
-    "        ],\n",
-    "        check=True\n",
-    "    )\n",
-    "\n",
     "    print(\"Converting checkpoint from HuggingFace...\")\n",
     "    env = os.environ.copy()\n",
     "    env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
@@ -607,4 +597,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

src/maxtext/examples/multimodal_gemma3_demo.ipynb

@@ -193,17 +193,6 @@
    "outputs": [],
    "source": [
     "if not epath.Path(MODEL_CHECKPOINT_PATH).exists():\n",
-    "    # Install torch for the conversion script\n",
-    "    print(\"Installing torch...\")\n",
-    "    subprocess.run(\n",
-    "        [\n",
-    "            sys.executable, \"-m\", \"pip\", \"install\",\n",
-    "            \"torch\", \"--index-url\", \"https://download.pytorch.org/whl/cpu\"\n",
-    "        ],\n",
-    "        check=True\n",
-    "    )\n",
-    "\n",
-    "    # Run checkpoint conversion with environment variables\n",
     "    print(\"Converting checkpoint from HuggingFace...\")\n",
     "    env = os.environ.copy()\n",
     "    env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
@@ -222,7 +211,7 @@
     "        check=True,\n",
     "        env=env\n",
     "    )\n",
-    "    \n",
+    "\n",
     "    MODEL_CHECKPOINT_PATH = os.path.join(MODEL_CHECKPOINT_PATH, \"0/items\")\n",
     "else:\n",
     "    print(f\"Model checkpoint exists at {MODEL_CHECKPOINT_PATH}\")"

src/maxtext/examples/rl_llama3_demo.ipynb

@@ -206,17 +206,6 @@
    "outputs": [],
    "source": [
     "if not epath.Path(MODEL_CHECKPOINT_PATH).exists():\n",
-    "    # Install torch for the conversion script\n",
-    "    print(\"Installing torch...\")\n",
-    "    subprocess.run(\n",
-    "        [\n",
-    "            sys.executable, \"-m\", \"pip\", \"install\",\n",
-    "            \"torch\", \"--index-url\", \"https://download.pytorch.org/whl/cpu\"\n",
-    "        ],\n",
-    "        check=True\n",
-    "    )\n",
-    "\n",
-    "    # Run checkpoint conversion with environment variables\n",
     "    print(\"Converting checkpoint from HuggingFace...\")\n",
     "    env = os.environ.copy()\n",
     "    env[\"JAX_PLATFORMS\"] = \"cpu\"\n",