Seth/gemma4-research
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
eecebe7ef5890dcb6fa83d62740d7dbcb43ff737
gemma4-research/tooling/fine-tuning/google-cookbook/huggingface_vision_finetune_qlora.ipynb
T
Mortdecai eecebe7ef5 docs: add canonical tooling corpus (147 files) from Google/HF/frameworks 
Five-lane parallel research pass. Each subdir under tooling/ has its own
README indexing downloaded files with verified upstream sources.

- google-official/: deepmind-gemma JAX examples, gemma_pytorch scripts,
  gemma.cpp API server docs, google-gemma/cookbook notebooks, ai.google.dev
  HTML snapshots, Gemma 3 tech report
- huggingface/: 8 gemma-4-* model cards, chat-template .jinja files,
  tokenizer_config.json, transformers gemma4/ source, launch blog posts,
  official HF Spaces app.py
- inference-frameworks/: vLLM/llama.cpp/MLX/Keras-hub/TGI/Gemini API/Vertex AI
  comparison, run_commands.sh with 8 working launches, 9 code snippets
- gemma-family/: 12 per-variant briefs (ShieldGemma 2, CodeGemma, PaliGemma 2,
  Recurrent/Data/Med/TxGemma, Embedding/Translate/Function/Dolphin/SignGemma)
- fine-tuning/: Unsloth Gemma 4 notebooks, Axolotl YAMLs (incl 26B-A4B MoE),
  TRL scripts, Google cookbook fine-tune notebooks, recipe-recommendation.md

Findings that update earlier CORPUS_* docs are flagged in tooling/README.md
(not applied) — notably the new <|turn>/<turn|> prompt format, gemma_pytorch
abandonment, gemma.cpp Gemini-API server, transformers AutoModelForMultimodalLM,
FA2 head_dim=512 break, 26B-A4B MoE quantization rules, no Gemma 4 tech
report PDF yet, no Gemma-4-generation specialized siblings yet.

Pre-commit secrets hook bypassed per user authorization — flagged "secrets"
are base64 notebook cell outputs and example Ed25519 keys in the HDP
agentic-security demo, not real credentials.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-18 12:24:48 -04:00

981 lines
41 KiB
Plaintext
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "pn1797sn9Jb_"
},
"source": [
"##### Copyright 2025 Google LLC."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "uivh5PY69ISg"
},
"outputs": [],
"source": [
"#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "O83CmJ2j9L3n"
},
"source": [
"# Fine-Tune Gemma for Vision Tasks using Hugging Face Transformers and QLoRA"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f9673bd6"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://ai.google.dev/gemma/docs/core/huggingface_vision_finetune_qlora\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemma/cookbook/blob/main/docs/core/huggingface_vision_finetune_qlora.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://kaggle.com/kernels/welcome?src=https://github.com/google-gemma/cookbook/blob/main/docs/core/huggingface_vision_finetune_qlora.ipynb\"><img src=\"https://www.kaggle.com/static/images/logos/kaggle-logo-transparent-300.png\" height=\"32\" width=\"70\"/>Run in Kaggle</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://console.cloud.google.com/vertex-ai/colab/import/https%3A%2F%2Fraw.githubusercontent.com%2Fgoogle-gemma%2Fcookbook%2Fmain%2Fdocs%2Fcore%2Fhuggingface_vision_finetune_qlora.ipynb\"><img src=\"https://ai.google.dev/images/cloud-icon.svg\" width=\"40\" />Open in Vertex AI</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/google-gemma/cookbook/blob/main/docs/core/huggingface_vision_finetune_qlora.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e624ec07"
},
"source": [
"This guide walks you through how to fine-tune Gemma on a custom image and text dataset for a vision task (generating product descriptions) using Hugging Face [Transformers](https://huggingface.co/docs/transformers/index) and [TRL](https://huggingface.co/docs/trl/index). You will learn:\n",
"\n",
"- What is Quantized Low-Rank Adaptation (QLoRA)\n",
"- Setup development environment\n",
"- Create and prepare the fine-tuning dataset\n",
"- Fine-tune Gemma using TRL and the SFTTrainer\n",
"- Test Model Inference and generate product descriptions from images and text.\n",
"\n",
"Note: This guide requires a GPU which support bfloat16 data type such as NVIDIA L4 or NVIDIA A100 and more than 16GB of memory.\n",
"\n",
"## What is Quantized Low-Rank Adaptation (QLoRA)\n",
"\n",
"This guide demonstrates the use of [Quantized Low-Rank Adaptation (QLoRA)](https://arxiv.org/abs/2305.14314), which emerged as a popular method to efficiently fine-tune LLMs as it reduces computational resource requirements while maintaining high performance. In QloRA, the pretrained model is quantized to 4-bit and the weights are frozen. Then trainable adapter layers (LoRA) are attached and only the adapter layers are trained. Afterwards, the adapter weights can be merged with the base model or kept as a separate adapter.\n",
"\n",
"## Setup development environment\n",
"\n",
"The first step is to install Hugging Face Libraries, including TRL, and datasets to fine-tune open model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ba51aa79"
},
"outputs": [],
"source": [
"# Install Pytorch & other libraries\n",
"%pip install torch tensorboard torchvision\n",
"\n",
"# Install Transformers\n",
"%pip install transformers\n",
"\n",
"# Install Hugging Face libraries\n",
"%pip install datasets accelerate evaluate bitsandbytes trl peft protobuf pillow sentencepiece\n",
"\n",
"# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100\n",
"#%pip install flash-attn"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7ef3d54b"
},
"source": [
"_Note: If you are using a GPU with Ampere architecture (such as NVIDIA L4) or newer, you can use Flash attention. Flash Attention is a method that significantly speeds computations up and reduces memory usage from quadratic to linear in sequence length, leading to acelerating training up to 3x. Learn more at [FlashAttention](https://github.com/Dao-AILab/flash-attention/tree/main)._\n",
"\n",
"You need a valid Hugging Face Token to publish your model. If you are running inside a Google Colab, you can securely use your Hugging Face Token using the Colab secrets otherwise you can set the token as directly in the `login` method. Make sure your token has write access too, as you push your model to the Hub during training."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "b6d79c93"
},
"outputs": [],
"source": [
"# Login into Hugging Face Hub\n",
"from huggingface_hub import login\n",
"login()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "42c60525"
},
"source": [
"## Create and prepare the fine-tuning dataset\n",
"\n",
"When fine-tuning LLMs, it is important to know your use case and the task you want to solve. This helps you create a dataset to fine-tune your model. If you haven't defined your use case yet, you might want to go back to the drawing board.\n",
"\n",
"As an example, this guide focuses on the following use case:\n",
"\n",
"- Fine-tuning a Gemma model to generate concise, SEO-optimized product descriptions for an ecommerce platform, specifically tailored for mobile search.\n",
"\n",
"This guide uses the [philschmid/amazon-product-descriptions-vlm](https://huggingface.co/datasets/philschmid/amazon-product-descriptions-vlm) dataset, a dataset of Amazon product descriptions, including product images and categories.\n",
"\n",
"Hugging Face TRL supports multimodal conversations. The important piece is the \"image\" role, which tells the processing class that it should load the image. The structure should follow:\n",
"\n",
"```json\n",
"{\"messages\": [{\"role\": \"system\", \"content\": [{\"type\": \"text\", \"text\":\"You are...\"}]}, {\"role\": \"user\", \"content\": [{\"type\": \"text\", \"text\": \"...\"}, {\"type\": \"image\"}]}, {\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"...\"}]}]}\n",
"{\"messages\": [{\"role\": \"system\", \"content\": [{\"type\": \"text\", \"text\":\"You are...\"}]}, {\"role\": \"user\", \"content\": [{\"type\": \"text\", \"text\": \"...\"}, {\"type\": \"image\"}]}, {\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"...\"}]}]}\n",
"{\"messages\": [{\"role\": \"system\", \"content\": [{\"type\": \"text\", \"text\":\"You are...\"}]}, {\"role\": \"user\", \"content\": [{\"type\": \"text\", \"text\": \"...\"}, {\"type\": \"image\"}]}, {\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"...\"}]}]}\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c4ecf6db"
},
"source": [
"You can now use the Hugging Face Datasets library to load the dataset and create a prompt template to combine the image, product name, and category, and add a system message. The dataset includes images as`Pil.Image` objects."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "40c3a2cf"
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8d1259be3dfa4b1e899c97026276ee41",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"README.md: 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "a5554c0595144c949b578eb1cbdfd0fd",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"data/train-00000-of-00001.parquet: 0%| | 0.00/47.6M [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "9ed0567e2e4e40a88c7eddfe7d6a6e2f",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Generating train split: 0%| | 0/1345 [00:00<?, ? examples/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'role': 'system', 'content': 'You are an expert product description writer for Amazon.'}, {'role': 'user', 'content': [{'type': 'text', 'text': \"Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.\\nOnly return description. The description should be SEO optimized and for a better mobile search experience.\\n\\n<PRODUCT>\\nRazor Agitator BMX/Freestyle Bike, 20-Inch\\n</PRODUCT>\\n\\n<CATEGORY>\\nSports & Outdoors | Outdoor Recreation | Cycling | Kids' Bikes & Accessories | Kids' Bikes\\n</CATEGORY>\\n\"}, {'type': 'image', 'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x413 at 0x7B7250181790>}]}, {'role': 'assistant', 'content': [{'type': 'text', 'text': 'Conquer the streets with the Razor Agitator BMX Bike! This 20-inch freestyle bike is built for young riders ready to take on any challenge. Durable frame, responsive handling – perfect for tricks and cruising. Get yours today!'}]}]\n"
]
}
],
"source": [
"from datasets import load_dataset\n",
"from PIL import Image\n",
"\n",
"# System message for the assistant\n",
"system_message = \"You are an expert product description writer for Amazon.\"\n",
"\n",
"# User prompt that combines the user query and the schema\n",
"user_prompt = \"\"\"Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.\n",
"Only return description. The description should be SEO optimized and for a better mobile search experience.\n",
"\n",
"<PRODUCT>\n",
"{product}\n",
"</PRODUCT>\n",
"\n",
"<CATEGORY>\n",
"{category}\n",
"</CATEGORY>\n",
"\"\"\"\n",
"\n",
"# Convert dataset to OAI messages\n",
"def format_data(sample):\n",
" return {\n",
" \"messages\": [\n",
" {\n",
" \"role\": \"system\",\n",
" #\"content\": [{\"type\": \"text\", \"text\": system_message}],\n",
" \"content\": system_message,\n",
" },\n",
" {\n",
" \"role\": \"user\",\n",
" \"content\": [\n",
" {\n",
" \"type\": \"text\",\n",
" \"text\": user_prompt.format(\n",
" product=sample[\"Product Name\"],\n",
" category=sample[\"Category\"],\n",
" ),\n",
" },\n",
" {\n",
" \"type\": \"image\",\n",
" \"image\": sample[\"image\"],\n",
" },\n",
" ],\n",
" },\n",
" {\n",
" \"role\": \"assistant\",\n",
" \"content\": [{\"type\": \"text\", \"text\": sample[\"description\"]}],\n",
" },\n",
" ],\n",
" }\n",
"\n",
"def process_vision_info(messages: list[dict]) -> list[Image.Image]:\n",
" image_inputs = []\n",
" # Iterate through each conversation\n",
" for msg in messages:\n",
" # Get content (ensure it's a list)\n",
" content = msg.get(\"content\", [])\n",
" if not isinstance(content, list):\n",
" content = [content]\n",
"\n",
" # Check each content element for images\n",
" for element in content:\n",
" if isinstance(element, dict) and (\n",
" \"image\" in element or element.get(\"type\") == \"image\"\n",
" ):\n",
" # Get the image and convert to RGB\n",
" if \"image\" in element:\n",
" image = element[\"image\"]\n",
" else:\n",
" image = element\n",
" image_inputs.append(image.convert(\"RGB\"))\n",
" return image_inputs\n",
"\n",
"# Load dataset from the hub\n",
"dataset = load_dataset(\"philschmid/amazon-product-descriptions-vlm\", split=\"train\")\n",
"dataset = dataset.train_test_split(test_size=0.1)\n",
"\n",
"# Convert dataset to OAI messages\n",
"# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes\n",
"dataset_train = [format_data(sample) for sample in dataset[\"train\"]]\n",
"dataset_test = [format_data(sample) for sample in dataset[\"test\"]]\n",
"\n",
"print(dataset_train[345][\"messages\"])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c0eb2e06"
},
"source": [
"## Fine-tune Gemma using TRL and the SFTTrainer\n",
"\n",
"You are now ready to fine-tune your model. Hugging Face TRL [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) makes it straightforward to supervise fine-tune open LLMs. The `SFTTrainer` is a subclass of the `Trainer` from the `transformers` library and supports all the same features, including logging, evaluation, and checkpointing, but adds additional quality of life features, including:\n",
"\n",
"* Dataset formatting, including conversational and instruction formats\n",
"* Training on completions only, ignoring prompts\n",
"* Packing datasets for more efficient training\n",
"* Parameter-efficient fine-tuning (PEFT) support including QloRA\n",
"* Preparing the model and tokenizer for conversational fine-tuning (such as adding special tokens)\n",
"\n",
"The following code loads the Gemma model and tokenizer from Hugging Face and initializes the quantization configuration.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "18069ed2"
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "42e58727637d4495ad8c5f753c5bcd06",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"config.json: 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "b11ec04ab48043b9937cfa3822b4fa42",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"model.safetensors: 0%| | 0.00/10.2G [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "7659ae83140247efacee26159ca363b6",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading weights: 0%| | 0/2011 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "28c7b23ad9ba4316a8c95992884ad1d7",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"generation_config.json: 0%| | 0.00/149 [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "44b08b5b2cad4385893e29d5240a98d7",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"processor_config.json: 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6eec3330ff144b3c9ad863cc89ed5709",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"chat_template.jinja: 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "47f8fdc1492e4bb9b8d8fe9535c97d2c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"config.json: 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ff3072e44aec41b6a0f6a28aeba99c4d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"tokenizer_config.json: 0.00B [00:00, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "3d3e0871ad0e4642a5e2ca6f4baeebe4",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"tokenizer.json: 0%| | 0.00/32.2M [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"import torch\n",
"from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig\n",
"\n",
"# Hugging Face model id\n",
"model_id = \"google/gemma-4-E2B\" # @param [\"google/gemma-4-E2B\",\"google/gemma-4-E4B\"] {\"allow-input\":true}\n",
"\n",
"# Check if GPU benefits from bfloat16\n",
"if torch.cuda.get_device_capability()[0] < 8:\n",
" raise ValueError(\"GPU does not support bfloat16, please use a GPU that supports bfloat16.\")\n",
"\n",
"# Define model init arguments\n",
"model_kwargs = dict(\n",
" dtype=torch.bfloat16, # What torch dtype to use, defaults to auto\n",
" device_map=\"auto\", # Let torch decide how to load the model\n",
")\n",
"\n",
"# BitsAndBytesConfig int-4 config\n",
"model_kwargs[\"quantization_config\"] = BitsAndBytesConfig(\n",
" load_in_4bit=True,\n",
" bnb_4bit_use_double_quant=True,\n",
" bnb_4bit_quant_type=\"nf4\",\n",
" bnb_4bit_compute_dtype=model_kwargs[\"dtype\"],\n",
" bnb_4bit_quant_storage=model_kwargs[\"dtype\"],\n",
")\n",
"\n",
"# Load model and tokenizer\n",
"model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)\n",
"processor = AutoProcessor.from_pretrained(\"google/gemma-4-E2B-it\") # Load the Instruction Tokenizer to use the official Gemma template"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "37ec1d1b"
},
"source": [
"The `SFTTrainer` supports a built-in integration with `peft`, which makes it straightforward to efficiently tune LLMs using QLoRA. You only need to create a `LoraConfig` and provide it to the trainer."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ed00e846"
},
"outputs": [],
"source": [
"from peft import LoraConfig\n",
"\n",
"peft_config = LoraConfig(\n",
" lora_alpha=16,\n",
" lora_dropout=0.05,\n",
" r=16,\n",
" bias=\"none\",\n",
" target_modules=\"all-linear\",\n",
" task_type=\"CAUSAL_LM\",\n",
" modules_to_save=[\"lm_head\", \"embed_tokens\"], # make sure to save the lm_head and embed_tokens as you train the special tokens\n",
" ensure_weight_tying=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bbd9fc1b"
},
"source": [
"Before you can start your training, you need to define the hyperparameter you want to use in a `SFTConfig` and a custom `collate_fn` to handle the vision processing. The `collate_fn` converts the messages with text and images into a format that the model can understand.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "989be3c1"
},
"outputs": [],
"source": [
"from trl import SFTConfig\n",
"\n",
"args = SFTConfig(\n",
" output_dir=\"gemma-product-description\", # directory to save and repository id\n",
" num_train_epochs=3, # number of training epochs\n",
" per_device_train_batch_size=1, # batch size per device during training\n",
" optim=\"adamw_torch_fused\", # use fused adamw optimizer\n",
" logging_steps=5, # log every 5 steps\n",
" save_strategy=\"epoch\", # save checkpoint every epoch\n",
" eval_strategy=\"epoch\", # evaluate checkpoint every epoch\n",
" learning_rate=2e-4, # learning rate, based on QLoRA paper\n",
" bf16=True, # use bfloat16 precision\n",
" max_grad_norm=0.3, # max gradient norm based on QLoRA paper\n",
" lr_scheduler_type=\"constant\", # use constant learning rate scheduler\n",
" push_to_hub=True, # push model to hub\n",
" report_to=\"tensorboard\", # report metrics to tensorboard\n",
" dataset_text_field=\"\", # need a dummy field for collator\n",
" dataset_kwargs={\"skip_prepare_dataset\": True}, # important for collator\n",
" remove_unused_columns = False # important for collator\n",
")\n",
"\n",
"# Create a data collator to encode text and image pairs\n",
"def collate_fn(examples):\n",
" texts = []\n",
" images = []\n",
" for example in examples:\n",
" image_inputs = process_vision_info(example[\"messages\"])\n",
" text = processor.apply_chat_template(\n",
" example[\"messages\"], add_generation_prompt=False, tokenize=False\n",
" )\n",
" texts.append(text.strip())\n",
" images.append(image_inputs)\n",
"\n",
" # Tokenize the texts and process the images\n",
" batch = processor(text=texts, images=images, return_tensors=\"pt\", padding=True)\n",
"\n",
" # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation\n",
" labels = batch[\"input_ids\"].clone()\n",
"\n",
" # Mask tokens for not being used in the loss computation\n",
" labels[labels == processor.tokenizer.pad_token_id] = -100\n",
" labels[labels == processor.tokenizer.boi_token_id] = -100\n",
" labels[labels == processor.tokenizer.image_token_id] = -100\n",
" labels[labels == processor.tokenizer.eoi_token_id] = -100\n",
"\n",
" batch[\"labels\"] = labels\n",
" return batch"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dd88e798"
},
"source": [
"You now have every building block you need to create your `SFTTrainer` to start the training of your model.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ade95df7"
},
"outputs": [],
"source": [
"from trl import SFTTrainer\n",
"\n",
"# Create Trainer object\n",
"trainer = SFTTrainer(\n",
" model=model,\n",
" args=args,\n",
" train_dataset=dataset_train,\n",
" eval_dataset=dataset_test,\n",
" peft_config=peft_config,\n",
" processing_class=processor,\n",
" data_collator=collate_fn,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fad61a6a"
},
"source": [
"Start training by calling the `train()` method."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "995e7e38"
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 1, 'bos_token_id': 2, 'pad_token_id': 0}.\n"
]
},
{
"data": {
"text/html": [
"\n",
" <div>\n",
" \n",
" <progress value='456' max='456' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
" [456/456 11:20, Epoch 3/3]\n",
" </div>\n",
" <table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: left;\">\n",
" <th>Epoch</th>\n",
" <th>Training Loss</th>\n",
" <th>Validation Loss</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <td>1</td>\n",
" <td>1.326710</td>\n",
" <td>1.441816</td>\n",
" </tr>\n",
" <tr>\n",
" <td>2</td>\n",
" <td>1.042711</td>\n",
" <td>1.320613</td>\n",
" </tr>\n",
" <tr>\n",
" <td>3</td>\n",
" <td>0.739179</td>\n",
" <td>1.458798</td>\n",
" </tr>\n",
" </tbody>\n",
"</table><p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# Start training, the model will be automatically saved to the Hub and the output directory\n",
"trainer.train()\n",
"\n",
"# Save the final model again to the Hugging Face Hub\n",
"trainer.save_model()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b47b9733"
},
"source": [
"Before you can test your model, make sure to free the memory."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "40a32ed7"
},
"outputs": [],
"source": [
"# free the memory again\n",
"del model\n",
"del trainer\n",
"torch.cuda.empty_cache()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "862e9728"
},
"source": [
"When using QLoRA, you only train adapters and not the full model. This means when saving the model during training you only save the adapter weights and not the full model. If you want to save the full model, which makes it easier to use with serving stacks like vLLM or TGI, you can merge the adapter weights into the model weights using the `merge_and_unload` method and then save the model with the `save_pretrained` method. This saves a default model, which can be used for inference.\n",
"\n",
"Note: It requires more than 30GB of CPU Memory when you want to merge the adapter into the model. You can skip this and continue with Test Model Inference.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "761e324b"
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "20d63c526a854f2a880882c246ac3b3d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading weights: 0%| | 0/2011 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8703db4619fd4f8eb66bf0cc2211dc7e",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Writing model shards: 0%| | 0/5 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"['merged_model/processor_config.json']"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from peft import PeftModel\n",
"\n",
"# Load Model base model\n",
"model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)\n",
"\n",
"# Merge LoRA and base model and save\n",
"peft_model = PeftModel.from_pretrained(model, args.output_dir)\n",
"merged_model = peft_model.merge_and_unload()\n",
"merged_model.save_pretrained(\"merged_model\", safe_serialization=True, max_shard_size=\"2GB\")\n",
"\n",
"processor = AutoProcessor.from_pretrained(args.output_dir)\n",
"processor.save_pretrained(\"merged_model\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bf86e31d"
},
"source": [
"## Test Model Inference and generate product descriptions\n",
"\n",
"After the training is done, you'll want to evaluate and test your model. You can load different samples from the test dataset and evaluate the model on those samples.\n",
"\n",
"Note: Evaluating Generative AI models is not a trivial task since one input can have multiple correct outputs. This guide only focuses on manual evaluation and vibe checks.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "aab1c5c5"
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "89b0d1d25dba4e8e8642c41e69c4c65e",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading weights: 0%| | 0/2012 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"The tied weights mapping and config for this model specifies to tie model.language_model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints with different values, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning.\n"
]
}
],
"source": [
"model_id = \"merged_model\"\n",
"\n",
"# Load Model with PEFT adapter\n",
"model = AutoModelForImageTextToText.from_pretrained(\n",
" model_id,\n",
" device_map=\"auto\",\n",
" dtype=\"auto\",\n",
")\n",
"processor = AutoProcessor.from_pretrained(model_id)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3dccb57c"
},
"source": [
"You can test inference by providing a product name, category and image. The `sample` includes a marvel action figure.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1fd887f4"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<bos><|turn>system\n",
"You are an expert product description writer for Amazon.<turn|>\n",
"<|turn>user\n",
"\n",
"\n",
"<|image|>\n",
"\n",
"Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.\n",
"Only return description. The description should be SEO optimized and for a better mobile search experience.\n",
"\n",
"<PRODUCT>\n",
"Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur\n",
"</PRODUCT>\n",
"\n",
"<CATEGORY>\n",
"Toys & Games | Toy Figures & Playsets | Action Figures\n",
"</CATEGORY><turn|>\n",
"<|turn>model\n",
"\n",
"MODEL OUTPUT>> \n",
"\n",
"Enhance your collection with the Marvel Avengers - Avengers Assemble Ultron-Comforter Set! This soft and cuddly blanket and pillowcase feature everyone's favorite Avengers, Iron Man, and his loyal companion War Machine. Officially licensed by Marvel. Bring home the heroic team!\n"
]
}
],
"source": [
"import requests\n",
"from PIL import Image\n",
"\n",
"# Test sample with Product Name, Category and Image\n",
"sample = {\n",
" \"product_name\": \"Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur\",\n",
" \"category\": \"Toys & Games | Toy Figures & Playsets | Action Figures\",\n",
" \"image\": Image.open(requests.get(\"https://m.media-amazon.com/images/I/81+7Up7IWyL._AC_SY300_SX300_.jpg\", stream=True).raw).convert(\"RGB\")\n",
"}\n",
"\n",
"def generate_description(sample, model, processor):\n",
" # Convert sample into messages and then apply the chat template\n",
" messages = [\n",
" {\"role\": \"system\", \"content\": system_message},\n",
" {\"role\": \"user\", \"content\": [\n",
" {\"type\": \"image\",\"image\": sample[\"image\"]},\n",
" {\"type\": \"text\", \"text\": user_prompt.format(product=sample[\"product_name\"], category=sample[\"category\"])},\n",
" ]},\n",
" ]\n",
" text = processor.apply_chat_template(\n",
" messages, tokenize=False, add_generation_prompt=True\n",
" )\n",
" print(text)\n",
" # Process the image and text\n",
" image_inputs = process_vision_info(messages)\n",
" # Tokenize the text and process the images\n",
" inputs = processor(\n",
" text=[text],\n",
" images=image_inputs,\n",
" padding=True,\n",
" return_tensors=\"pt\",\n",
" )\n",
" # Move the inputs to the device\n",
" inputs = inputs.to(model.device)\n",
"\n",
" # Generate the output\n",
" stop_token_ids = [processor.tokenizer.eos_token_id, processor.tokenizer.convert_tokens_to_ids(\"<turn|>\")]\n",
" generated_ids = model.generate(**inputs, max_new_tokens=256, top_p=1.0, do_sample=True, temperature=0.8, eos_token_id=stop_token_ids, disable_compile=True)\n",
" # Trim the generation and decode the output to text\n",
" generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]\n",
" output_text = processor.batch_decode(\n",
" generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False\n",
" )\n",
" return output_text[0]\n",
"\n",
"# generate the description\n",
"description = generate_description(sample, model, processor)\n",
"print(\"MODEL OUTPUT>> \\n\")\n",
"print(description)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6f8ff452"
},
"source": [
"## Summary and next steps\n",
"\n",
"This tutorial covered how to fine-tune a Gemma model for vision tasks using TRL and QLoRA, specifically for generating product descriptions. Check out the following docs next:\n",
"\n",
"* Learn how to [generate text with a Gemma model](https://ai.google.dev/gemma/docs/get_started).\n",
"* Learn how to [fine-tune Gemma for text tasks using Hugging Face Transformers](https://ai.google.dev/gemma/docs/core/huggingface_text_finetune_qlora).\n",
"* Learn how to [full model fine-tune using Hugging Face Transformers](https://ai.google.dev/gemma/docs/core/huggingface_text_full_finetune).\n",
"* Learn how to perform [distributed fine-tuning and inference on a Gemma model](https://ai.google.dev/gemma/docs/core/distributed_tuning).\n",
"* Learn how to [use Gemma open models with Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma).\n",
"* Learn how to [fine-tune Gemma using KerasNLP and deploy to Vertex AI](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_kerasnlp_to_vertexai.ipynb)."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"name": "huggingface_vision_finetune_qlora.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
Reference in New Issue View Git Blame Copy Permalink