Tinh chỉnh Gemma cho các tác vụ thị giác bằng cách sử dụng Hugging Face Transformers và QLoRA

Hướng dẫn này sẽ hướng dẫn bạn cách tinh chỉnh Gemma trên một tập dữ liệu hình ảnh và văn bản tuỳ chỉnh cho một nhiệm vụ thị giác (tạo nội dung mô tả sản phẩm) bằng cách sử dụng TransformerTRL của Hugging Face. Bạn sẽ tìm hiểu:

  • Điều chỉnh thứ hạng thấp được lượng tử hoá (QLoRA) là gì?
  • Thiết lập môi trường phát triển
  • Tạo và chuẩn bị tập dữ liệu tinh chỉnh cho các tác vụ thị giác
  • Tinh chỉnh Gemma bằng TRL và SFTTrainer
  • Kiểm thử tính năng Suy luận mô hình và tạo nội dung mô tả sản phẩm từ hình ảnh và văn bản.

Điều chỉnh thứ hạng thấp được lượng tử hoá (QLoRA) là gì?

Hướng dẫn này minh hoạ cách sử dụng Tính năng thích ứng theo thứ hạng thấp được lượng tử hoá (QLoRA). Đây là một phương pháp phổ biến để tinh chỉnh hiệu quả các LLM vì phương pháp này làm giảm các yêu cầu về tài nguyên điện toán trong khi vẫn duy trì hiệu suất cao. Trong QloRA, mô hình được huấn luyện trước được lượng tử hoá thành 4 bit và các trọng số được cố định. Sau đó, các lớp chuyển đổi có thể huấn luyện (LoRA) sẽ được đính kèm và chỉ các lớp chuyển đổi mới được huấn luyện. Sau đó, bạn có thể hợp nhất trọng số của bộ chuyển đổi với mô hình cơ sở hoặc giữ lại dưới dạng một bộ chuyển đổi riêng biệt.

Thiết lập môi trường phát triển

Bước đầu tiên là cài đặt Thư viện Hugging Face, bao gồm cả TRL và các tập dữ liệu để tinh chỉnh mô hình mở.

# Install Pytorch & other libraries
%pip install "torch>=2.4.0" tensorboard torchvision

# Install Gemma release branch from Hugging Face
%pip install "transformers>=4.51.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

Trước khi bắt đầu huấn luyện, bạn phải đảm bảo rằng bạn đã chấp nhận điều khoản sử dụng của Gemma. Bạn có thể chấp nhận giấy phép trên Hugging Face bằng cách nhấp vào nút Đồng ý và truy cập vào kho lưu trữ trên trang mô hình tại: http://huggingface.co/google/gemma-3-4b-pt (hoặc trang mô hình thích hợp cho mô hình Gemma có khả năng thị giác mà bạn đang sử dụng).

Sau khi chấp nhận giấy phép, bạn cần có Mã thông báo Hugging Face hợp lệ để truy cập vào mô hình. Nếu đang chạy bên trong Google Colab, bạn có thể sử dụng mã thông báo Hugging Face một cách an toàn bằng cách sử dụng các khoá Colab; nếu không, bạn có thể đặt mã thông báo trực tiếp trong phương thức login. Đảm bảo mã thông báo của bạn cũng có quyền ghi khi bạn đẩy mô hình vào Trung tâm trong quá trình huấn luyện.

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)

Tạo và chuẩn bị tập dữ liệu tinh chỉnh

Khi tinh chỉnh LLM, điều quan trọng là bạn phải biết trường hợp sử dụng và nhiệm vụ bạn muốn giải quyết. Điều này giúp bạn tạo một tập dữ liệu để tinh chỉnh mô hình. Nếu chưa xác định trường hợp sử dụng, bạn nên quay lại bảng vẽ.

Ví dụ: hướng dẫn này tập trung vào trường hợp sử dụng sau:

  • Điều chỉnh mô hình Gemma để tạo nội dung mô tả sản phẩm ngắn gọn, được tối ưu hoá cho SEO cho một nền tảng thương mại điện tử, được điều chỉnh riêng cho hoạt động tìm kiếm trên thiết bị di động.

Hướng dẫn này sử dụng tập dữ liệu philschmid/amazon-product-descriptions-vlm, một tập dữ liệu về nội dung mô tả sản phẩm trên Amazon, bao gồm cả hình ảnh và danh mục sản phẩm.

Hugging Face TRL hỗ trợ các cuộc trò chuyện đa phương thức. Phần quan trọng là vai trò "image" (hình ảnh), cho lớp xử lý biết rằng lớp này sẽ tải hình ảnh. Cấu trúc phải tuân theo:

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

Giờ đây, bạn có thể sử dụng thư viện Tập dữ liệu Hugging Face để tải tập dữ liệu và tạo mẫu lời nhắc để kết hợp hình ảnh, tên sản phẩm và danh mục, đồng thời thêm thông báo hệ thống. Tập dữ liệu này bao gồm hình ảnh dưới dạng đối tượng 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"])

Tinh chỉnh Gemma bằng TRL và SFTTrainer

Giờ thì bạn đã sẵn sàng tinh chỉnh mô hình. SFTTrainer của Hugging Face TRL giúp bạn dễ dàng giám sát việc tinh chỉnh các LLM mở. SFTTrainer là lớp con của Trainer trong thư viện transformers và hỗ trợ tất cả các tính năng tương tự, bao gồm cả tính năng ghi nhật ký, đánh giá và kiểm tra điểm, nhưng bổ sung thêm các tính năng chất lượng cuộc sống, bao gồm:

  • Định dạng tập dữ liệu, bao gồm cả định dạng trò chuyện và hướng dẫn
  • Chỉ huấn luyện về các lượt hoàn thành, bỏ qua lời nhắc
  • Gói tập dữ liệu để huấn luyện hiệu quả hơn
  • Hỗ trợ tinh chỉnh hiệu quả theo tham số (PEFT) bao gồm cả QloRA
  • Chuẩn bị mô hình và trình tạo mã thông báo để tinh chỉnh cuộc trò chuyện (chẳng hạn như thêm mã thông báo đặc biệt)

Mã sau đây tải mô hình Gemma và trình tạo mã thông báo từ Hugging Face, đồng thời khởi chạy cấu hình lượng tử hoá.

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 hỗ trợ tích hợp sẵn với peft, giúp bạn dễ dàng điều chỉnh LLM một cách hiệu quả bằng QLoRA. Bạn chỉ cần tạo một LoraConfig và cung cấp cho trình huấn luyện.

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

Trước khi bắt đầu huấn luyện, bạn cần xác định tham số siêu dữ liệu mà bạn muốn sử dụng trong SFTConfigcollate_fn tuỳ chỉnh để xử lý quá trình xử lý hình ảnh. collate_fn chuyển đổi các thông báo có văn bản và hình ảnh thành định dạng mà mô hình có thể hiểu được.

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

Giờ đây, bạn đã có mọi thành phần cần thiết để tạo SFTTrainer nhằm bắt đầu huấn luyện mô hình.

from trl import SFTTrainer

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

Bắt đầu huấn luyện bằng cách gọi phương thức 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()

Trước khi kiểm thử mô hình, hãy nhớ giải phóng bộ nhớ.

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

Khi sử dụng QLoRA, bạn chỉ huấn luyện bộ chuyển đổi chứ không phải toàn bộ mô hình. Điều này có nghĩa là khi lưu mô hình trong quá trình huấn luyện, bạn chỉ lưu trọng số của bộ chuyển đổi chứ không phải toàn bộ mô hình. Nếu muốn lưu toàn bộ mô hình để dễ dàng sử dụng với các ngăn xếp phân phát như vLLM hoặc TGI, bạn có thể hợp nhất trọng số của bộ chuyển đổi vào trọng số của mô hình bằng phương thức merge_and_unload, sau đó lưu mô hình bằng phương thức save_pretrained. Thao tác này sẽ lưu một mô hình mặc định có thể dùng để suy luận.

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

Kiểm thử tính năng Suy luận mô hình và tạo nội dung mô tả sản phẩm

Sau khi quá trình huấn luyện hoàn tất, bạn nên đánh giá và kiểm thử mô hình. Bạn có thể tải nhiều mẫu từ tập dữ liệu kiểm thử và đánh giá mô hình trên các mẫu đó.

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)

Bạn có thể thử nghiệm tính năng suy luận bằng cách cung cấp tên sản phẩm, danh mục và hình ảnh. sample bao gồm một mô hình hành động 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)

Tóm tắt và các bước tiếp theo

Hướng dẫn này trình bày cách tinh chỉnh mô hình Gemma cho các tác vụ thị giác bằng TRL và QLoRA, cụ thể là để tạo nội dung mô tả sản phẩm. Tiếp theo, hãy xem các tài liệu sau: