Hugging Face Transformers と QLoRA を使用して Gemma をファインチューニングしてビジョン タスクを実行する

ai.google.dev で表示 Google Colab で実行 Kaggle で実行 Vertex AI で開く GitHub 上のソースを見る

このガイドでは、Hugging Face の TransformersTRL を使用して、ビジョン タスク(商品紹介文の生成)用にカスタム画像とテキストのデータセットで Gemma をファインチューニングする方法について説明します。学習内容:

  • 量子化 Low-Rank Adaptation(QLoRA)とは
  • 開発環境をセットアップする
  • ファインチューニング データセットを作成して準備する
  • TRL と SFTTrainer を使用して Gemma をファインチューニングする
  • モデルの推論をテストし、画像とテキストから商品説明を生成します。

量子化 Low-Rank Adaptation(QLoRA)とは

このガイドでは、量子化された低ランク適応(QLoRA)の使用方法について説明します。QLoRA は、高いパフォーマンスを維持しながらコンピューティング リソースの要件を削減するため、LLM を効率的にファインチューニングする一般的な方法として登場しました。QloRA では、事前トレーニング済みモデルが 4 ビットに量子化され、重みが固定されます。次に、トレーニング可能なアダプタレイヤ(LoRA)が接続され、アダプタレイヤのみがトレーニングされます。その後、アダプタの重みをベースモデルと統合するか、別のアダプタとして保持できます。

開発環境をセットアップする

まず、TRL やオープンモデルをファインチューニングするためのデータセットなど、Hugging Face ライブラリをインストールします。

# 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

注: Ampere アーキテクチャ(NVIDIA L4 など)以降の GPU を使用している場合は、Flash アテンションを使用できます。Flash Attention は、計算を大幅に高速化し、メモリ使用量をシーケンス長の 2 次関数から線形関数に削減する手法です。これにより、トレーニングを最大 3 倍高速化できます。詳しくは、FlashAttention をご覧ください。

モデルを公開するには、有効な Hugging Face トークンが必要です。Google Colab 内で実行している場合は、Colab シークレットを使用して Hugging Face トークンを安全に使用できます。それ以外の場合は、login メソッドでトークンを直接設定できます。トレーニング中にモデルを Hub に push するため、トークンに書き込みアクセス権があることを確認してください。

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

ファインチューニング データセットを作成して準備する

LLM をファインチューニングする際は、ユースケースと解決したいタスクを把握することが重要です。これにより、モデルをファインチューニングするためのデータセットを作成できます。ユースケースをまだ定義していない場合は、振り出しに戻ることをおすすめします。

このガイドでは、次のユースケースを例として説明します。

  • Gemma モデルをファインチューニングして、モバイル検索に特化した、簡潔で SEO 対策された e コマース プラットフォームの商品説明を生成します。

このガイドでは、商品画像やカテゴリを含む Amazon 商品説明のデータセットである philschmid/amazon-product-descriptions-vlm データセットを使用します。

Hugging Face TRL はマルチモーダルな会話をサポートしています。重要なのは「image」ロールです。これにより、処理クラスに画像を読み込むよう指示します。構造は次のとおりです。

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

Hugging Face Datasets ライブラリを使用してデータセットを読み込み、画像、商品名、カテゴリを組み合わせてシステム メッセージを追加するプロンプト テンプレートを作成できます。データセットには、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}],
                "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!'}]}]

TRL と SFTTrainer を使用して Gemma をファインチューニングする

これで、モデルをファインチューニングする準備が整いました。Hugging Face TRL の SFTTrainer を使用すると、オープン LLM の教師ありファインチューニングを簡単に行うことができます。SFTTrainertransformers ライブラリの Trainer のサブクラスであり、ロギング、評価、チェックポイントなどのすべての同じ機能をサポートしていますが、次のような QOL 機能をさらに追加しています。

  • 会話形式や指示形式などのデータセットの形式設定
  • プロンプトを無視して完了のみをトレーニングする
  • トレーニング効率を高めるためのデータセットのパッキング
  • QloRA を含むパラメータ エフィシエント ファインチューニング(PEFT)のサポート
  • 会話型ファインチューニング用のモデルとトークナイザーの準備(特殊トークンの追加など)

次のコードは、Hugging Face から Gemma モデルとトークナイザーを読み込み、量子化構成を初期化します。

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]

SFTTrainerpeft との組み込み統合をサポートしているため、QLoRA を使用して LLM を効率的にチューニングできます。LoraConfig を作成してトレーナーに提供するだけで済みます。

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

トレーニングを開始する前に、SFTConfig とカスタム collate_fn で使用するハイパーパラメータを定義して、ビジョン処理を処理する必要があります。collate_fn は、テキストと画像を含むメッセージをモデルが理解できる形式に変換します。

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

これで、モデルのトレーニングを開始するために SFTTrainer を作成するために必要な構成要素がすべて揃いました。

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

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

モデルをテストする前に、必ずメモリを解放してください。

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

QLoRA を使用する場合、アダプタのみをトレーニングし、モデル全体はトレーニングしません。つまり、トレーニング中にモデルを保存するときに、完全なモデルではなくアダプタの重みのみを保存します。モデル全体を保存して、vLLM や TGI などのサービング スタックで簡単に使用できるようにする場合は、merge_and_unload メソッドを使用してアダプタの重みをモデルの重みに統合し、save_pretrained メソッドでモデルを保存します。これにより、推論に使用できるデフォルトのモデルが保存されます。

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

モデルの推論をテストして商品説明を生成する

トレーニングが完了したら、モデルを評価してテストします。テスト データセットからさまざまなサンプルを読み込み、それらのサンプルでモデルを評価できます。

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.

商品名、カテゴリ、画像を指定して、推論をテストできます。sample にはマーベルのアクション フィギュアが含まれています。

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!

まとめと次のステップ

このチュートリアルでは、TRL と QLoRA を使用して、特に商品説明の生成を目的として、ビジョン タスク用に Gemma モデルをファインチューニングする方法について説明しました。次のドキュメントもご覧ください。