Fine-tuning

The official training entry is Trainer. This chapter first fine-tunes DistilBERT on rotten_tomatoes (a short run works on CPU), then sketches PEFT LoRA. Full-parameter 7B training is not a beginner task.

VRAM warning: DistilBERT classification fits a 4 GB GPU or a CPU (slowly). Full-parameter causal-LM fine-tunes OOM easily—use LoRA, or revisit batches and gradients in the PyTorch course. Without a GPU, set a tiny max_steps so the script merely runs.


Trainer: review sentiment

import numpy as np
import evaluate
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    TrainingArguments,
    Trainer,
)

model_id = "distilbert/distilbert-base-uncased"
raw = load_dataset("rotten_tomatoes")
tokenizer = AutoTokenizer.from_pretrained(model_id)

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True)

tokenized = raw.map(tokenize, batched=True)
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)
accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return accuracy.compute(predictions=preds, references=labels)

args = TrainingArguments(
    output_dir="./rt-distilbert",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=2,
    weight_decay=0.01,
    report_to="none",
    fp16=False,  # set True on NVIDIA GPUs
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["validation"],
    processing_class=tokenizer,
    data_collator=DataCollatorWithPadding(tokenizer),
    compute_metrics=compute_metrics,
)
# trainer.train()
# trainer.save_model("./rt-distilbert-final")

On CPU, try per_device_train_batch_size=4 and raw["train"].select(range(512)). For loop internals, compare this site’s PyTorch tutorial.


Infer the fine-tuned model

from transformers import pipeline

pipe = pipeline("text-classification", model="./rt-distilbert-final")
print(pipe("A warm, clever, and generous film."))

Labels may show up as LABEL_0 / LABEL_1. Pass id2label={0: "NEGATIVE", 1: "POSITIVE"} into from_pretrained.


Minimal PEFT LoRA sketch

Train low-rank adapters; freeze the base. Fits a small causal LM such as Qwen/Qwen2.5-0.5B-Instruct (GPU recommended):

from peft import LoraConfig, TaskType, get_peft_model
from transformers import AutoModelForCausalLM

base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B-Instruct",
    device_map="auto",
    dtype="auto",
)
peft_cfg = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    task_type=TaskType.CAUSAL_LM,
    target_modules=["q_proj", "v_proj"],
)
model = get_peft_model(base, peft_cfg)
model.print_trainable_parameters()
# Then hand `model` to Trainer; data must provide input_ids + labels

pip install peft. After training, model.save_pretrained("./qwen-lora") stores only the adapter (MBs). Reload with PeftModel.from_pretrained(base, "./qwen-lora"). Data wrangling is outlined in Practical Examples.


When not to use Trainer

SituationAlternative
Chat onlyOllama
Custom loop / DeepSpeedAccelerate + plain PyTorch
7B+ full params on 8 GBLoRA / QLoRA, or skip full FT

Agents can draft training scripts via Hugging Face Skills; you still own VRAM and licenses.


Next steps

评论