手搓minigpt 预训练+SFT

1. 前言

本文基于github开源项目：minimind在ai快速发展的现在，我有点儿不满足只写一大堆提示词去调 llm 的接口。心里挺好奇，llm 是怎么造出来的？于是动手试了试自己练一个 minigpt 出来。这篇笔记，记录了我从小白到训练出一个 minigpt 的学习过程。

2. 机器学习

2.1. 为什么要用机器学习

我的工作内容是后端开发，开发的流程一般是 确定需求 -> 设计架构 -> 开发测试运维。可以认为是把一个复杂的 规则 翻译成代码。

当一个事情的规则难以捋清楚时，机器学习上场了。

2.2. 机器学习是什么？

机器学习可以认为是基于 y=wx+b 而来（当然实际上比这个基础公式复杂多了）。用已有的数据去预测新的数据。其实从y=wx+b这个公式上就能看出为什么现在的大语言模型这么火，之前在不同的领域想用模型得训练w和b，在一个新的领域就得练一次，但是大语言模型把w和b训练得很强大，让用户可以只改变x就输出预期的y。emm也是有点儿力大飞砖的感觉。

2.3. 一个极简机器学习的例子

前面介绍了机器学习的背景，接着介绍下一个最小的机器学习的流程。

这里用机器学习来预测一个函数 y=a * x**3+b * x**2 + c * x + d。手动计算的梯度。（在pytorch框架中，是不需要手动计算梯度的。）

参考代码：

import matplotlib.pyplot as pyplotimport mathimport sysX = [0.01 * x for x in range(100)]Y = [4*x**3 + 2*x**2 + 3*x + 4 for x in X]def func(x):y = w0 * x**3 + w1 * x ** 2 + w2 * x + w3return ydef loss(y_pred, y_true):return (y_pred - y_true) ** 2# 权重随机初始化w0, w1, w2, w3 = 0,0,0,0# 学习率设置lr = 0.1#batch sizebatch_size = 10# 训练过程for epoch in range(5000):epoch_loss = 0grad_w0 = 0grad_w1 = 0grad_w2 = 0grad_w3 = 0counter = 0for x, y_true in zip(X, Y):y_pred = func(x)epoch_loss += loss(y_pred, y_true)counter += 1#梯度计算grad_w0 += 2 * (y_pred - y_true) * x ** 3grad_w1 += 2 * (y_pred - y_true) * x ** 2grad_w2 += 2 * (y_pred - y_true) * x grad_w3 += 2 * (y_pred - y_true) if counter == batch_size:#权重更新# print(f"权重更新前：{w1} {w2} {w3}")w0 = w0 - lr * grad_w0/batch_sizew1 = w1 - lr * grad_w1/batch_size #sgd w2 = w2 - lr * grad_w2/batch_sizew3 = w3 - lr * grad_w3/batch_size# print(f"权重更新后：{w1} {w2} {w3}")counter = 0grad_w0 = 0grad_w1 = 0grad_w2 = 0grad_w3 = 0# print(f"权重更新后：{w1} {w2} {w3}")epoch_loss =epoch_loss/len(X)print("第%d轮， loss %f" %(epoch, epoch_loss))if epoch_loss < 0.00001:breakprint(f"训练后权重:w0:{w0} w1:{w1} w2:{w2} w3:{w3}")#使用训练后模型输出预测值Yp = [func(i) for i in X] #预测值与真实值比对数据分布pyplot.scatter(X, Y, color="red")pyplot.scatter(X, Yp)pyplot.show()# 求导说明：# 对损失函数 ：(y_pred- y_true) ** 2，求w1的梯度# 链式法则# 令 u = y_pred- y_true = w1 * x ** 2 + w2 * x + w3- y_true， 即损失函数为 = u ** 2# 所以根据链式法则，导函数 = 2u * du/dw1 （u对w1的偏导数）# du/dw1 = x ** 2(除w1外，其他字母都看做常数，常数导数为0)# 所以导函数 = 2u * x ** 2# 同理可得w2,w3的导函数# 所以梯度 = [2u * x ** 2, 2u * x, 2u]

3. Transformer

整个项目最核心的就是transformer-decoder块，所以先提到前面解释下。

3.1. 单个 Decoder 块

gpt结构是 Decoder-only 模型，只用到了decoder部分。这里面最核心的是多头自注意力机制。

3.2. Multi-Head Self-Attention

多头自注意力机制。接下来我解释为什么要自注意力机制、为什么要多头。在更早期语言模型用的RNN这种，劣势是处理不好长距离依赖。比如这个例子 “狗没有穿过街道，因为它很累。” 句子里面的 “它” 指的是狗，但是用RNN这种模型就很难做到。于是引入了自注意力机制，用一个二维矩阵去记录每两个token之间的相关性。公式是：

Q*K^T图示：

然后人们希望模型能从不同的角度去关注句子，比如能同时从语法、近义词等等角度关注句子。单个注意力头只能关注一种特定关系，多头让模型同时关注来自不同子空间的信息。

3.3. Attention Mask

gpt这种模型是根据前面的话推测出之后的话，所以需要有一个方式不让前面的话看到后面的话。让字只能看到自己和前面的。

4. minigpt 的实现过程

先预训练，让模型学会往后接话；然后SFT，让模型的回答看起来像人在回答。

4.1. 预训练

4.1.1. 选择数据集

预训练的数据集选择的是minimind里面的 pretrain_t2t_mini.jsonl。数据格式是这样：

{"text": "给我生成一首有关秋天的诗歌。秋日早晨，清风拂面。n金色的叶子，似火在燃烧。n露珠晶莹，如珍珠般美丽。n秋的气息，弥漫在空气中。n余音袅袅，如鸟儿的歌唱。n美丽的秋天，是大自然的馈赠。帮我想一些创意，给即将到来的公司年会准备节目。一些节目比如能否请一位表演者为我们表演一曲钢琴曲，或者请一位小提琴手为大家演奏一首古典曲目。如果想要画面更具有视觉冲击力，可以安排一个魔术师或者杂技演员的表演。另外，也可以设计一些小游戏或者有奖竞猜来增添活动的趣味性，这些小游戏可以和公司的文化、发展历程等相关。请问给我讲一个清净的法则。无为而治是一项清净的法则。即在处理问题时，不要强行干预，反而尽可能地减少干预，坚持自然的发展趋势。比如，让植物自然生长，照顾它们就行了，不需要过度地修剪和整齐地排列。类似地，让人们自由发展，而不是通过繁琐的管制、政策干涉等方式，去要求、指导人们的行为。这种无为而治的法则旨在维持一个有秩序的、和平的治理状态，却不需要大规模的干预和指挥。"}

4.1.2. 分词器处理文本

我选择的分词器是 Qwen2-0.5B-Instruct。原始的词表大小是151643，官方加了三个特殊字符。在预训练的时候，给每个句子的结尾加上 "<|endoftext|>"。然后按 seq_len+1=257 切块（每块包含 input + target 的右移），然后按 95/5 比例分为训练集和验证集。

4.1.3. 实现模型

多头自注意力机制的代码如下：

class CausalSelfAttention(nn.Module):"""带mask的多头自注意力机制"""def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):super().__init__()assert d_model % n_heads == 0self.n_heads = n_headsself.d_head = d_model // n_headsself.qkv_proj = nn.Linear(d_model, 3* d_model, bias=False) # qkv合并self.out_proj = nn.Linear(d_model, d_model, bias=False)self.attn_dropout = nn.Dropout(dropout)self.resid_dropout = nn.Dropout(dropout)def forward(self, x: torch.Tensor) -> torch.Tensor:B, T, C = x.shape # batch, seq_len, d_modelqkv = self.qkv_proj(x) # (B, T, 3*C)q, k, v = qkv.split(C, dim=-1) # 各 (B, T, C)# 多头，(B, T, C) -> (B, n_heads, T, d_head)q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)# QK^T / sqrt(d_head)scale = math.sqrt(self.d_head)attn = (q @ k.transpose(-2, -1)) / scale # (B, n_heads, T, T)# 上三角矩阵，不含对角线。把上三角置为-inf，softmax后趋近于0causal_mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()attn = attn.masked_fill(causal_mask, float("-inf"))# softmaxattn = F.softmax(attn, dim=-1)attn = self.attn_dropout(attn)# *Vout = (attn @ v).transpose(1, 2).reshape(B, T, C)return self.resid_dropout(self.out_proj(out))

前馈神经网络的代码如下：

class FeedForward(nn.Module):"""前馈网络Linear -> GELU -> Linear"""def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):super().__init__()self.net = nn.Sequential(nn.Linear(d_model, d_ff),nn.GELU(),nn.Linear(d_ff, d_model),nn.Dropout(dropout))def forward(self, x: torch.Tensor) -> torch.Tensor:return self.net(x)

单个decoder块用代码表示如下：

class TransformerBlock(nn.Module):"""单个block块，transformer decoder"""def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1):super().__init__()self.ln1 = nn.LayerNorm(d_model) # 归一化，防止计算爆炸self.attn = CausalSelfAttention(d_model, n_heads, dropout)self.ln2 = nn.LayerNorm(d_model)self.ffn = FeedForward(d_model, d_ff, dropout)def forward(self, x: torch.Tensor) -> torch.Tensor:x = x + self.attn(self.ln1(x)) # 残差x = x + self.ffn(self.ln2(x))return x

整体的模型结构如下图所示：

用代码表示：

class MiniGPT(nn.Module):"""Decoder-only GPT 模型"""def __init__(self,vocab_size: int,seq_len: int,d_model: int,n_heads: int,n_layers: int,d_ff: int,dropout: float,):super().__init__()self.seq_len = seq_lenself.d_model = d_modelself.token_emb = nn.Embedding(vocab_size, d_model)self.pos_emb = nn.Embedding(seq_len, d_model)self.emb_dropout = nn.Dropout(dropout)self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, dropout)for _ in range(n_layers)])self.ln_final = nn.LayerNorm(d_model)self.lm_head = nn.Linear(d_model, vocab_size, bias=False)self.apply(self._init_weights)def _init_weights(self, module: nn.Module): # 不用自带的初始值if isinstance(module, nn.Linear):nn.init.normal_(module.weight, mean=0.0, std=0.02)if module.bias is not None:nn.init.zeros_(module.bias)elif isinstance(module, nn.Embedding):nn.init.normal_(module.weight, mean=0.0, std=0.02)def forward(self, input_ids: torch.Tensor) -> torch.Tensor:B, T = input_ids.shape # batch, seq_lenassert T <= self.seq_len, f"输入长度 {T} 超过最大序列长度 {self.seq_len}"positions = torch.arange(T, device=input_ids.device)x = self.emb_dropout(self.token_emb(input_ids) + self.pos_emb(positions))for block in self.blocks:x = block(x)x = self.ln_final(x)return self.lm_head(x) # (B, T, vocab_size)def count_parameters(self) -> int:return sum(p.numel() for p in self.parameters() if p.requires_grad)

4.1.4. 训练

比如现在有一段话 abcd，训练的目的就是拿 abc 预测出 bcd。先加载已经处理好的预训练文本：

class TokenDataset(Dataset):"""从 .pt 文件加载加载数据集每个样本是 (seq_len+1, ) 的token序列input = sample[:-1]target = sample[1:]"""def __init__(self, pt_path: Path):ckpt = torch.load(pt_path, weights_only=True)self.data = ckpt["data"]self.vocab_size = ckpt["vocab_size"]self.seq_len = ckpt["seq_len"]logger.info(f"加载数据集：{pt_path.name}，共 {len(self.data):,} 个样本")def __len__(self):return len(self.data)def __getitem__(self, idx):chunk = self.data[idx]return chunk[:-1], chunk[1:]

然后把input喂给模型，再把结果和target做loss计算，再更新梯度：

input_ids = input_ids.to(device)targets = targets.to(device)logits = model(input_ids) # transformer计算后的结果B, T, V = logits.shape # (batch_size, seq_len, vocab_size)# 自回归损失：每个位置预测下一个tokenloss = loss_fn(logits.view(B*T, V), targets.view(B*T))optimizer.zero_grad()loss.backward()# Gradient clipping：将梯度的 L2 范数截断到 grad_clipgrad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)optimizer.step()scheduler.step()

4.2. 采样策略

常见的有四种策略。"greedy"- 贪心，每步取 argmax"temperature" - 温度采样，temperature 参数控制分布平滑度。温度越小越保守。"top_k" - Top-K 采样，从最高 K 个 token 中随机采样"top_p" - Top-P（Nucleus）采样，从累积概率 ≥ p 的最小集合中采样

for _ in range(max_new_tokens):# 如果序列超过 seq_len，截取最后 seq_len 个 token 作为上下文context = generated[:, -seq_len:]logits = model(context)# (1, T, V)next_logits = logits[0, -1, :] # (V,) 只取最后一步的预测if strategy == "greedy":# 贪心：直接取最高分next_token = next_logits.argmax(dim=-1, keepdim=True)# (1,)else:# 采样前统一做 temperature 缩放next_logits = next_logits / max(temperature, 1e-8)if strategy == "top_k":# Top-K：将低于第 K 大的 logits 置为 -infvalues, _ = torch.topk(next_logits, top_k)threshold = values[-1]next_logits = next_logits.masked_fill(next_logits < threshold, float("-inf"))elif strategy == "top_p":# Top-P（Nucleus）：按概率从大到小排序，累积到 ≥ p 时截断sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)cumprobs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)# 找到刚好超过 p 的位置，将该位置之后的 token 排除sorted_indices_to_remove = cumprobs > top_p# 保留至少 1 个 token：右移一位，使超过阈值的第一个位置不被移除sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()sorted_indices_to_remove[0] = Falseindices_to_remove = sorted_indices[sorted_indices_to_remove]next_logits[indices_to_remove] = float("-inf")elif strategy == "temperature":pass# 只做 temperature 缩放，不做 top-k/p 截断probs = F.softmax(next_logits, dim=-1)next_token = torch.multinomial(probs, num_samples=1)# (1,)# logger.info(f"next_token: {next_token}")if next_token.item() == 151643:breakgenerated = torch.cat([generated, next_token.unsqueeze(0)], dim=1)# (1, T+1)return generated[0]# (total_len,)

预训练的效果如下：

4.3. SFT

SFT的数据格式如下所示：

{"conversations": [{"role": "user", "content": "你好"},{"role": "assistant", "content": "你好！"},{"role": "user", "content": "再见"},{"role": "assistant", "content": "再见！"}]}

需要将原始文本处理成：

<|im_start|>user你好<|im_end|><|im_start|>assistant你好！<|im_end|><|im_start|>user再见<|im_end|><|im_start|>assistant再见！<|im_end|><|endoftext|>

然后在计算loss时和预训练有区别，SFT的loss只计算回答的：

再见！<|im_end|><|endoftext|>

以预训练的流程为基础，修改以下几点适配SFT：

1. 处理文本的代码，原先是给每句话的结尾加上<|endoftext|>，变成了加上<|im_start|>、<|im_end|>。预训练一段话是能切分的，但是sft需要用完整的文本去训练。处理sft文本时，长度需要与预训练的时候一样。
2. loss计算只算回答的
3. 预训练时学习率设置的1e-4，SFT应该改小比如改成1e-5、避免覆盖了预训练时学到的知识
4. 查看效果时，按照sft的格式先给问题加上<|im_start|>这些标签再发给模型

核心是更改源数据的处理代码：

def process_single_dialog(conversations, tokenizer, seq_len):"""单条对话转 input_ids + 带-100掩码的labels，统一padding到seq_len"""# 脏数据过滤if len(conversations) < 2 or conversations[-1]["role"] != "assistant":return None# 拆分历史对话 + 最后一组问答history_turns = conversations[:-2]final_user = conversations[-2]final_assistant = conversations[-1]# 前缀文本：全部mask=-100（历史对话 + 最后一轮user + assistant开头标记）prefix_text = ""for turn in history_turns:prefix_text += f"{IM_START}{turn['role']}n{turn['content']}{IM_END}n"prefix_text += f"{IM_START}{final_user['role']}n{final_user['content']}{IM_END}n"prefix_text += f"{IM_START}{final_assistant['role']}n"# 回答文本：参与loss计算（回答内容 + 结束标记 + EOS）response_text = f"{final_assistant['content']}{IM_END}{EOS_TOKEN}"# 分词prefix_ids = tokenizer.encode(prefix_text, add_special_tokens=False)response_ids = tokenizer.encode(response_text, add_special_tokens=False)full_input_ids = prefix_ids + response_ids# 超长直接丢弃if len(full_input_ids) > seq_len:return None# 构造掩码标签labels = [-100] * len(prefix_ids) + response_ids# 统一padding到seq_len，和预训练数据集保持定长格式pad_len = seq_len - len(full_input_ids)full_input_ids += [0] * pad_lenlabels += [-100] * pad_lenreturn (torch.tensor(full_input_ids, dtype=torch.long),torch.tensor(labels, dtype=torch.long))

sft后的效果如下：sft后的回答是

相比于预训练的回答

更加适合人阅读了。

5. 小结

从推导梯度，到亲手把Transformer的矩阵乘法跑通，再到看着模型说出像样的回答，这件事终于从“想做”变成了“做到”。

很早之前就想做这个小项目了，可是很长一段时间里，精力总是不够用。不是真的没时间，而是心被太多东西占着。那时候总觉得，要做的事情必须有用，必须能被认可。于是很多没有直接回报的念头，就被一次次搁置了。包括这个手搓minigpt、练电吉他等等。

于是慢慢的开始做减法，少做一些无关紧要的事情，开始去做更多自己想做的。

嗯，舒了口气。