Optimiza Gemma para tareas de visión con Hugging Face Transformers y QLoRA

Ver en ai.google.dev Ejecutar en Google Colab Ejecutar en Kaggle Abrir en Vertex AI Ver el código fuente en GitHub

En esta guía, se explica cómo ajustar Gemma en un conjunto de datos personalizado de imágenes y texto para una tarea de visión (generar descripciones de productos) con Transformers y TRL de Hugging Face. Aprenderás a hacer lo siguiente:

  • ¿Qué es la adaptación de bajo rango cuantificada (QLoRA)?
  • Configura el entorno de desarrollo
  • Crea y prepara el conjunto de datos de ajuste
  • Optimiza Gemma con TRL y SFTTrainer
  • Probar la inferencia del modelo y generar descripciones de productos a partir de imágenes y texto

¿Qué es la adaptación de bajo rango cuantificada (QLoRA)?

En esta guía, se demuestra el uso de la adaptación de bajo rango cuantificada (QLoRA), que surgió como un método popular para ajustar LLMs de manera eficiente, ya que reduce los requisitos de recursos de procesamiento y, al mismo tiempo, mantiene un alto rendimiento. En QLoRA, el modelo previamente entrenado se cuantifica en 4 bits y los pesos se inmovilizan. Luego, se adjuntan capas de adaptadores entrenables (LoRA) y solo se entrenan las capas de adaptadores. Luego, los pesos del adaptador se pueden combinar con el modelo base o mantener como un adaptador independiente.

Configura el entorno de desarrollo

El primer paso es instalar las bibliotecas de Hugging Face, incluidas TRL y los conjuntos de datos para ajustar el modelo abierto.

# Install Pytorch & other libraries
%pip install torch tensorboard torchvision

# Install Transformers
%pip install transformers

# Install Hugging Face libraries
%pip install datasets accelerate evaluate bitsandbytes trl peft protobuf pillow sentencepiece

# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100
#%pip install flash-attn

Nota: Si usas una GPU con arquitectura Ampere (como la NVIDIA L4) o una más reciente, puedes usar la atención flash. Flash Attention es un método que acelera significativamente los cálculos y reduce el uso de memoria de cuadrático a lineal en la longitud de la secuencia, lo que acelera el entrenamiento hasta 3 veces. Obtén más información en FlashAttention.

Necesitas un token de Hugging Face válido para publicar tu modelo. Si ejecutas el código dentro de Google Colab, puedes usar tu token de Hugging Face de forma segura con los secretos de Colab. De lo contrario, puedes establecer el token directamente en el método login. Asegúrate de que tu token también tenga acceso de escritura, ya que subirás tu modelo al Hub durante el entrenamiento.

# Login into Hugging Face Hub
from huggingface_hub import login
login()

Crea y prepara el conjunto de datos de ajuste

Cuando ajustas LLMs, es importante que conozcas tu caso de uso y la tarea que deseas resolver. Esto te ayuda a crear un conjunto de datos para ajustar tu modelo. Si aún no definiste tu caso de uso, tal vez debas volver a empezar.

Como ejemplo, esta guía se enfoca en el siguiente caso de uso:

  • Ajustar un modelo de Gemma para generar descripciones de productos concisas y optimizadas para SEO en una plataforma de comercio electrónico, específicamente diseñadas para la búsqueda en dispositivos móviles

En esta guía, se usa el conjunto de datos philschmid/amazon-product-descriptions-vlm, que contiene descripciones de productos de Amazon, incluidas imágenes y categorías de productos.

Hugging Face TRL admite conversaciones multimodales. La parte importante es el rol "image", que le indica a la clase de procesamiento que debe cargar la imagen. La estructura debe seguir el siguiente formato:

{"messages": [{"role": "system", "content": [{"type": "text", "text":"You are..."}]}, {"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image"}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]}]}
{"messages": [{"role": "system", "content": [{"type": "text", "text":"You are..."}]}, {"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image"}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]}]}
{"messages": [{"role": "system", "content": [{"type": "text", "text":"You are..."}]}, {"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image"}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]}]}

Ahora puedes usar la biblioteca de Hugging Face Datasets para cargar el conjunto de datos y crear una plantilla de instrucciones para combinar la imagen, el nombre del producto y la categoría, y agregar un mensaje del sistema. El conjunto de datos incluye imágenes como objetosPil.Image.

from datasets import load_dataset
from PIL import Image

# System message for the assistant
system_message = "You are an expert product description writer for Amazon."

# User prompt that combines the user query and the schema
user_prompt = """Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
Only return description. The description should be SEO optimized and for a better mobile search experience.

<PRODUCT>
{product}
</PRODUCT>

<CATEGORY>
{category}
</CATEGORY>
"""

# Convert dataset to OAI messages
def format_data(sample):
    return {
        "messages": [
            {
                "role": "system",
                #"content": [{"type": "text", "text": system_message}],
                "content": system_message,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_prompt.format(
                            product=sample["Product Name"],
                            category=sample["Category"],
                        ),
                    },
                    {
                        "type": "image",
                        "image": sample["image"],
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": sample["description"]}],
            },
        ],
    }

def process_vision_info(messages: list[dict]) -> list[Image.Image]:
    image_inputs = []
    # Iterate through each conversation
    for msg in messages:
        # Get content (ensure it's a list)
        content = msg.get("content", [])
        if not isinstance(content, list):
            content = [content]

        # Check each content element for images
        for element in content:
            if isinstance(element, dict) and (
                "image" in element or element.get("type") == "image"
            ):
                # Get the image and convert to RGB
                if "image" in element:
                    image = element["image"]
                else:
                    image = element
                image_inputs.append(image.convert("RGB"))
    return image_inputs

# Load dataset from the hub
dataset = load_dataset("philschmid/amazon-product-descriptions-vlm", split="train")
dataset = dataset.train_test_split(test_size=0.1)

# Convert dataset to OAI messages
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
dataset_train = [format_data(sample) for sample in dataset["train"]]
dataset_test = [format_data(sample) for sample in dataset["test"]]

print(dataset_train[345]["messages"])

README.md: 0.00B [00:00, ?B/s]
data/train-00000-of-00001.parquet:   0%|          | 0.00/47.6M [00:00<?, ?B/s]
Generating train split:   0%|          | 0/1345 [00:00<?, ? examples/s]
[{'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!'}]}]

Optimiza Gemma con TRL y SFTTrainer

Ahora ya puedes ajustar tu modelo. El SFTTrainer de Hugging Face TRL facilita la supervisión del ajuste fino de los LLM abiertos. SFTTrainer es una subclase de Trainer de la biblioteca transformers y admite las mismas funciones, como el registro, la evaluación y la creación de puntos de control, pero agrega funciones adicionales para mejorar la calidad de vida, como las siguientes:

  • Formato de conjuntos de datos, incluidos los formatos conversacionales y de instrucciones
  • Solo se entrena en función de las finalizaciones, sin tener en cuenta las instrucciones.
  • Cómo empaquetar conjuntos de datos para un entrenamiento más eficiente
  • Compatibilidad con el ajuste eficiente de parámetros (PEFT), incluido QLoRA
  • Preparar el modelo y el tokenizador para el ajuste conversacional (por ejemplo, agregar tokens especiales)

El siguiente código carga el modelo y el tokenizador de Gemma desde Hugging Face, y luego inicializa la configuración de cuantización.

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-4-E2B" # @param ["google/gemma-4-E2B","google/gemma-4-E4B"] {"allow-input":true}

# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] < 8:
    raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")

# Define model init arguments
model_kwargs = dict(
    dtype=torch.bfloat16, # What torch dtype to use, defaults to auto
    device_map="auto", # Let torch decide how to load the model
)

# BitsAndBytesConfig int-4 config
model_kwargs["quantization_config"] = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=model_kwargs["dtype"],
    bnb_4bit_quant_storage=model_kwargs["dtype"],
)

# Load model and tokenizer
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
processor = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") # Load the Instruction Tokenizer to use the official Gemma template

config.json: 0.00B [00:00, ?B/s]
model.safetensors:   0%|          | 0.00/10.2G [00:00<?, ?B/s]
Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
generation_config.json:   0%|          | 0.00/149 [00:00<?, ?B/s]
processor_config.json: 0.00B [00:00, ?B/s]
chat_template.jinja: 0.00B [00:00, ?B/s]
config.json: 0.00B [00:00, ?B/s]
tokenizer_config.json: 0.00B [00:00, ?B/s]
tokenizer.json:   0%|          | 0.00/32.2M [00:00<?, ?B/s]

SFTTrainer admite una integración incorporada con peft, lo que facilita el ajuste eficiente de LLMs con QLoRA. Solo necesitas crear un LoraConfig y proporcionárselo al entrenador.

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=["lm_head", "embed_tokens"], # make sure to save the lm_head and embed_tokens as you train the special tokens
    ensure_weight_tying=True,
)

Antes de comenzar el entrenamiento, debes definir el hiperparámetro que deseas usar en un SFTConfig y un collate_fn personalizado para controlar el procesamiento de visión. El collate_fn convierte los mensajes con texto e imágenes en un formato que el modelo puede comprender.

from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-product-description",     # directory to save and repository id
    num_train_epochs=3,                         # number of training epochs
    per_device_train_batch_size=1,              # batch size per device during training
    optim="adamw_torch_fused",                  # use fused adamw optimizer
    logging_steps=5,                            # log every 5 steps
    save_strategy="epoch",                      # save checkpoint every epoch
    eval_strategy="epoch",                      # evaluate checkpoint every epoch
    learning_rate=2e-4,                         # learning rate, based on QLoRA paper
    bf16=True,                                  # use bfloat16 precision
    max_grad_norm=0.3,                          # max gradient norm based on QLoRA paper
    lr_scheduler_type="constant",               # use constant learning rate scheduler
    push_to_hub=True,                           # push model to hub
    report_to="tensorboard",                    # report metrics to tensorboard
    dataset_text_field="",                      # need a dummy field for collator
    dataset_kwargs={"skip_prepare_dataset": True}, # important for collator
    remove_unused_columns = False               # important for collator
)

# Create a data collator to encode text and image pairs
def collate_fn(examples):
    texts = []
    images = []
    for example in examples:
        image_inputs = process_vision_info(example["messages"])
        text = processor.apply_chat_template(
            example["messages"], add_generation_prompt=False, tokenize=False
        )
        texts.append(text.strip())
        images.append(image_inputs)

    # Tokenize the texts and process the images
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)

    # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation
    labels = batch["input_ids"].clone()

    # Mask tokens for not being used in the loss computation
    labels[labels == processor.tokenizer.pad_token_id] = -100
    labels[labels == processor.tokenizer.boi_token_id] = -100
    labels[labels == processor.tokenizer.image_token_id] = -100
    labels[labels == processor.tokenizer.eoi_token_id] = -100

    batch["labels"] = labels
    return batch

Ahora tienes todos los componentes básicos que necesitas para crear tu SFTTrainer y comenzar el entrenamiento de tu modelo.

from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset_train,
    eval_dataset=dataset_test,
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)

Para comenzar el entrenamiento, llama al método train().

# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()

# Save the final model again to the Hugging Face Hub
trainer.save_model()

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}.

Antes de probar el modelo, asegúrate de liberar la memoria.

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

Cuando usas QLoRA, solo entrenas adaptadores y no el modelo completo. Esto significa que, cuando guardes el modelo durante el entrenamiento, solo guardarás los pesos del adaptador y no el modelo completo. Si deseas guardar el modelo completo, lo que facilita su uso con pilas de entrega como vLLM o TGI, puedes combinar los pesos del adaptador con los pesos del modelo usando el método merge_and_unload y, luego, guardar el modelo con el método save_pretrained. Esto guarda un modelo predeterminado que se puede usar para la inferencia.

from peft import PeftModel

# Load Model base model
model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)

# Merge LoRA and base model and save
peft_model = PeftModel.from_pretrained(model, args.output_dir)
merged_model = peft_model.merge_and_unload()
merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")

processor = AutoProcessor.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")

Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
Writing model shards:   0%|          | 0/5 [00:00<?, ?it/s]
['merged_model/processor_config.json']

Probar la inferencia del modelo y generar descripciones de productos

Una vez que finalice el entrenamiento, querrás evaluar y probar tu modelo. Puedes cargar diferentes muestras del conjunto de datos de prueba y evaluar el modelo en esas muestras.

model_id = "merged_model"

# Load Model with PEFT adapter
model = AutoModelForImageTextToText.from_pretrained(
  model_id,
  device_map="auto",
  dtype="auto",
)
processor = AutoProcessor.from_pretrained(model_id)

Loading weights:   0%|          | 0/2012 [00:00<?, ?it/s]
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.

Para probar la inferencia, proporciona un nombre, una categoría y una imagen del producto. El sample incluye una figura de acción de Marvel.

import requests
from PIL import Image

# Test sample with Product Name, Category and Image
sample = {
  "product_name": "Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur",
  "category": "Toys & Games | Toy Figures & Playsets | Action Figures",
  "image": Image.open(requests.get("https://m.media-amazon.com/images/I/81+7Up7IWyL._AC_SY300_SX300_.jpg", stream=True).raw).convert("RGB")
}

def generate_description(sample, model, processor):
    # Convert sample into messages and then apply the chat template
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": [
            {"type": "image","image": sample["image"]},
            {"type": "text", "text": user_prompt.format(product=sample["product_name"], category=sample["category"])},
        ]},
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    print(text)
    # Process the image and text
    image_inputs = process_vision_info(messages)
    # Tokenize the text and process the images
    inputs = processor(
        text=[text],
        images=image_inputs,
        padding=True,
        return_tensors="pt",
    )
    # Move the inputs to the device
    inputs = inputs.to(model.device)

    # Generate the output
    stop_token_ids = [processor.tokenizer.eos_token_id, processor.tokenizer.convert_tokens_to_ids("<turn|>")]
    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)
    # Trim the generation and decode the output to text
    generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return output_text[0]

# generate the description
description = generate_description(sample, model, processor)
print("MODEL OUTPUT>> \n")
print(description)

<bos><|turn>system
You are an expert product description writer for Amazon.<turn|>
<|turn>user


<|image|>

Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
Only return description. The description should be SEO optimized and for a better mobile search experience.

<PRODUCT>
Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur
</PRODUCT>

<CATEGORY>
Toys & Games | Toy Figures & Playsets | Action Figures
</CATEGORY><turn|>
<|turn>model

MODEL OUTPUT>> 

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!

Resumen y próximos pasos

En este instructivo, se explicó cómo ajustar un modelo de Gemma para tareas de visión con TRL y QLoRA, específicamente para generar descripciones de productos. A continuación, consulta los siguientes documentos: