【Datawhale学习笔记】使用AMD GPU15分钟部署&运行Gemma4大模型

JeffDingAI

18人浏览 · 2026-06-13 19:06:22

JeffDingAI · 2026-06-13 19:06:22 发布

跑通交互式实验

安装依赖

uv pip install -U vllm modelscope transformers accelerate datasets trl peft scikit-learn pandas tqdm torchvision --no-cache -i https://mirrors.cloud.tencent.com/pypi/simple/ --extra-index-url https://wheels.vllm.ai/rocm/

导入依赖与全局配置

import os
import re
import json
import random
import warnings

import numpy as np
import pandas as pd
import torch

from tqdm.auto import tqdm
from datasets import Dataset, DatasetDict, ClassLabel, load_dataset
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score

from modelscope import snapshot_download
from modelscope.hub.snapshot_download import dataset_snapshot_download

from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
from peft import LoraConfig, PeftModel
from trl import SFTConfig, SFTTrainer

warnings.filterwarnings("ignore")

# -----------------------------
# 基础配置
# -----------------------------
# 魔搭上的模型 ID（Gemma 4 E4B-it 在 ModelScope 上的官方仓库,instruction-tuned 版本,
# 仓库内自带官方 chat_template.jinja,无需手动处理 chat template)。
# 仓库地址: https://www.modelscope.cn/models/google/gemma-4-E4B-it
MODELSCOPE_MODEL_ID = "google/gemma-4-E4B-it"

# 魔搭上的数据集 ID（dair-ai/emotion 在 ModelScope 上的官方镜像）。
MODELSCOPE_DATASET_ID = "AI-ModelScope/emotion"

# 微调输出目录
OUTPUT_DIR = "./gemma4-it-emotion-lora-ms-single-gpu"

# 数据量控制。先用小数据跑通，确认没问题后再加大。
TRAIN_LIMIT = 4000
VALIDATION_LIMIT = 400
TEST_LIMIT = 400
EVAL_LIMIT = 400

SEED = 42
MODEL_DTYPE = torch.bfloat16
BF16 = True
FP16 = False

SYSTEM_PROMPT = """You are an emotion classification assistant.
Read the user's text and answer with exactly one label.
Only choose from: sadness, joy, love, anger, fear, surprise.
Return only the label and nothing else."""

LABEL_PATTERN = re.compile(r"(sadness|joy|love|anger|fear|surprise)", re.IGNORECASE)

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs("./models", exist_ok=True)
os.makedirs("./datasets", exist_ok=True)

print("torch version:", torch.__version__)
print("torch.cuda.is_available():", torch.cuda.is_available())
print("torch.cuda.device_count():", torch.cuda.device_count())
if torch.cuda.is_available():
    print("current device:", torch.cuda.current_device())
    print("device name:", torch.cuda.get_device_name(0))

输出信息

torch version: 2.10.0+git8514f05
torch.cuda.is_available(): True
torch.cuda.device_count(): 1
current device: 0
device name: AMD Radeon Graphics

固定随机种子

def setup_seed(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    set_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

setup_seed(SEED)

下载 Gemma 模型

MODELSCOPE_MODEL_ID = "google/gemma-4-E4B-it"
print("Downloading model from ModelScope...")
print("ModelScope model id:", MODELSCOPE_MODEL_ID)

model_dir = snapshot_download(
    MODELSCOPE_MODEL_ID,
    cache_dir="./models",
)

print("Downloaded model dir:", model_dir)

# 后续统一使用本地路径加载
LOCAL_MODEL_DIR = model_dir

加载 AI-ModelScope/emotion 数据集

import glob

EMOTION_LABEL_NAMES = ["sadness", "joy", "love", "anger", "fear", "surprise"]


# 直接把魔搭上的数据集仓库(parquet 文件)整体下载到本地,然后用 datasets 库从本地 parquet 加载。
# 不走 MsDataset.load -> datasets.load_dataset 的桥接路径,可以规避 modelscope 与 datasets 之间
# `as_dataset() got an unexpected keyword argument 'verification_mode'` 这类版本错配错误。
print("Downloading dataset from ModelScope...")
print("ModelScope dataset id:", MODELSCOPE_DATASET_ID)

dataset_dir = dataset_snapshot_download(
    MODELSCOPE_DATASET_ID,
    cache_dir="./datasets",
)
print("Downloaded dataset dir:", dataset_dir)


def _parquet_files_for(split_name: str):
    pattern = os.path.join(dataset_dir, "data", f"{split_name}-*.parquet")
    files = sorted(glob.glob(pattern))
    if not files:
        raise FileNotFoundError(
            f"No parquet files matched pattern: {pattern}. "
            f"Please check the dataset repo layout under {dataset_dir}."
        )
    return files


raw_dataset = load_dataset(
    "parquet",
    data_files={
        "train": _parquet_files_for("train"),
        "validation": _parquet_files_for("validation"),
        "test": _parquet_files_for("test"),
    },
)

# 从 parquet 加载时,label 字段类型会退化成普通整数,这里显式 cast 成 ClassLabel,
# 这样后续 `dataset["train"].features["label"].names` 和原始 HF 版接口完全一致。
for split_name in list(raw_dataset.keys()):
    if not isinstance(raw_dataset[split_name].features.get("label"), ClassLabel):
        raw_dataset[split_name] = raw_dataset[split_name].cast_column(
            "label", ClassLabel(names=EMOTION_LABEL_NAMES)
        )

print("Raw dataset:", raw_dataset)


def maybe_limit(split, limit):
    split = split.shuffle(seed=SEED)
    if limit is None:
        return split
    return split.select(range(min(limit, len(split))))


dataset = DatasetDict({
    "train": maybe_limit(raw_dataset["train"], TRAIN_LIMIT),
    "validation": maybe_limit(raw_dataset["validation"], VALIDATION_LIMIT),
    "test": maybe_limit(raw_dataset["test"], TEST_LIMIT),
})

label_names = dataset["train"].features["label"].names
VALID_LABELS = set(label_names)
ALL_EVAL_LABELS = label_names + ["INVALID"]

print(dataset)
print("label_names:", label_names)
print("example:", dataset["train"][0])

构造 prompt-completion 格式数据

def to_prompt_completion(example):
    text = example["text"]
    label = label_names[example["label"]]
    user_content = f"Classify the emotion of this text:\n\n{text}"
    return {
        "prompt": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
        "completion": [
            {"role": "assistant", "content": label},
        ],
    }

sft_dataset = dataset.map(
    to_prompt_completion,
    remove_columns=dataset["train"].column_names,
)

print(sft_dataset)
print(sft_dataset["train"][0])

加载 tokenizer 和基础模型

print("Loading tokenizer from:", LOCAL_MODEL_DIR)

tokenizer = AutoTokenizer.from_pretrained(
    LOCAL_MODEL_DIR,
    use_fast=True,
    trust_remote_code=True,
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("pad_token:", tokenizer.pad_token)
print("eos_token:", tokenizer.eos_token)

# `google/gemma-4-E4B-it` 的 tokenizer 通常会自带 chat_template。
# 若缺失（缓存不完整等），从同一魔搭仓库拉取官方 chat_template.jinja 注入（权重已在上面整仓下载时可跳过额外拉取）。
TEMPLATE_SOURCE_MODEL_ID = "google/gemma-4-E4B-it"

def _load_official_gemma_chat_template() -> str:
    """从 gemma-4-E4B-it 仓库下载官方 chat_template.jinja 并返回字符串。

    主路径:modelscope.snapshot_download(allow_file_pattern=["chat_template.jinja"])
    兜底:ModelScope raw file API 直接 HTTP GET
    """
    try:
        template_dir = snapshot_download(
            TEMPLATE_SOURCE_MODEL_ID,
            cache_dir="./models",
            allow_file_pattern=["chat_template.jinja"],
        )
        path = os.path.join(template_dir, "chat_template.jinja")
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return f.read()
    except Exception as e:
        print("snapshot_download(allow_file_pattern) failed, fallback to HTTP. err =", e)

    import urllib.request
    url = (
        "https://www.modelscope.cn/api/v1/models/"
        f"{TEMPLATE_SOURCE_MODEL_ID}/repo?Revision=master&FilePath=chat_template.jinja"
    )
    with urllib.request.urlopen(url, timeout=60) as resp:
        return resp.read().decode("utf-8")


if not getattr(tokenizer, "chat_template", None):
    print(f"Loading official chat_template.jinja from {TEMPLATE_SOURCE_MODEL_ID} ...")
    tokenizer.chat_template = _load_official_gemma_chat_template()
    print("Loaded official chat_template, length =", len(tokenizer.chat_template))
else:
    print("tokenizer.chat_template already set, leaving as-is.")

# 自检:跑一次 apply_chat_template,确保模板可用。
_probe = tokenizer.apply_chat_template(
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello"},
    ],
    tokenize=False,
    add_generation_prompt=True,
)
print("chat_template probe output:\n" + _probe)

print("Loading base model from:", LOCAL_MODEL_DIR)

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)
print("HIP version:", getattr(torch.version, "hip", None))

base_model = AutoModelForCausalLM.from_pretrained(
    LOCAL_MODEL_DIR,
    torch_dtype=MODEL_DTYPE,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)

base_model.to(device)

base_model.config.use_cache = False
base_model.config.pad_token_id = tokenizer.pad_token_id
base_model.config.bos_token_id = tokenizer.bos_token_id
base_model.config.eos_token_id = tokenizer.eos_token_id

base_model.generation_config.pad_token_id = tokenizer.pad_token_id
base_model.generation_config.bos_token_id = tokenizer.bos_token_id
base_model.generation_config.eos_token_id = tokenizer.eos_token_id

print("Base model loaded.")
print("Base model device:", next(base_model.parameters()).device)

推理辅助函数

def extract_label(raw_text: str) -> str:
    raw_text = raw_text.strip().lower()
    match = LABEL_PATTERN.search(raw_text)
    if match:
        return match.group(1)

    tokens = raw_text.split()
    if not tokens:
        return "INVALID"

    return tokens[0].strip(".,!?:;\"'()[]{}")


def generate_label(model, tokenizer, user_text: str, system_prompt: str = SYSTEM_PROMPT, max_new_tokens: int = 4) -> str:
    user_content = f"Classify the emotion of this text:\n\n{user_text}"
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]

    device = next(model.parameters()).device

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

    input_len = inputs["input_ids"].shape[-1]
    model.eval()

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    raw_pred = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
    return extract_label(raw_pred)


def predict_emotion(text: str, model=None):
    model = model or base_model
    return generate_label(model, tokenizer, text)

predict_emotion("I feel so happy and excited today!")

评估函数

def evaluate_model(model, tokenizer, split="test", limit=EVAL_LIMIT):
    y_true, y_pred, rows = [], [], []
    raw_source = dataset[split]

    if limit is not None:
        raw_source = raw_source.select(range(min(limit, len(raw_source))))

    model.eval()

    for ex in tqdm(raw_source, desc=f"Evaluating {split}", leave=False):
        true_label = label_names[ex["label"]]
        raw_pred_label = generate_label(model, tokenizer, ex["text"], SYSTEM_PROMPT)
        pred_label = raw_pred_label if raw_pred_label in VALID_LABELS else "INVALID"

        y_true.append(true_label)
        y_pred.append(pred_label)
        rows.append({
            "text": ex["text"],
            "true_label": true_label,
            "pred_label": pred_label,
            "raw_pred_label": raw_pred_label,
            "correct": true_label == pred_label,
        })

    metrics = {
        "accuracy": accuracy_score(y_true, y_pred),
        "macro_f1": f1_score(y_true, y_pred, labels=label_names, average="macro", zero_division=0),
        "invalid_predictions": sum(1 for p in y_pred if p == "INVALID"),
        "evaluated_examples": len(y_true),
    }

    report = classification_report(
        y_true,
        y_pred,
        labels=label_names,
        output_dict=True,
        zero_division=0,
    )

    return metrics, report, pd.DataFrame(rows)


def confusion_matrix_df(pred_df):
    return pd.DataFrame(
        confusion_matrix(pred_df["true_label"], pred_df["pred_label"], labels=ALL_EVAL_LABELS),
        index=ALL_EVAL_LABELS,
        columns=ALL_EVAL_LABELS,
    )

微调前评估

pre_metrics, pre_report, pre_preds = evaluate_model(base_model, tokenizer, split="test", limit=EVAL_LIMIT)
pre_metrics

配置 LoRA

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)

定义训练参数

单卡版本默认：

per_device_train_batch_size=4
gradient_accumulation_steps=4
等效 batch size 为 16.
使用 adamw_torch，避免 AMD ROCm 下 bitsandbytes 优化器兼容问题

training_args = SFTConfig(
    output_dir=OUTPUT_DIR,

    per_device_train_batch_size=4,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=4,

    learning_rate=1e-4,
    weight_decay=0.01,
    lr_scheduler_type="linear",
    warmup_steps=50,
    num_train_epochs=1,

    logging_steps=5,
    eval_strategy="steps",
    eval_steps=25,
    save_strategy="steps",
    save_steps=25,
    save_total_limit=2,

    metric_for_best_model="eval_loss",
    greater_is_better=False,

    gradient_checkpointing=True,

    bf16=BF16,
    fp16=FP16,
    tf32=False,

    max_length=256,
    packing=False,
    completion_only_loss=True,

    remove_unused_columns=False,
    dataloader_num_workers=2,

    optim="adamw_torch",
    report_to="none",

    seed=SEED,
    data_seed=SEED,
)

开始 LoRA 微调

if isinstance(base_model, PeftModel):
    base_model = base_model.unload()
    base_model.config.use_cache = False

trainer = SFTTrainer(
    model=base_model,
    train_dataset=sft_dataset["train"],
    eval_dataset=sft_dataset["validation"],
    peft_config=lora_config,
    args=training_args,
    processing_class=tokenizer,
)

trainable_params = 0
total_params = 0
trainable_param_names = []

for name, param in trainer.model.named_parameters():
    total_params += param.numel()
    if param.requires_grad:
        trainable_params += param.numel()
        trainable_param_names.append(name)

if trainable_params == 0:
    raise RuntimeError("No trainable LoRA parameters were attached. Check target_modules before training.")

print(f"Trainable LoRA parameters: {trainable_params:,}")
print(f"Total parameters: {total_params:,}")
print(f"Trainable ratio: {100 * trainable_params / total_params:.4f}%")
print("Example trainable parameters:")
print(trainable_param_names[:20])

train_result = trainer.train()

trainer.model.eval()
trainer.model.config.use_cache = True
train_result

保存 LoRA adapter 和 tokenizer

trainer.model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

with open(os.path.join(OUTPUT_DIR, "train_metrics.json"), "w", encoding="utf-8") as f:
    json.dump(train_result.metrics, f, ensure_ascii=False, indent=2)

print("Saved adapter and tokenizer to:", OUTPUT_DIR)

微调后评估

ft_model = trainer.model
ft_model.eval()
ft_model.config.use_cache = True

post_metrics, post_report, post_preds = evaluate_model(ft_model, tokenizer, split="test", limit=EVAL_LIMIT)
post_metrics

pd.DataFrame(post_report).transpose()

confusion_matrix_df(post_preds)

对比微调前后效果

comparison_df = pd.DataFrame([
    {"stage": "pre_finetuning", **pre_metrics},
    {"stage": "post_finetuning", **post_metrics},
])

comparison_df

merged_examples = pre_preds.copy()
merged_examples = merged_examples.rename(columns={
    "pred_label": "pre_pred",
    "correct": "pre_correct",
    "raw_pred_label": "pre_raw_pred_label",
})

merged_examples["post_pred"] = post_preds["pred_label"]
merged_examples["post_raw_pred_label"] = post_preds["raw_pred_label"]
merged_examples["post_correct"] = post_preds["correct"]

changed_predictions = merged_examples[merged_examples["pre_pred"] != merged_examples["post_pred"]]
changed_predictions.head(20)

在这里插入图片描述

手动测试微调后的模型

def predict_emotion_ft(user_text: str) -> str:
    return generate_label(ft_model, tokenizer, user_text, SYSTEM_PROMPT)

test_texts = [
    "I feel completely heartbroken and alone.",
    "This is the best day of my life!",
    "I am really scared about what might happen tomorrow.",
    "I can't believe they remembered my birthday!",
    "I am so angry that nobody listened to me.",
    "I really love spending time with my family.",
]

for text in test_texts:
    print(text, "=>", predict_emotion_ft(text))

保存评估结果

comparison_df.to_csv(os.path.join(OUTPUT_DIR, "gemma4_emotion_before_after_metrics.csv"), index=False)
merged_examples.to_csv(os.path.join(OUTPUT_DIR, "gemma4_emotion_prediction_examples.csv"), index=False)
changed_predictions.to_csv(os.path.join(OUTPUT_DIR, "gemma4_emotion_changed_predictions.csv"), index=False)
pre_preds.to_csv(os.path.join(OUTPUT_DIR, "pre_finetuning_predictions.csv"), index=False)
post_preds.to_csv(os.path.join(OUTPUT_DIR, "post_finetuning_predictions.csv"), index=False)

pd.DataFrame(pre_report).transpose().to_csv(os.path.join(OUTPUT_DIR, "pre_finetuning_classification_report.csv"))
pd.DataFrame(post_report).transpose().to_csv(os.path.join(OUTPUT_DIR, "post_finetuning_classification_report.csv"))
confusion_matrix_df(pre_preds).to_csv(os.path.join(OUTPUT_DIR, "pre_finetuning_confusion_matrix.csv"))
confusion_matrix_df(post_preds).to_csv(os.path.join(OUTPUT_DIR, "post_finetuning_confusion_matrix.csv"))

print("Saved all outputs to:", OUTPUT_DIR)