diff --git a/examples/vl_dpo_gemma3.ipynb b/examples/vl_dpo_gemma3.ipynb index 080cbebb7..096171e65 100644 --- a/examples/vl_dpo_gemma3.ipynb +++ b/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,9 +485,18 @@ }, "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": { @@ -464,8 +504,7 @@ "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 -} +} \ No newline at end of file