ScratchLLMStepByStep：SFT之分类微调

1. 引言

前面我们花了三节内容来介绍预训练，包括如何从零搭建、如何加速运算、如何分布式加速训练，本节开始我们将进入监督微调（SFT）阶段。

常见语言模型的微调任务有两类，分类微调和指令微调。

• 分类微调模型通常是一个专用模型，它只能用来进行特定分类标签的预测，例如：输入一封邮件文本，输出这封邮件是垃圾邮件、非垃圾邮件。
• 指令微调模型通常是一个多任务模型，它可以同时在多个任务上表现良好，训练的难度也更大。

在这两类微调任务中，分类微调更为专注单一任务的优化，它在处理特定应用场景（例如情感分析、主题分类或意图识别等任务时）时，能够表现出更高的准确率，同时对训练数据量的要求也更少。并且，由于是单一任务，分类微调训练时不需要指令，直接输入数据进行训练，如下所示：

注：分类微调在传统的机器学习中非常常见，例如经典的手写数字识别，只不过手写数字识别是对图片进行分类，我们这里是对文本进行分类。

本节我们将以垃圾邮件检测这个任务为例，来介绍如何对预训练模型进行分类微调。

2. 数据准备

2.1 数据下载

先下载数据集：

!curl -o ../../dataset/minigpt/smsspam.txt "https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"

下载完后进行解压，数据集是一个文本文件，可以通过pandas库来预览下载后的数据。

import pandas as pd

df = pd.read_csv("/data2/minigpt/dataset/sft/smsspam.txt", sep='\t', header=None, names=["label", "text"])
df.head()

| | label | text |
| 0 | ham | Go until jurong point, crazy.. Available only ... |
| 1 | ham | Ok lar... Joking wif u oni... |
| 2 | spam | Free entry in 2 a wkly comp to win FA Cup fina... |
| 3 | ham | U dun say so early hor... U c already then say... |

4	ham	Nah I don't think he goes to usf, he lives aro...

可以看到，数据集比较简单，每个数据只有两个信息：文本(text)和标签分类(label)，标签分类也只有两个类别：ham表示正常邮件，spam表示垃圾邮件。

2.2 数据处理

下面，我们来查看下两个标签分类的数据分布。

df['label'].value_counts()

label
ham     4825
spam     747
Name: count, dtype: int64

可以看到，两种标签的数据分布是不均衡的，正常邮件ham的数据量远比垃圾邮件spam类别要多。这不利于模型的训练学习，需要对数据分布作平衡操作。

数据平衡通常有两种方式：

1. 对偏多的数据做欠采样；
2. 对偏少的数据做过采样。

我们本节主要是为了演示，所以采用比较简单的方式，对ham类别的数据进行欠采样，使之与spam类别的数据量相等，均为747条。

注：如果想要进行过采样，可以参考这篇文章中介绍的方法：over_sampling

def balance_dataset(df):
    spam_set = df[df["label"] == "spam"]
    ham_sub_set = df[df["label"] == "ham"].sample(spam_set.shape[0], random_state=123)
    return pd.concat([spam_set, ham_sub_set])

balanced_df = balance_dataset(df)
balanced_df["label"].value_counts()

label
spam    747
ham     747
Name: count, dtype: int64

将字符串分类转换为数字分类0和1，便于模型预测。

balanced_df['label'] = balanced_df['label'].map({'spam': 1, 'ham': 0})
balanced_df['label'].value_counts()

label
1    747
0    747
Name: count, dtype: int64

接下来，我们写一个切分函数random_split，采用8:1:1的比例，随机的将数据集分割为训练集、验证集和测试集。

def random_split(df, train_ratio, eval_ratio):
    df = df.sample(frac=1, random_state=123).reset_index(drop=True)
    train_len = int(len(df) * train_ratio)
    eval_len = train_len + int(len(df) * eval_ratio)
    return df[:train_len], df[train_len: eval_len], df[eval_len:]

train_df, eval_df, test_df = random_split(balanced_df, 0.8, 0.1)
len(train_df), len(eval_df), len(test_df)

(1195, 149, 150)

2.3 数据加载器

模型训练时一般都会使用小批量梯度下降的方法，这要求同一批训练数据必须具有相同的序列长度，我们可以使用前面训练的分词器，对前10条文本进行序列化查看下长度分布。

from transformers  import AutoTokenizer

tokenizer_path = "/data2/minigpt/models/tokenizer_v3"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=False)
inputs = [tokenizer.encode(text) for text in train_df['text']]
for i, v in enumerate(inputs[:10]):
    print(f"seq {i} length: {len(v)}")

seq 0 length: 42
seq 1 length: 45
seq 2 length: 11
seq 3 length: 23
seq 4 length: 5
seq 5 length: 49
seq 6 length: 53
seq 7 length: 6
seq 8 length: 79
seq 9 length: 12

可以看到，训练数据集中每条文本序列化后的长度基本都不相同。我们想要每个小批量都具有相同的长度，就只能有两种做法：

1. 长序列变短：将一个数据集或小批量中的所有文本都切割成与最短文本的序列长度相同；
2. 短序列变长：将一个数据集或小批量中的所有文本都填充到与最长文本的序列长度相同；

通常情况下，第二种方法用的更多，它能很好保留每条文本的内容完整。下面是使用<|endoftext|>作为填充token的长度填充示例。

下面自定义一个数据集类，来封装文本的序列化和长度填充逻辑，以pytorch中标准化的形式来提供数据集迭代功能，每次迭代时会自动调用__getitem__内置方法，返回一条输入序列和一个目标分类标签。

from torch.utils.data import Dataset

class SpamDataset(Dataset):
    def __init__(self, raw_df, tokenizer, max_tokens=1024):
        self.labels = raw_df['label']
        self.tokenizer = tokenizer
        pad_token_id = tokenizer.unk_token_id
        inputs = [tokenizer.encode(text)[:max_tokens] for text in raw_df['text']]
        max_length = max([len(item) for item in inputs])
        self.inputs = [item + [pad_token_id] * (max_length - len(item)) for item in inputs]

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, i):
        return (torch.tensor(self.inputs[i], dtype=torch.int64), 
            torch.tensor(self.labels.iloc[i], dtype=torch.int64))

接下来，我们将前面切割后的训练集、验证集、测试集分别用SpamDataset封装，得到train_set、eval_set、test_set三个标准数据集，并以此为基础为训练、验证、则试三个场景分别创建三个小批量数据加载器train_loader、eval_loader、test_loader。

import torch
from torch.utils.data import DataLoader

torch.manual_seed(123)
batch_size = 8

train_set = SpamDataset(train_df, tokenizer)
eval_set = SpamDataset(eval_df, tokenizer)
test_set = SpamDataset(test_df, tokenizer)

train_loader = DataLoader(train_set, shuffle=True, batch_size=batch_size)
eval_loader = DataLoader(eval_set, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_set, shuffle=True, batch_size=batch_size)

对DataLoader进行迭代时，会自动按照batch_size指定的小批量大小，一次性从Dataset中取多条数据（如下所示inputs和labels分别是8条）。

inputs, labels = next(iter(train_loader))
inputs, labels

(tensor([[ 1148,  3476,  2948,  ...,     0,     0,     0],
         [ 8150,    68,  6409,  ...,     0,     0,     0],
         [12203,  6527,  1017,  ...,     0,     0,     0],
         ...,
         [17647,  2511,   276,  ...,     0,     0,     0],
         [   40, 14387,    39,  ...,     0,     0,     0],
         [   43,  2341,   290,  ...,     0,     0,     0]]),
 tensor([1, 0, 1, 1, 1, 0, 1, 0]))

3. 准备模型

在这一部分，我们需要对之前预训练的模型的输出头进行改造，用二分类代替词表长度的多分类。

3.1 加载模型状态

首先，加载之前预训练的模型，来看下模型的结构。

%run transformer.py

device = 'cpu'
checkpoint_path = "/data2/minigpt/models/20241015/checkpoint-180000.pth"

model = MiniGPT(GPTConfig(flash_attn=True))
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
model.load_state_dict(checkpoint['model_state'])
model

MiniGPT(
  (token_emb): Embedding(32000, 768)
  (drop_emb): Dropout(p=0.1, inplace=False)
  (decode_layers): ModuleList(
    (0-11): 12 x TransformerBlock(
      (atten): FlashMultiHeadAttention(
        (Wq): Linear(in_features=768, out_features=768, bias=False)
        (Wk): Linear(in_features=768, out_features=768, bias=False)
        (Wv): Linear(in_features=768, out_features=768, bias=False)
        (Wo): Linear(in_features=768, out_features=768, bias=True)
        (dropout): Dropout(p=0.1, inplace=False)
      )
      (ffn): FeedForward(
        (layers): Sequential(
          (0): Linear(in_features=768, out_features=3072, bias=True)
          (1): GELU(approximate='none')
          (2): Linear(in_features=3072, out_features=768, bias=True)
        )
      )
      (drop): Dropout(p=0.1, inplace=False)
      (layernorm1): LayerNorm()
      (layernorm2): LayerNorm()
    )
  )
  (final_norm): LayerNorm()
  (out_head): Linear(in_features=768, out_features=32000, bias=True)
)

目前的输出头out_head将嵌入维度768映射到32000（词表大小），但是在这个垃圾分类任务中我们只需要二个目标分类0和1，因此，我们需要自己定义一个二分类输出头并替换掉原来的out_head，如下所示。

3.2 自定义输出头

torch.manual_seed(123)
classes = 2

model.out_head = nn.Linear(in_features=model.out_head.in_features, out_features=classes)
model.out_head

Linear(in_features=768, out_features=2, bias=True)

理论上来讲，只需要训练新的输出头out_head，但业界一些实验表明，训练更多的层有助于提高性能。因此，我们将最后一个transformer block和连接它的layernorm设置为可训练，其它层设置为不可训练（效果如下图和代码所示）。

for param in model.parameters():
    param.requires_grad = False

for param in model.decode_layers[-1].parameters():
    param.requires_grad = True

for param in model.final_norm.parameters():
    param.requires_grad = True

for param in model.out_head.parameters():
    param.requires_grad = True

trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad==True)
all_params = sum(p.numel() for p in model.parameters())
trainable_params, all_params

(7088642, 109605890)

如上打印信息所示，需要训练的参数总共有708万，占所有参数109605890的比例约6.47%。

3.3 模型测试

拿一个输入文本来测试二分类输出。

input_text = "Do you have time?"
input_tokens = tokenizer.encode(input_text)
inputs = torch.tensor(input_tokens).unsqueeze(0)
inputs.shape

torch.Size([1, 5])

with torch.no_grad():
    outputs = model(inputs)

outputs.shape

torch.Size([1, 5, 2])

可以看到，输入与之前并没有什么区别，是一个batch_size=1,seq_len=5的单条序列，输出与之前有所不同，变成了维度为2的二分类输出。

根据因果注意力掩码的特性，每个token都只能看到此token本身以及当前位置之前的token，而最后一个token会是唯一一个能包含所有token上下文的向量，因此，我们只用最后一个token的向量来计算最终的输出分类结果。

probs = torch.softmax(outputs[:, -1, :], dim=-1)
preds = torch.argmax(probs, dim=-1)
probs, preds

(tensor([[0.4616, 0.5384]]), tensor([1]))

可以看到，在未经训练情况下，模型预测的分类为是1（垃圾邮件），但实际上这个示例文本Do you have time?与垃圾邮件并无关系，这正是我们要通过微调模型来改善的地方。

接下来，我们编写一个函数calc_accuracy，用来计算模型在指定数据集上的分类准确率。它通过收集正确预测的数据条数占比来实现。

import torch.nn.functional as f

labels = torch.tensor([0])
loss = f.cross_entropy(outputs[:, -1, :], labels)
loss.item()

0.7731215953826904

def calc_accuracy(model, dataloader, device):
    num_datas = 0
    num_corrects = 0
    for inputs, labels in dataloader:
        inputs, labels = inputs.to(device), labels.to(device)
        with torch.no_grad():
            logits = model(inputs)[:, -1, :]
        probs = torch.softmax(logits, dim=-1)
        preds = torch.argmax(probs, dim=-1)
        num_datas += len(labels)
        num_corrects += (preds == labels).sum().item()
    return num_corrects / num_datas

计算模型在验证数据集上的分类准确率。

calc_accuracy(model, eval_loader, 'cpu')

0.5637583892617449

计算模型在测试数据集上的分类准确率。

calc_accuracy(model, test_loader, 'cpu')

0.5733333333333334

可以看到，未经过微调的模型，不论是在验证数据集还是在测试数据上准确率都不是很好，这也是我们需要通过微调模型来改善的地方。

4. 微调

4.1 训练代码准备

在微调开始之前，我们需要编写几个函数，以实现单步训练、模型评估和主训练循环。

单步训练和前面预训练阶段的实现基本是一致的，步骤基本固定为以下这几步：

• 梯度清零
• 模型推理
• 损失计算
• 反向传播
• 更新参数

import torch.nn.functional as f

def train_step(model, optimizer, X, Y):
    # 清零梯度
    optimizer.zero_grad()
    # 模型调用收集logits
    logits = model(X)
    # print(f"logits.shape: {logits.shape}, target.shape: {Y.shape}")
    # 计算损失
    loss = f.cross_entropy(logits[:, -1, :], Y.flatten())
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    return loss.item()

模型评估与预训练阶段也基本类似，统计模型在整个数据集上的平均损失。

def evaluate(model, dataloader, device):
    model.eval()
    num_batches = len(dataloader)
    total_loss = 0
    with torch.no_grad():
        for (X, Y) in dataloader:
            X, Y = X.to(device), Y.to(device)
            logits = model(X)
            loss = f.cross_entropy(logits[:, -1, :], Y.flatten())
            total_loss += loss.item()
    model.train()
    return total_loss/num_batches

主训练方法的实现也与前面大同小异，对所有的数据训练num_epochs指定的轮数，并在每轮训练中跟踪训练损失和验证损失的变化。

def train(model, optimizer, train_loader, eval_loader, device, num_epochs):
    train_losses, eval_losses = [], []
    for epoch in range(num_epochs):
        model.train()
        for inputs, targets in train_loader:
            train_loss = train_step(model, optimizer, inputs.to(device), targets.to(device))
            train_losses.append(train_loss)
        eval_loss = evaluate(model, eval_loader, device)
        eval_losses.append(eval_loss)
        print(f"Epoch: {epoch}, Train Loss: {train_loss:.4f}, Eval Loss: {eval_loss:.4f}")

    return train_losses, eval_losses

4.2 开始训练

%%time
device = 'cuda:0'
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)
train(model, optimizer, train_loader, eval_loader, device, 20)

    Epoch: 0, Train Loss: 0.3042, Eval Loss: 0.4855
    Epoch: 1, Train Loss: 0.0237, Eval Loss: 0.2542
    Epoch: 2, Train Loss: 0.2716, Eval Loss: 0.2975
    Epoch: 3, Train Loss: 0.0011, Eval Loss: 0.2413
    Epoch: 4, Train Loss: 0.0008, Eval Loss: 0.2194
    Epoch: 5, Train Loss: 0.7194, Eval Loss: 0.2193
    Epoch: 6, Train Loss: 0.0016, Eval Loss: 0.2332
    Epoch: 7, Train Loss: 0.2460, Eval Loss: 0.2788
    Epoch: 8, Train Loss: 0.0699, Eval Loss: 0.3474
    Epoch: 9, Train Loss: 0.0037, Eval Loss: 0.3063
    Epoch: 10, Train Loss: 0.0004, Eval Loss: 0.2577
    Epoch: 11, Train Loss: 0.0004, Eval Loss: 0.4328
    Epoch: 12, Train Loss: 0.0935, Eval Loss: 0.2229
    Epoch: 13, Train Loss: 0.0000, Eval Loss: 0.2417
    Epoch: 14, Train Loss: 0.0022, Eval Loss: 0.5061
    Epoch: 15, Train Loss: 0.0369, Eval Loss: 0.2403
    Epoch: 16, Train Loss: 0.0000, Eval Loss: 0.6432
    Epoch: 17, Train Loss: 0.0008, Eval Loss: 0.2384
    Epoch: 18, Train Loss: 0.0000, Eval Loss: 0.2711
    Epoch: 19, Train Loss: 0.0000, Eval Loss: 0.3409
    CPU times: user 1min 7s, sys: 647 ms, total: 1min 8s
    Wall time: 1min 8s

注：从损失数据中可以看到，虽然训练了20轮，但跑到第6轮之后，就有些过拟合了，表现为训练损失很低，但验证损失不降反升的现象，虽然是我们的数据量太少，又重复训练了多轮的缘故。

4.3 准确率评估

下面来统计下模型在不同数据集上的准确率。

train_acc = calc_accuracy(model, train_loader, device)
eval_acc = calc_accuracy(model, eval_loader, device)
test_acc = calc_accuracy(model, test_loader, device)
print(f"train_acc: {train_acc}, eval_acc: {eval_acc}, test_acc: {test_acc}")

train_acc: 0.999163179916318, eval_acc: 0.9328859060402684, test_acc: 0.9133333333333333

结果分析如下：

• 在训练数据集准确率接近1.0，不太正常，原因在于训练轮数过多，过度拟合了训练数据集。
• 在验证数据集上0.9328%，较为正常。
• 在测试数据集上准确率为0.913，这是模型从未见过的数据集，这个数字最为客观。

注：基于前面的损失下降数据，模型性能最好的点理论上应该是在epoch=5(eval_loss=0.2193)的时候，但是这个点的模型状态我们未能保存下来，这里主要为了演示，就没有再次训练，你如果感兴趣，可以修改代码将每个epoch的结果都保存下来，来找到准确率表现最佳的模型状态。

下面我们将模型状态保存下来，以备后面使用。

import os
output_dir = "/data2/minigpt/models/classify/"
os.makedirs(output_dir, exist_ok=True)
torch.save(model.state_dict(), os.path.join(output_dir, "model_classify_v1.pth"))

4.4 模型使用

我们先编写一个predict函数，用于对文本进行垃圾邮件分类预测，并将结果转换为人类可以理解的文本形式spam和not spam。

def predict(text, model, tokenizer, device, max_length):
    token_ids = tokenizer.encode(text)
    input_ids = torch.tensor(token_ids).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(input_ids)
    probs = torch.softmax(logits[:, -1, :], dim=-1)
    pred = torch.argmax(probs, dim=-1)
    return "spam" if pred.item() == 1 else "not spam"

我们使用刚微调的垃圾邮件分类模型，对下面两个样例文本进行实际预测。

texts = [
    "You are a winner you have been specially",
    "selected to receive $1000 cash or a $2000 award.",
]

for text in texts:
    print(text, " --> ", predict(text, model, tokenizer, device, 131))

You are a winner you have been specially  -->  not spam
selected to receive $1000 cash or a $2000 award.  -->  spam

可以看到，对于这两个样例文本，在微调后的模型上分类均是正确的，基本达到了分类训练的目标。

小结：本节我们主要介绍了基于预训练模型进行分类微调的实施过程，先对数据进行预处理并封装小批量数据加载器，再对模型结构进行改造，替换了一个二分类输出头，最后对模型进行微调训练，得到一个在测试数据集上准确率为91%的垃圾邮件分类模型，这个准确率应该还可以继续提高，有兴趣可以按照自己的想法动手修改验证。

相关阅读

• 预训练模型之分布式训练加速
• 带你构建MiniGPT