从零构建并微调一个类大模型

从零构建并微调一个类大模型:以数字比较为例

本文将手把手教你从零构建一个带有 Transformer 编码器的简易“大模型”,用于学习“两个数字谁更大”的任务,并进一步扩展到微调与多轮训练场景。适合希望理解大模型基本原理、动手构建微型语言模型的开发者。


任务定义

我们希望构建一个模型,输入两个整数 ab,模型输出判断结果:

  • 1:表示 a > b
  • 0:表示 a <= b

例如:

1
2
输入: [3, 7] → 输出: 0  
输入: [9, 2] → 输出: 1

第一阶段:构建最小可用模型(MLP)

使用 PyTorch 实现一个最简单的神经网络模型,直接对 [a, b] 输入进行处理。

📦 数据生成

1
2
3
4
5
6
7
8
9

def generate_data(n=1000):
    data = []
    for _ in range(n):
        a = random.randint(0, 9)
        b = random.randint(0, 9)
        label = 1 if a > b else 0
        data.append(([a, b], label))
    return data

🧱 模型结构

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class SimpleComparator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 1),
            nn.Sigmoid()
        )
    def forward(self, x):
        return self.model(x)

🚀 模型训练

1
2
3
4
5
6
for epoch in range(20):
    model.train()
    pred = model(x_train)
    loss = loss_fn(pred, y_train)
    loss.backward()
    optimizer.step()

第二阶段:构建类 GPT 的简化 Transformer 模型

我们升级模型,引入:

  • Embedding 层

  • 位置编码

  • Transformer Encoder

  • 分类头

🔧 模型结构

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class TransformerComparator(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(10, 32)
        self.position = nn.Parameter(torch.randn(1, 2, 32))
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=32, nhead=2, dim_feedforward=64, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=1)
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 2, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.embedding(x) + self.position
        x = self.transformer(x)
        return self.classifier(x)

🔁 第三阶段:模型微调(Fine-tuning)

我们以“判断两个数是否相等”的任务为例,对已经训练好的模型进行微调。

新任务数据

1
2
3
def generate_equal_data(n=500):
    ...
    label = 1 if a == b else 0

♻️ 微调步骤

1
2
3
4
5
finetune_model = TransformerComparator()
finetune_model.load_state_dict(torch.load("transformer_comparator.pt"))

for epoch in range(10):
    ...

经过几轮微调,模型可以很好地迁移到新任务,例如从“大小比较” → “相等判断”。

🔄 第四阶段:支持多轮训练(对话上下文)

我们进一步升级模型为支持“上下文记忆”,模拟多轮对话任务:

3 > 7? → No; 7 > 5? → ?

📦 构造序列数据

1
2
3
def make_conversation_data(n=500):
    prompt = f"{a1} > {b1}? → {res1}; {a2} > {b2}? →"
    label = res2

🧱 字符级模型结构

1
2
3
4
5
6
class CharTransformer(nn.Module):
    ...
    def forward(self, x):
        x = self.embed(x) + self.pos[:, :x.size(1), :]
        out = self.transformer(x)
        return torch.sigmoid(self.classifier(out[:, -1, :]))

✏️ Tokenizer + Padding

1
2
def tokenize(text):
    return [char2idx[ch] for ch in text if ch in char2idx]

🚀 多轮训练结果

训练后,模型能理解上下文并输出合理判断:

1
3 > 7? → No; 7 > 5? → Yes ✅

🧩 拓展方向建议

类型思路
增加数据复杂度处理三数比较、加法判断、等式验证等任务
使用自然语言将数字比较转化为自然语言输入 “Is 3 greater than 5?”
中文支持输入 “3 比 5 大吗?",使用字级 tokenizer 和 Embedding
接入大模型替换为 LLaMA / GPT 微调,使用 LoRA / PEFT 微调工具
数据格式迁移使用 ChatML / Alpaca 格式制作更多轮推理任务数据集

完整代码

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import torch
import torch.nn as nn
import torch.optim as optim
import random

# -------------------
# 数据生成函数
# -------------------
def generate_data(n=1000):
    data = []
    for _ in range(n):
        a = random.randint(0, 9)
        b = random.randint(0, 9)
        label = 1 if a > b else 0
        data.append(([a, b], label))
    return data

def generate_equal_data(n=500):
    data = []
    for _ in range(n):
        a = random.randint(0, 9)
        b = random.randint(0, 9)
        label = 1 if a == b else 0
        data.append(([a, b], label))
    return data

# -------------------
# 数据转换为Tensor
# -------------------
def to_tensor(data):
    x = torch.tensor([item[0] for item in data], dtype=torch.long)
    y = torch.tensor([item[1] for item in data], dtype=torch.float32).unsqueeze(1)
    return x, y

# -------------------
# 简单Transformer比较器模型
# -------------------
class TransformerComparator(nn.Module):
    def __init__(self):
        super().__init__()
        vocab_size = 10       # 数字范围 0-9
        embedding_dim = 32
        max_seq_len = 2

        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.position = nn.Parameter(torch.randn(1, max_seq_len, embedding_dim))

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embedding_dim,
            nhead=2,
            dim_feedforward=64,
            batch_first=True,
            activation="relu"
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=1)

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(embedding_dim * max_seq_len, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.embedding(x) + self.position  # [B, 2, D]
        x = self.transformer(x)                # [B, 2, D]
        return self.classifier(x)              # [B, 1]

# -------------------
# 训练函数
# -------------------
def train(model, optimizer, loss_fn, x_train, y_train, epochs=20):
    model.train()
    for epoch in range(epochs):
        pred = model(x_train)
        loss = loss_fn(pred, y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch % 5 == 0 or epoch == epochs-1:
            print(f"Epoch {epoch}: loss={loss.item():.4f}")

# -------------------
# 测试函数
# -------------------
def test(model, x_test, y_test):
    model.eval()
    with torch.no_grad():
        pred = model(x_test)
        acc = ((pred > 0.5) == y_test).float().mean().item()
    print(f"Test Accuracy: {acc:.4f}")
    return acc

# -------------------
# 主程序
# -------------------
if __name__ == "__main__":
    # 1. 训练初始模型判断 a > b
    train_data = generate_data(1000)
    test_data = generate_data(200)

    x_train, y_train = to_tensor(train_data)
    x_test, y_test = to_tensor(test_data)

    model = TransformerComparator()
    loss_fn = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.005)

    print("训练判断 a > b 的模型")
    train(model, optimizer, loss_fn, x_train, y_train, epochs=30)
    test(model, x_test, y_test)

    # 保存模型
    torch.save(model.state_dict(), "transformer_comparator.pt")

    # 2. 微调模型判断 a == b
    equal_train_data = generate_equal_data(500)
    equal_test_data = generate_equal_data(100)

    x_eq_train, y_eq_train = to_tensor(equal_train_data)
    x_eq_test, y_eq_test = to_tensor(equal_test_data)

    finetune_model = TransformerComparator()
    finetune_model.load_state_dict(torch.load("transformer_comparator.pt"))

    optimizer_ft = optim.Adam(finetune_model.parameters(), lr=0.001)
    loss_fn_ft = nn.BCELoss()

    print("微调模型判断 a == b")
    for epoch in range(10):
        finetune_model.train()
        pred = finetune_model(x_eq_train)
        loss = loss_fn_ft(pred, y_eq_train)
        optimizer_ft.zero_grad()
        loss.backward()
        optimizer_ft.step()

        if epoch % 2 == 0 or epoch == 9:
            finetune_model.eval()
            with torch.no_grad():
                test_pred = finetune_model(x_eq_test)
                acc = ((test_pred > 0.5) == y_eq_test).float().mean().item()
            print(f"[FT Epoch {epoch}] loss={loss.item():.4f}, accuracy={acc:.4f}")

总结

我们从零构建了一个可训练的“数字比较模型”,并逐步演化出:

  • 基础 MLP 二分类模型;

  • 类 GPT 的 Transformer 编码结构;

  • 微调流程(Fine-tuning 到相等判断);

  • 多轮对话理解能力的训练方法。

  • 这种从小模型到“大模型训练范式”的过程,不仅适用于教学和实验,也能作为大模型任务工程化前的原型验证参考。

Licensed under CC BY-NC-SA 4.0
使用 Hugo 构建
主题 StackJimmy 设计