En esta guía, se explica cómo ajustar Gemma en un conjunto de datos de texto e imágenes personalizado 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 fino para tareas de visión
- Optimiza Gemma con TRL y SFTTrainer
- Prueba la inferencia de modelos y genera 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 muestra el uso de la adaptación de clasificación baja cuantificada (QLoRA), que surgió como un método popular para ajustar con eficiencia los LLM, ya que reduce los requisitos de recursos computacionales y, al mismo tiempo, mantiene un alto rendimiento. En QloRA, el modelo previamente entrenado se cuantifica a 4 bits y los pesos se congelan. Luego, se adjuntan capas de adaptadores entrenables (LoRA) y solo se entrenan las capas de adaptadores. Luego, las ponderaciones del adaptador se pueden combinar con el modelo de base o mantenerse como un adaptador independiente.
Configura el entorno de desarrollo
El primer paso es instalar las bibliotecas de Hugging Face, incluida la TRL, y los conjuntos de datos para ajustar el modelo abierto.
# Install Pytorch & other libraries
%pip install "torch>=2.4.0" tensorboard torchvision
# Install Gemma release branch from Hugging Face
%pip install git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3
# Install Hugging Face libraries
%pip install --upgrade \
"datasets==3.3.2" \
"accelerate==1.4.0" \
"evaluate==0.4.3" \
"bitsandbytes==0.45.3" \
"trl==0.15.2" \
"peft==0.14.0" \
"pillow==11.1.0" \
protobuf \
sentencepiece
Antes de comenzar a entrenar, debes asegurarte de haber aceptado las condiciones de uso de Gemma. Para aceptar la licencia en Hugging Face, haz clic en el botón Aceptar y acceder al repositorio en la página del modelo, en http://huggingface.co/google/gemma-3-4b-pt (o la página del modelo adecuada para el modelo de Gemma compatible con la visión que estés usando).
Después de aceptar la licencia, necesitas un token de Hugging Face válido para acceder al modelo. Si ejecutas el código en Google Colab, puedes usar de forma segura tu token de Hugging Face con los secretos de Colab. De lo contrario, puedes configurar el token directamente en el método login. Asegúrate de que tu token también tenga acceso de escritura cuando envíes tu modelo al Hub durante el entrenamiento.
from google.colab import userdata
from huggingface_hub import login
# Login into Hugging Face Hub
hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab
login(hf_token)
Crea y prepara el conjunto de datos de ajuste
Cuando ajustas los LLM, es importante conocer tu caso de uso y la tarea que quieres resolver. Esto te ayuda a crear un conjunto de datos para ajustar tu modelo. Si aún no definiste tu caso de uso, te recomendamos que vuelvas al principio.
A modo de 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 para una plataforma de comercio electrónico, específicamente adaptadas para la búsqueda en dispositivos móviles
En esta guía, se usa el conjunto de datos philschmid/amazon-product-descriptions-vlm, que incluye descripciones de productos de Amazon, imágenes y categorías.
El TRL de Hugging Face admite conversaciones multimodales. El elemento importante es el rol "image", que le indica a la clase de procesamiento que debe cargar la imagen. La estructura debe seguir lo siguiente:
{"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 conjuntos de datos de Hugging Face para cargar el conjunto de datos y crear una plantilla de instrucción 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 objetos Pil.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}],
},
{
"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")
# Convert dataset to OAI messages
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
dataset = [format_data(sample) for sample in dataset]
print(dataset[345]["messages"])
Optimiza Gemma con TRL y SFTTrainer
Ya puedes ajustar tu modelo. SFTTrainer de TRL de Hugging Face permite supervisar y ajustar con precisión los LLM abiertos de forma sencilla. SFTTrainer es una subclase de Trainer de la biblioteca transformers y admite las mismas funciones, como el registro, la evaluación y los puntos de control, pero agrega funciones adicionales de calidad de vida, como las siguientes:
- Aplicación de formato a conjuntos de datos, incluidos los formatos de instrucción y de conversación
- Entrenamiento solo en las acciones de finalización, sin tener en cuenta las instrucciones
- Empaquetamiento de 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 analizador para el ajuste fino de la conversación (como agregar tokens especiales)
El siguiente código carga el modelo y el analizador de Gemma desde Hugging Face y, luego, inicializa la configuración de cuantificación.
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
# Hugging Face model id
model_id = "google/gemma-3-4b-pt" # or `google/gemma-3-12b-pt`, `google/gemma-3-27-pt`
# 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(
attn_implementation="eager", # Use "flash_attention_2" when running on Ampere or newer GPU
torch_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["torch_dtype"],
bnb_4bit_quant_storage=model_kwargs["torch_dtype"],
)
# Load model and tokenizer
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
processor = AutoProcessor.from_pretrained("google/gemma-3-4b-it")
SFTTrainer admite una integración integrada con peft, lo que facilita la configuración eficiente de los LLM con QLoRA. Solo debes crear un LoraConfig y proporcionarlo al capacitador.
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",
],
)
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 la visión. collate_fn convierte los mensajes con texto e imágenes en un formato que el modelo puede entender.
from trl import SFTConfig
args = SFTConfig(
output_dir="gemma-product-description", # directory to save and repository id
num_train_epochs=1, # number of training epochs
per_device_train_batch_size=1, # batch size per device during training
gradient_accumulation_steps=4, # number of steps before performing a backward/update pass
gradient_checkpointing=True, # use gradient checkpointing to save memory
optim="adamw_torch_fused", # use fused adamw optimizer
logging_steps=5, # log every 5 steps
save_strategy="epoch", # save 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
warmup_ratio=0.03, # warmup ratio 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
gradient_checkpointing_kwargs={
"use_reentrant": False
}, # use reentrant checkpointing
dataset_text_field="", # need a dummy field for collator
dataset_kwargs={"skip_prepare_dataset": True}, # important for collator
)
args.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 image tokens
image_token_id = [
processor.tokenizer.convert_tokens_to_ids(
processor.tokenizer.special_tokens_map["boi_token"]
)
]
# Mask tokens for not being used in the loss computation
labels[labels == processor.tokenizer.pad_token_id] = -100
labels[labels == image_token_id] = -100
labels[labels == 262144] = -100
batch["labels"] = labels
return batch
Ahora tienes todos los componentes básicos que necesitas para crear tu SFTTrainer y comenzar a entrenar el modelo.
from trl import SFTTrainer
trainer = SFTTrainer(
model=model,
args=args,
train_dataset=dataset,
peft_config=peft_config,
processing_class=processor,
data_collator=collate_fn,
)
Para iniciar 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()
Antes de que puedas probar tu 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 guardas el modelo durante el entrenamiento, solo guardas las ponderaciones 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 con 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")
Prueba la inferencia de modelos y genera descripciones de productos
Una vez que se complete el entrenamiento, deberás evaluar y probar tu modelo. Puedes cargar diferentes muestras del conjunto de datos de prueba y evaluar el modelo en ellas.
import torch
# Load Model with PEFT adapter
model = AutoModelForImageTextToText.from_pretrained(
args.output_dir,
device_map="auto",
torch_dtype=torch.bfloat16,
attn_implementation="eager",
)
processor = AutoProcessor.from_pretrained(args.output_dir)
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": [{"type": "text", "text": 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
)
# 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("<end_of_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(description)
Resumen y próximos pasos
En este instructivo, se analizó cómo ajustar un modelo de Gemma para tareas de visión con TRL y QLoRA, específicamente para generar descripciones de productos. Consulta los siguientes documentos:
- Obtén más información para generar texto con un modelo de Gemma.
- Obtén información para ajustar Gemma para tareas de texto con Hugging Face Transformers.
- Obtén información para realizar ajustes y una inferencia distribuidos en un modelo de Gemma.
- Obtén información para usar modelos abiertos de Gemma con Vertex AI.
- Aprende a optimizar Gemma con KerasNLP y, luego, impleméntalo en Vertex AI.