Merge pull request #4362 from AI-Hypercomputer:igorts/dpo-notebook-with-eval

PiperOrigin-RevId: 948420388

Author: Google-ML-Automation

.github/workflows/run_jupyter_notebooks.yml

@@ -111,7 +111,7 @@ jobs:
 
           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" ]]; then
+            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

docs/guides/run_python_notebook.md

@@ -186,6 +186,10 @@ jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
 - **`sft_qwen3_demo.ipynb`** → Qwen3-0.6B SFT training and evaluation on [OpenAI's GSM8K dataset](https://huggingface.co/datasets/openai/gsm8k). This notebook is friendly for beginners and runs successfully on Google Colab's free-tier v5e-1 TPU runtime.
 - **`sft_llama3_demo_tpu.ipynb`** → Llama3.1-8B SFT training on [Hugging Face ultrachat_200k dataset](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k). We recommend running this on a v5p-8 TPU VM using [Method 2](#method-2-visual-studio-code-with-tpu-recommended) or [Method 3](#method-3-local-jupyter-lab-with-tpu-recommended).
 
+### Direct Preference Optimization (DPO)
+
+- **`dpo_qwen3_demo.ipynb`** → Qwen3-0.6B DPO training and evaluation on [Argilla's distilabel-intel-orca-dpo-pairs dataset](https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs). This notebook is fully optimized to run evaluations across all 4 TPUs using Tensor Parallelism (TP=4) on a TPU v5p-8 VM.
+
 ### Reinforcement Learning (GRPO/GSPO) Training
 
 - **`rl_llama3_demo.ipynb`** → GRPO/GSPO training on [OpenAI's GSM8K dataset](https://huggingface.co/datasets/openai/gsm8k). We recommend running this on a v5p-8 TPU VM using [Method 2](#method-2-visual-studio-code-with-tpu-recommended) or [Method 3](#method-3-local-jupyter-lab-with-tpu-recommended).

src/maxtext/eval/runner/harness_runner.py

@@ -74,7 +74,23 @@ def _map_results(raw_results: dict, tasks: list[str]) -> dict:
     if acc_norm is not None:
       scores[f"{task}_accuracy_norm"] = round(float(acc_norm) * 100, 2)
 
-    if acc is None and task_r:
+    # Extract ifeval and other task-specific custom accuracy keys
+    has_custom_keys = False
+    for suffix in (
+        "prompt_level_strict_acc,none",
+        "inst_level_strict_acc,none",
+        "prompt_level_loose_acc,none",
+        "inst_level_loose_acc,none",
+        "prompt_level_strict_acc",
+        "inst_level_strict_acc",
+        "prompt_level_loose_acc",
+        "inst_level_loose_acc",
+    ):
+      if task_r.get(suffix) is not None:
+        scores[f"{task}_{suffix.replace(',none', '')}"] = round(float(task_r[suffix]) * 100, 2)
+        has_custom_keys = True
+
+    if acc is None and not has_custom_keys and task_r:
       logger.warning(
           "No known accuracy keys found for task '%s'. Available: %s",
           task,
@@ -167,6 +183,8 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
           "limit": num_samples,
           "log_samples": False,
       }
+      if cfg.get("max_num_seqs") is not None:
+        simple_eval_kwargs["batch_size"] = cfg["max_num_seqs"]
       if apply_chat_template:
         simple_eval_kwargs["apply_chat_template"] = True
       if fewshot_as_multiturn: