在自定义数据集上微调 Gemma 4 的全流程指南

import json

# 示例：宠物专家微调数据集
examples = [
    {
        "messages": [
            {"role": "system", "content": "你是一位资深的宠物行为专家。"},
            {"role": "user", "content": "为什么我的猫晚上总是乱叫？"},
            {"role": "assistant", "content": "猫在夜间鸣叫通常有三个原因：1. 天生的捕猎本能；2. 精力过剩；3. 寻求关注。建议在睡前进行 15 分钟的互动游戏，并固定喂食时间。"}
        ]
    }
]

# 保存为 JSONL
with open("train_data.jsonl", "w", encoding="utf-8") as f:
    for ex in examples:
        f.write(json.dumps(ex, ensure_ascii=False) + "\n")

pip install torch>=2.2.0 transformers>=4.40.0 trl>=0.8.0 peft>=0.10.0 accelerate bitsandbytes

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer
from datasets import load_dataset

# 配置 4-bit 量化以节省显存
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-9b-it",
    quantization_config=bnb_config,
    device_map="auto"
)

# 定义 LoRA 参数
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj"],
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)

# 启动训练器
trainer = SFTTrainer(
    model=model,
    train_dataset=load_dataset("json", data_files="train_data.jsonl", split="train"),
    max_seq_length=2048,
    args=TrainingArguments(
        output_dir="./gemma-4-output",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        num_train_epochs=3,
        fp16=True,
        logging_steps=10,
        optim="paged_adamw_8bit"
    )
)

trainer.train()

from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained("google/gemma-4-9b-it")
model = PeftModel.from_pretrained(base_model, "./gemma-4-output")
merged_model = model.merge_and_unload()
merged_model.save_pretrained("./gemma-4-final-model")