Fix CI notebook execution and add SFT training step verification

.github/workflows/ci_pipeline.yml

@@ -168,7 +168,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        flavor: [tpu-post-training-unit]
+        flavor: [tpu-post-training-unit, tpu-post-training-integration]
     with:
       flavor: ${{ matrix.flavor }}
       base_image: maxtext-unit-test-tpu:py312
@@ -241,7 +241,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        flavor: [cpu-post-training-unit]
+        flavor: [cpu-post-training-unit, cpu-post-training-integration]
     with:
       flavor: ${{ matrix.flavor }}
       base_image: maxtext-unit-test-tpu:py312

.github/workflows/run_ci_tests.yml

@@ -71,7 +71,7 @@ jobs:
         flavor: >-
             ${{ fromJSON('{
               "gpu-pre-training": ["gpu-unit", "gpu-integration"],
-              "tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit"],
+              "tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit", "cpu-post-training-integration"],
               "tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit", "cpu-integration"]
             }')[format('{0}-{1}', inputs.device, inputs.workflow)] }}
     uses: ./.github/workflows/run_tests_coordinator.yml

.github/workflows/run_jupyter_notebooks.yml

@@ -34,20 +34,77 @@ on:
       maxtext_sha:
         required: false
         type: string
-      # Flag to skip source checkout and wheel installation
-      maxtext_installed:
-        required: false
-        type: boolean
-        default: false
     secrets:
       HF_TOKEN:
         required: true
 
 permissions:
   contents: read
 jobs:
+  discover_notebooks:
+    name: Discover Notebooks to Run
+    runs-on: ubuntu-latest
+    outputs:
+      notebooks: ${{ steps.list.outputs.notebooks }}
+      count: ${{ steps.list.outputs.count }}
+    steps:
+      - name: Checkout MaxText
+        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
+        with:
+          ref: ${{ inputs.maxtext_sha }}
+          fetch-depth: 0
+          persist-credentials: false
+      - name: Determine Notebook List
+        id: list
+        shell: bash
+        run: |
+          MAXTEXT_NOTEBOOKS_ROOT="src/maxtext/examples"
+          SKIPPED=("sft_llama3_demo_gpu.ipynb" "maxtext_with_gepa.ipynb" "demo_decoding.ipynb" "dpo_qwen3_demo.ipynb")
+
+          is_skipped() {
+            local name="$1"
+            for s in "${SKIPPED[@]}"; do
+              [[ "$name" == "$s" ]] && return 0
+            done
+            return 1
+          }
+
+          SELECTED=()
+          if [ "${GITHUB_EVENT_NAME}" == "pull_request" ] && [ -n "${GITHUB_BASE_REF}" ]; then
+            git fetch origin "${GITHUB_BASE_REF}" --depth=50 2>/dev/null || true
+            while IFS= read -r file; do
+              if [[ -n "$file" && -f "$file" ]]; then
+                name=$(basename "$file")
+                if ! is_skipped "$name"; then
+                  SELECTED+=("$name")
+                fi
+              fi
+            done < <(git diff --name-only "origin/${GITHUB_BASE_REF}...HEAD" -- "${MAXTEXT_NOTEBOOKS_ROOT}"/*.ipynb 2>/dev/null || true)
+          fi
+
+          # If not a PR or no specific notebooks were modified (e.g. scheduled run or workflow file modified), run all active notebooks
+          if [ ${#SELECTED[@]} -eq 0 ]; then
+            for nb in "$MAXTEXT_NOTEBOOKS_ROOT"/*.ipynb; do
+              name=$(basename "$nb")
+              if ! is_skipped "$name" && [[ -f "$nb" ]]; then
+                SELECTED+=("$name")
+              fi
+            done
+          fi
+
+          JSON_ARRAY=$(jq -nc '$ARGS.positional' --args "${SELECTED[@]}")
+          echo "Discovered notebooks to run: $JSON_ARRAY"
+          echo "notebooks=$JSON_ARRAY" >> "$GITHUB_OUTPUT"
+          echo "count=${#SELECTED[@]}" >> "$GITHUB_OUTPUT"
+
   run:
-    name: Execute Notebooks
+    name: Execute ${{ matrix.notebook }}
+    needs: [discover_notebooks]
+    if: needs.discover_notebooks.outputs.count > 0
+    strategy:
+      fail-fast: false
+      matrix:
+        notebook: ${{ fromJson(needs.discover_notebooks.outputs.notebooks) }}
     runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
     container:
       image: gcr.io/tpu-prod-env-multipod/${{ inputs.base_image }} # zizmor: ignore[unpinned-images]
@@ -56,18 +113,15 @@ jobs:
         UV_TORCH_BACKEND: "cpu"
     steps:
       - name: Checkout MaxText
-        if: ${{ !inputs.maxtext_installed }}
         uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
         with:
           ref: ${{ inputs.maxtext_sha }}
           persist-credentials: false
       - name: Download the MaxText wheel
-        if: ${{ !inputs.maxtext_installed }}
         uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
         with:
           name: maxtext-wheel
       - name: Install MaxText and Dependencies
-        if: ${{ !inputs.maxtext_installed }}
         shell: bash
         run: |
           # 1. Create virtual environment
@@ -89,25 +143,18 @@ jobs:
           install_tpu_post_train_extra_deps
 
           python3 -m pip freeze
-      - name: Run Post-Training Notebooks
+      - name: Run Post-Training Notebook
         shell: bash
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-          MAXTEXT_INSTALLED: ${{ inputs.maxtext_installed }}
+          NOTEBOOK_NAME: ${{ matrix.notebook }}
           # TODO: Fix evaluation in sft_qwen3_demo.ipynb and remove this env variable
           RUN_EVALUATION: "false"
         run: |
-          if [ "${MAXTEXT_INSTALLED}" == "true" ]; then
-            # Move to the directory where code is baked into the image. See the Dockerfile.
-            # This is necessary because GHA sets an empty workspace by default.
-            cd /deps
-            PYTHON_EXE="python3"
-            PAPERMILL_EXE="papermill"
-          else
-            PYTHON_EXE=".venv/bin/python3"
-            PAPERMILL_EXE=".venv/bin/papermill"
-            source .venv/bin/activate
-          fi
+          PYTHON_EXE=".venv/bin/python3"
+          PAPERMILL_EXE=".venv/bin/papermill"
+          source .venv/bin/activate
+
           export PYTHONPATH="${PWD}/src${PYTHONPATH:+:${PYTHONPATH}}"
 
           MAXTEXT_REPO_ROOT=$(pwd)
@@ -121,28 +168,22 @@ jobs:
           # Run Hugging Face authentication
           hf auth login --token "$HF_TOKEN"
 
-          for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/*.ipynb; do
-            filename=$(basename "$notebook")
-            if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" || "$filename" == "maxtext_with_gepa.ipynb" || "$filename" == "demo_decoding.ipynb" || "$filename" == "dpo_qwen3_demo.ipynb" ]]; then
-              echo "Skipping $filename"
-              continue
-            fi
-            output_name="${filename%.ipynb}_output.ipynb"
+          notebook="$MAXTEXT_NOTEBOOKS_ROOT/$NOTEBOOK_NAME"
+          output_name="${NOTEBOOK_NAME%.ipynb}_output.ipynb"
 
-            echo "------------------------------------------------------"
-            echo "Running $filename ..."
-            echo "------------------------------------------------------"
+          echo "------------------------------------------------------"
+          echo "Running $NOTEBOOK_NAME ..."
+          echo "------------------------------------------------------"
 
-            $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv
+          $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
+          # Clean up any checkpoint directories created by the notebook to avoid filling up disk space
+          echo "Post-notebook disk cleanup for $NOTEBOOK_NAME ..."
+          rm -rf "$MAXTEXT_PKG_DIR"/*_output
+          rm -rf "$HOME/.cache/huggingface/hub"
       - name: Upload Outputs
         if: always()
         uses: actions/upload-artifact@5d5d22a31266ced268874388b861e4b58bb5c2f3 # v4.3.1
         with:
-          name: notebook-outputs-${{ inputs.device_name }}
+          name: notebook-outputs-${{ matrix.notebook }}-${{ inputs.device_name }}
           path: ./*_output.ipynb

.github/workflows/run_tests_coordinator.yml

@@ -119,7 +119,8 @@ jobs:
           "gpu-integration": "cuda12",
           "cpu-unit": "cpu",
           "cpu-integration": "cpu",
-          "cpu-post-training-unit": "cpu"
+          "cpu-post-training-unit": "cpu",
+          "cpu-post-training-integration": "cpu"
         }')[inputs.flavor] }}
 
       device_name: >-
@@ -136,7 +137,8 @@ jobs:
           "gpu-integration": "a100-40gb-4",
           "cpu-unit": "X64",
           "cpu-integration": "X64",
-          "cpu-post-training-unit": "X64"
+          "cpu-post-training-unit": "X64",
+          "cpu-post-training-integration": "X64"
         }')[inputs.flavor] }}
 
       cloud_runner: >-
@@ -153,7 +155,8 @@ jobs:
           "gpu-integration": "linux-x86-a2-48-a100-4gpu",
           "cpu-unit": "linux-x86-n2-32",
           "cpu-integration": "linux-x86-n2-32",
-          "cpu-post-training-unit": "linux-x86-n2-32"
+          "cpu-post-training-unit": "linux-x86-n2-32",
+          "cpu-post-training-integration": "linux-x86-n2-32"
         }')[inputs.flavor] }}
       # Pytest Marker Mapping
       pytest_marker: >-
@@ -170,7 +173,8 @@ jobs:
             "gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training",
             "cpu-unit": "cpu_only and not post_training and not integration_test",
             "cpu-integration": "cpu_only and not post_training and integration_test",
-            "cpu-post-training-unit": "cpu_only and post_training"
+            "cpu-post-training-unit": "cpu_only and post_training",
+            "cpu-post-training-integration": "cpu_only and post_training and integration_test"
           }')[inputs.flavor] }}
 
       pytest_addopts: >-
@@ -187,7 +191,8 @@ jobs:
             "gpu-integration": "",
             "cpu-unit": "",
             "cpu-integration": "",
-            "cpu-post-training-unit": "tests/post_training/unit tests/unit"
+            "cpu-post-training-unit": "tests/post_training/unit tests/unit",
+            "cpu-post-training-integration": "tests/post_training/integration"
           }')[inputs.flavor] }}
 
       pytest_extra_args: >-
@@ -204,7 +209,8 @@ jobs:
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
             "cpu-integration": "--ignore=tests/post_training",
-            "cpu-post-training-unit": ""
+            "cpu-post-training-unit": "",
+            "cpu-post-training-integration": ""
           }')[inputs.flavor] }}
 
       # Resource Scaling

src/maxtext/examples/lora_llama3_demo.ipynb

@@ -140,7 +140,11 @@
     "from flax import nnx\n",
     "from etils import epath\n",
     "\n",
-    "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
+    "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")\n",
+    "\n",
+    "from absl import flags\n",
+    "if not flags.FLAGS.is_parsed():\n",
+    "  flags.FLAGS.mark_as_parsed()"
    ]
   },
   {
@@ -301,7 +305,7 @@
     "        f\"train_split={TRAIN_DATA_SPLIT}\",\n",
     "        f\"hf_data_dir={HF_DATA_DIR}\",\n",
     "        f\"train_data_columns={TRAIN_DATA_COLUMNS}\",\n",
-    "        \"steps=200\",\n",
+    "        \"steps=5\",\n",
     "        \"per_device_batch_size=1\",\n",
     "        \"max_target_length=512\",\n",
     "        \"learning_rate=5e-5\", \n",
@@ -465,6 +469,12 @@
    "source": [
     "print(\"Starting LoRA SFT Training...\")\n",
     "trainer = train_sft.train_model(config, trainer, mesh)\n",
+    "# Verify that the expected number of steps actually executed\n",
+    "if trainer.train_steps < config.steps:\n",
+    "    raise RuntimeError(\n",
+    "        f\"Training ended prematurely! Expected {config.steps} steps, \"\n",
+    "        f\"but only completed {trainer.train_steps} steps.\"\n",
+    "    )\n",
     "print(\"LoRA SFT Training Complete!\")"
    ]
   },

src/maxtext/examples/sft_llama3_demo_tpu.ipynb

@@ -127,7 +127,11 @@
     "from etils import epath\n",
     "\n",
     "\n",
-    "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
+    "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")\n",
+    "\n",
+    "from absl import flags\n",
+    "if not flags.FLAGS.is_parsed():\n",
+    "  flags.FLAGS.mark_as_parsed()"
    ]
   },
   {
@@ -253,7 +257,7 @@
     "    f\"{MAXTEXT_PKG_DIR}/configs/post_train/sft.yml\",\n",
     "    f\"load_parameters_path={MODEL_CHECKPOINT_PATH}\",\n",
     "    f\"model_name={MODEL_NAME}\",\n",
-    "    \"steps=100\",\n",
+    "    \"steps=5\",\n",
     "    \"per_device_batch_size=1\",\n",
     "    \"max_target_length=1024\",\n",
     "    \"learning_rate=2.0e-5\",\n",
@@ -296,6 +300,12 @@
     "\n",
     "try:\n",
     "    trainer, mesh = train_sft.train(config)\n",
+    "    # Verify that the expected number of steps actually executed\n",
+    "    if trainer.train_steps < config.steps:\n",
+    "        raise RuntimeError(\n",
+    "            f\"Training ended prematurely! Expected {config.steps} steps, \"\n",
+    "            f\"but only completed {trainer.train_steps} steps.\"\n",
+    "        )\n",
     "    print(\"\\n\" + \"=\" * 60)\n",
     "    print(\"✅ Training Completed Successfully!\")\n",
     "    print(\"=\" * 60)\n",

src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb

@@ -236,7 +236,7 @@
     "    f\"load_parameters_path={MODEL_CHECKPOINT_PATH}\",\n",
     "    f\"model_name={MODEL_NAME}\",\n",
     "    f\"tokenizer_path={TOKENIZER_NAME}\",\n",
-    "    \"steps=10\",\n",
+    "    \"steps=5\",\n",
     "    \"attention=dot_product\",\n",
     "    \"per_device_batch_size=1\",\n",
     "    \"max_prefill_predict_length=1024\",\n",

src/maxtext/examples/sft_qwen3_demo.ipynb

@@ -144,7 +144,11 @@
     "from flax import nnx\n",
     "from etils import epath\n",
     "\n",
-    "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
+    "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")\n",
+    "\n",
+    "from absl import flags\n",
+    "if not flags.FLAGS.is_parsed():\n",
+    "  flags.FLAGS.mark_as_parsed()"
    ]
   },
   {
@@ -315,7 +319,7 @@
     "        f\"train_split={TRAIN_DATA_SPLIT}\",\n",
     "        f\"hf_data_dir={HF_DATA_DIR}\",\n",
     "        f\"train_data_columns={TRAIN_DATA_COLUMNS}\",\n",
-    "        \"steps=500\",\n",
+    "        \"steps=5\",\n",
     "        \"per_device_batch_size=1\",\n",
     "        \"max_target_length=1024\",\n",
     "        \"learning_rate=3e-6\",\n",
@@ -488,6 +492,12 @@
    "source": [
     "print(\"Starting SFT Training...\")\n",
     "trainer = train_sft.train_model(config, trainer, mesh)\n",
+    "# Verify that the expected number of steps actually executed\n",
+    "if trainer.train_steps < config.steps:\n",
+    "    raise RuntimeError(\n",
+    "        f\"Training ended prematurely! Expected {config.steps} steps, \"\n",
+    "        f\"but only completed {trainer.train_steps} steps.\"\n",
+    "    )\n",
     "print(\"SFT Training Complete!\")"
    ]
   },

src/maxtext/integration/vllm/maxtext_vllm_rollout.py

@@ -291,7 +291,6 @@ def __init__(
     engine_kwargs = {
         "max_model_len": cache_config_or_size,
         "model": rollout_config.rollout_vllm_model_version,
-        "swap_space": getattr(rollout_config, "rollout_vllm_swap_space_size_gb", maxtext_config.swap_space_vllm_gb),
         # Async scheduling causes KeyError in dp_scheduler on slow models
         # (30B+) where inference latency exceeds the scheduler's window.
         "async_scheduling": rollout_config.rollout_vllm_async_scheduling,