Affiner Gemma pour les tâches de vision à l'aide de Hugging Face Transformers et de QLoRA

 Afficher sur ai.google.dev  Exécuter dans Google Colab Exécuter dans Kaggle Ouvrir dans Vertex AI  Afficher la source sur GitHub

Ce guide vous explique comment affiner Gemma sur un ensemble de données d'images et de texte personnalisé pour une tâche de vision (génération de descriptions de produits) à l'aide de Transformers et TRL de Hugging Face. Vous pourrez découvrir :

  • Qu'est-ce que l'adaptation à faible rang quantifiée (QLoRA) ?
  • Configurer l'environnement de développement
  • Créer et préparer l'ensemble de données de réglage fin
  • Affiner Gemma à l'aide de TRL et de SFTTrainer
  • Testez l'inférence du modèle et générez des descriptions de produits à partir d'images et de texte.

Qu'est-ce que l'adaptation à faible rang quantifiée (QLoRA) ?

Ce guide explique comment utiliser Quantized Low-Rank Adaptation (QLoRA), une méthode populaire pour affiner efficacement les LLM, car elle réduit les besoins en ressources de calcul tout en maintenant des performances élevées. Dans QloRA, le modèle pré-entraîné est quantifié à 4 bits et les pondérations sont figées. Des calques d'adaptateur entraînables (LoRA) sont ensuite associés, et seuls les calques d'adaptateur sont entraînés. Les pondérations de l'adaptateur peuvent ensuite être fusionnées avec le modèle de base ou conservées en tant qu'adaptateur distinct.

Configurer l'environnement de développement

La première étape consiste à installer les bibliothèques Hugging Face, y compris TRL et les ensembles de données, pour affiner le modèle ouvert.

# 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

Remarque : Si vous utilisez un GPU avec une architecture Ampere (comme NVIDIA L4) ou plus récente, vous pouvez utiliser Flash Attention. Flash Attention est une méthode qui accélère considérablement les calculs et réduit l'utilisation de la mémoire de quadratique à linéaire en termes de longueur de séquence, ce qui permet d'accélérer l'entraînement jusqu'à trois fois. Pour en savoir plus, consultez FlashAttention.

Pour publier votre modèle, vous avez besoin d'un jeton Hugging Face valide. Si vous exécutez le code dans Google Colab, vous pouvez utiliser votre jeton Hugging Face de manière sécurisée à l'aide des secrets Colab. Sinon, vous pouvez définir le jeton directement dans la méthode login. Assurez-vous que votre jeton dispose également d'un accès en écriture, car vous allez transférer votre modèle vers le Hub pendant l'entraînement.

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

Créer et préparer l'ensemble de données de réglage fin

Lorsque vous ajustez des LLM, il est important de connaître votre cas d'utilisation et la tâche que vous souhaitez résoudre. Cela vous aide à créer un ensemble de données pour affiner votre modèle. Si vous n'avez pas encore défini votre cas d'utilisation, vous pouvez revenir en arrière.

Par exemple, ce guide se concentre sur le cas d'utilisation suivant :

  • Affinage d'un modèle Gemma pour générer des descriptions de produits concises et optimisées pour le SEO pour une plate-forme d'e-commerce, en particulier pour la recherche mobile.

Ce guide utilise l'ensemble de données philschmid/amazon-product-descriptions-vlm, qui contient des descriptions de produits Amazon, y compris des images et des catégories de produits.

Hugging Face TRL est compatible avec les conversations multimodales. L'élément important est le rôle "image", qui indique à la classe de traitement qu'elle doit charger l'image. La structure doit être la suivante :

{"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": "..."}]}]}

Vous pouvez désormais utiliser la bibliothèque Hugging Face Datasets pour charger l'ensemble de données et créer un modèle d'invite permettant de combiner l'image, le nom et la catégorie du produit, et d'ajouter un message système. L'ensemble de données inclut des images en tant qu'objetsPil.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!'}]}]

Affiner Gemma à l'aide de TRL et de SFTTrainer

Vous êtes maintenant prêt à affiner votre modèle. Hugging Face TRL SFTTrainer facilite le réglage fin supervisé des LLM Open Source. SFTTrainer est une sous-classe de Trainer de la bibliothèque transformers et prend en charge les mêmes fonctionnalités, y compris la journalisation, l'évaluation et la création de points de contrôle, mais ajoute des fonctionnalités supplémentaires de qualité de vie, y compris :

  • Mise en forme des ensembles de données, y compris les formats conversationnels et d'instructions
  • Entraînement sur les complétions uniquement, en ignorant les requêtes
  • Regrouper les ensembles de données pour un entraînement plus efficace
  • Prise en charge de l'affinage des paramètres avec optimisation (PEFT), y compris QLoRA
  • Préparer le modèle et le tokenizer pour l'affinage conversationnel (par exemple, en ajoutant des jetons spéciaux)

Le code suivant charge le modèle et le tokenizer Gemma depuis Hugging Face et initialise la configuration de quantification.

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 est compatible avec une intégration intégrée à peft, ce qui permet d'ajuster efficacement les LLM à l'aide de QLoRA. Il vous suffit de créer un LoraConfig et de le fournir au formateur.

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,
)

Avant de pouvoir commencer votre entraînement, vous devez définir l'hyperparamètre que vous souhaitez utiliser dans un SFTConfig et un collate_fn personnalisé pour gérer le traitement de la vision. collate_fn convertit les messages contenant du texte et des images dans un format que le modèle peut comprendre.

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

Vous disposez désormais de tous les éléments nécessaires pour créer votre SFTTrainer et commencer l'entraînement de votre modèle.

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,
)

Commencez l'entraînement en appelant la méthode 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}.

Avant de pouvoir tester votre modèle, assurez-vous de libérer de la mémoire.

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

Lorsque vous utilisez QLoRA, vous n'entraînez que les adaptateurs et non le modèle complet. Cela signifie que, lorsque vous enregistrez le modèle pendant l'entraînement, vous n'enregistrez que les pondérations de l'adaptateur et non le modèle complet. Si vous souhaitez enregistrer le modèle complet, ce qui facilite son utilisation avec des piles de diffusion telles que vLLM ou TGI, vous pouvez fusionner les pondérations de l'adaptateur dans les pondérations du modèle à l'aide de la méthode merge_and_unload, puis enregistrer le modèle avec la méthode save_pretrained. Cela permet d'enregistrer un modèle par défaut, qui peut être utilisé pour l'inférence.

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']

Tester l'inférence du modèle et générer des descriptions de produits

Une fois l'entraînement terminé, vous devez évaluer et tester votre modèle. Vous pouvez charger différents échantillons à partir de l'ensemble de données de test et évaluer le modèle sur ces échantillons.

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.

Vous pouvez tester l'inférence en fournissant un nom de produit, une catégorie et une image. Le sample inclut une figurine 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!

Résumé et étapes suivantes

Ce tutoriel vous a montré comment affiner un modèle Gemma pour des tâches de vision à l'aide de TRL et QLoRA, en particulier pour générer des descriptions de produits. Consultez ensuite les documents suivants :