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

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érer des descriptions de produits) à l'aide de Transformers et de TRL de Hugging Face. Vous allez découvrir comment :

  • Qu'est-ce que l'adaptation à rang faible quantifiée (QLoRA) ?
  • Configurer l'environnement de développement
  • Créer et préparer l'ensemble de données de réglage fin pour les tâches de vision
  • Affiner Gemma à l'aide de TRL et de SFTTrainer
  • Testez l'inférence de 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 présente l'utilisation de l'adaptation faible à quantification (QLoRA), qui s'est imposée comme une méthode populaire pour affiner efficacement les LLM, car elle réduit les exigences 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 poids sont figés. Des couches d'adaptateur enregistrables (LoRA) sont ensuite associées, et seules les couches d'adaptateur sont entraînées. Ensuite, les poids de l'adaptateur peuvent être fusionnés avec le modèle de base ou conservés 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>=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

Avant de pouvoir commencer l'entraînement, vous devez vous assurer d'avoir accepté les conditions d'utilisation de Gemma. Vous pouvez accepter la licence sur Hugging Face en cliquant sur le bouton "Accepter et accéder au dépôt" sur la page du modèle à l'adresse http://huggingface.co/google/gemma-3-4b-pt (ou sur la page du modèle appropriée pour le modèle Gemma compatible avec la vision que vous utilisez).

Une fois que vous avez accepté la licence, vous avez besoin d'un jeton Hugging Face valide pour accéder au modèle. Si vous exécutez votre code dans un Google Colab, vous pouvez utiliser votre jeton Hugging Face de manière sécurisée à l'aide de secrets Colab. Sinon, vous pouvez définir le jeton directement dans la méthode login. Assurez-vous également que votre jeton dispose d'un accès en écriture, car vous envoyez votre modèle au hub pendant l'entraînement.

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)

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

Lorsque vous affinez 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 à la planche à dessin.

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

  • Affiner un modèle Gemma pour générer des descriptions de produits concises et optimisées pour le référencement naturel pour une plate-forme d'e-commerce, spécifiquement adaptées à la recherche sur 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 d'ensembles de données Hugging Face pour charger l'ensemble de données et créer un modèle d'invite pour combiner l'image, le nom du produit et la catégorie, et ajouter un message système. L'ensemble de données inclut des images en tant qu'objets 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"])

Affiner Gemma à l'aide de TRL et de SFTTrainer

Vous êtes maintenant prêt à affiner votre modèle. Le SFTTrainer de Hugging Face TRL permet de superviser facilement l'ajustement fin des LLM ouverts. SFTTrainer est une sous-classe de Trainer de la bibliothèque transformers. Elle est compatible avec toutes les mêmes fonctionnalités, y compris la journalisation, l'évaluation et le point de contrôle, mais ajoute des fonctionnalités supplémentaires pour améliorer la qualité de vie, y compris les suivantes:

  • Mise en forme des ensembles de données, y compris les formats conversationnel et d'instruction
  • Entraînement sur les finalisations uniquement, en ignorant les requêtes
  • Empaqueter des ensembles de données pour un entraînement plus efficace
  • Prise en charge de l'optimisation du réglage des paramètres (PEFT, Parameter-efficient Fine-Tuning), y compris QloRA
  • Préparer le modèle et le tokenizer pour le réglage fin de la conversation (par exemple, en ajoutant des jetons spéciaux)

Le code suivant charge le modèle et le tokenizer Gemma à partir de Hugging Face, puis initialise la configuration de quantification.

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 est compatible avec une intégration intégrée à peft, ce qui permet de régler facilement et 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",
    ],
)

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 visuel. 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=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

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

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)

Démarrez 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()

Avant de pouvoir tester votre modèle, assurez-vous de libérer 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 poids de l'adaptateur et non le modèle complet. Si vous souhaitez enregistrer le modèle complet, ce qui le rend plus facile à utiliser avec des piles de diffusion telles que vLLM ou TGI, vous pouvez fusionner les poids de l'adaptateur dans les poids 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")

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.

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)

Vous pouvez tester l'inférence en fournissant un nom, une catégorie et une image de produit. 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": [{"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)

Résumé et prochaines étapes

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