Fix vl_dpo_gemma3 notebook execution bugs, JAX deprecations, and OOM memory allocation

examples/vl_dpo_gemma3.ipynb

@@ -8,6 +8,9 @@
       "source": [
         "# Fine-tuning a Vision Language Model (VLM) using DPO\n",
         "\n",
+        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/tunix/blob/main/examples/vl_dpo_gemma3.ipynb) [![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/google/tunix/blob/main/examples/vl_dpo_gemma3.ipynb) [![View on GitHub](https://img.shields.io/badge/GitHub-181717?style=flat&logo=github&logoColor=white)](https://github.com/google/tunix/blob/main/examples/vl_dpo_gemma3.ipynb)\n",
+        "\n",
+        "\n",
         "This notebook demonstrates how to fine-tune a Vision Language Model (VLM),\n",
         "specifically the Gemma 3 4B IT model, using Direct Preference Optimization\n",
         "(DPO). DPO is a method for training language models to align with human\n",
@@ -26,14 +29,18 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "tG2noh8nB7mH"
+      },
       "source": [
         "## Setup"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "fV_4i37kB7mH"
+      },
       "source": [
         "### Installing libraries (restart Colab after running the cell below)"
       ]
@@ -46,24 +53,26 @@
       },
       "outputs": [],
       "source": [
-        "!pip install -q kagglehub\n",
+        "!pip install -Uq kagglehub\n",
         "\n",
-        "!pip install -q tensorflow\n",
-        "!pip install -q tensorboardX\n",
-        "!pip install -q grain\n",
-        "!pip install -q git+https://github.com/google/tunix\n",
-        "!pip install -q git+https://github.com/google/qwix\n",
+        "!pip install -Uq tensorflow\n",
+        "!pip install -Uq tensorboard\n",
+        "!pip install -Uq grain\n",
+        "!pip install -Uq git+https://github.com/google/tunix\n",
+        "!pip install -Uq git+https://github.com/google/qwix\n",
         "\n",
         "!pip uninstall -q -y flax\n",
-        "!pip install -q git+https://github.com/google/flax.git\n",
+        "!pip install -Uq git+https://github.com/google/flax.git\n",
         "\n",
-        "!pip install -q huggingface_hub\n",
-        "!pip install -q datasets"
+        "!pip install -Uq huggingface_hub\n",
+        "!pip install -Uq datasets"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "OsKnnz9AB7mI"
+      },
       "source": [
         "### Imports"
       ]
@@ -87,6 +96,7 @@
         "import numpy as np\n",
         "import grain\n",
         "import optax\n",
+        "import orbax.checkpoint as ocp\n",
         "import qwix\n",
         "import datasets\n",
         "from flax import nnx\n",
@@ -97,13 +107,15 @@
         "from tunix.models.gemma3 import model as gemma3_model_lib\n",
         "from tunix.models.gemma3 import params as params_lib\n",
         "from tunix.processors import image_processor as image_processor_lib\n",
-        "from tunix.sft import metrics_logger\n",
+        "from tunix.sft import checkpoint_options, metrics_logger\n",
         "from tunix.sft.dpo.dpo_trainer import DPOTrainer, DPOTrainingConfig, TrainingInput"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "7skQidDfB7mI"
+      },
       "source": [
         "### Configuration and hyperparameters"
       ]
@@ -156,11 +168,11 @@
         "# Linearly increase learning rate from 0. to 5e-6 in the first 10% training\n",
         "# steps, and then gradually decrease the learning rate to 0 using cosine\n",
         "# scheduler.\n",
-        "BATCH_SIZE = 2\n",
+        "BATCH_SIZE = 1\n",
         "NUM_SAMPLES = 5000\n",
         "NUM_BATCHES = NUM_SAMPLES // BATCH_SIZE\n",
         "EVAL_EVERY_N_STEPS = 1000\n",
-        "NUM_EPOCHS = 10\n",
+        "NUM_EPOCHS = 2\n",
         "MAX_STEPS = int(NUM_BATCHES * NUM_EPOCHS)\n",
         "\n",
         "WARMUP_STEPS = 0.1 * MAX_STEPS\n",
@@ -187,7 +199,9 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "YfxRIdc2B7mJ"
+      },
       "source": [
         "### Load reference model and LoRA model\n",
         "\n",
@@ -205,8 +219,8 @@
       "source": [
         "model_config = gemma3_model_lib.ModelConfig.gemma3_4b_it(text_only=False)\n",
         "\n",
-        "mesh = jax.make_mesh(*MESH)\n",
-        "with mesh:\n",
+        "mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0]))\n",
+        "with jax.set_mesh(mesh):\n",
         "  gemma3 = params_lib.create_model_from_checkpoint(\n",
         "    MODEL_CKPT_PATH, model_config, mesh, dtype=jnp.bfloat16\n",
         "  )\n",
@@ -246,7 +260,7 @@
         "      base_model, lora_provider, **base_model.get_model_input()\n",
         "  )\n",
         "\n",
-        "  with mesh:\n",
+        "  with jax.set_mesh(mesh):\n",
         "    state = nnx.state(lora_model)\n",
         "    pspecs = nnx.get_partition_spec(state)\n",
         "    nnx.update(lora_model, jax.lax.with_sharding_constraint(state, pspecs))\n",
@@ -269,7 +283,9 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "dyMJDV0sB7mK"
+      },
       "source": [
         "## Load dataset\n",
         "\n",
@@ -281,7 +297,9 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "metadata": {},
+      "metadata": {
+        "id": "_PH7dxYAB7mK"
+      },
       "outputs": [],
       "source": [
         "TEMPLATE = (\n",
@@ -330,7 +348,9 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "metadata": {},
+      "metadata": {
+        "id": "2daXgwDuB7mK"
+      },
       "outputs": [],
       "source": [
         "train_ds = load_dataset(\n",
@@ -339,12 +359,14 @@
         "        config=model_config.vision_config\n",
         "    )\n",
         ")\n",
-        "train_ds = train_ds.batch(BATCH_SIZE)"
+        "train_ds = train_ds.repeat(NUM_EPOCHS).batch(BATCH_SIZE)\n"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "se_lD9uIB7mK"
+      },
       "source": [
         "### Define optimizer and initialize trainer\n",
         "\n",
@@ -355,25 +377,33 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "metadata": {},
+      "metadata": {
+        "id": "MzrSkymBB7mK"
+      },
       "outputs": [],
       "source": [
         "# Ckpt saving\n",
-        "checkpointing_options = ocp.CheckpointManagerOptions(\n",
-        "    save_interval_steps=SAVE_INTERVAL_STEPS, max_to_keep=MAX_TO_KEEP\n",
+        "checkpointing_options = checkpoint_options.TunixCheckpointingOptions(\n",
+        "    save_decision_policy=ocp.v1.training.save_decision_policies.ContinuousCheckpointingPolicy(\n",
+        "        minimum_interval_secs=SAVE_INTERVAL_STEPS,\n",
+        "    ),\n",
+        "    preservation_policy=ocp.v1.training.preservation_policies.LatestN(\n",
+        "        n=MAX_TO_KEEP\n",
+        "    )\n",
         ")\n",
         "\n",
         "# Metrics logger\n",
         "metrics_logging_options = metrics_logger.MetricsLoggerOptions(\n",
         "    log_dir=\"/tmp/tensorboard/dpo\", flush_every_n_steps=20\n",
-        ")\n",
-        "   "
+        ")\n"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
-      "metadata": {},
+      "metadata": {
+        "id": "tkeZf_MUB7mK"
+      },
       "outputs": [],
       "source": [
         "# Logs\n",
@@ -384,7 +414,9 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "metadata": {},
+      "metadata": {
+        "id": "4hKZE_txB7mK"
+      },
       "outputs": [],
       "source": [
         "# Optimizer, learning rate scheduler, gradient clipping\n",
@@ -432,16 +464,15 @@
         "    ref_model=gemma3,\n",
         "    optimizer=optimizer,\n",
         "    training_config=dpo_config,\n",
-        "    tokenizer=gemma_tokenizer,\n",
-        "    image_processor=image_processor_lib.ImageProcessor(\n",
-        "      config=model_config.vision_config\n",
-        "    )\n",
+        "    tokenizer=gemma_tokenizer\n",
         ")"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "qwcIvk_ZB7mK"
+      },
       "source": [
         "## Train!"
       ]
@@ -454,18 +485,26 @@
       },
       "outputs": [],
       "source": [
-        "with mesh:\n",
-        "  trainer.train(train_ds)"
+        "with jax.set_mesh(mesh):\n",
+        "  dpo_trainer.train(train_ds)"
       ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "8gVhlHWzEbx7"
+      },
+      "execution_count": null,
+      "outputs": []
     }
   ],
   "metadata": {
     "accelerator": "TPU",
     "colab": {
       "gpuType": "V6E1",
       "machine_shape": "hm",
-      "provenance": [],
-      "toc_visible": true
+      "provenance": []
     },
     "kernelspec": {
       "display_name": "Python 3",
@@ -477,4 +516,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}