1. 前言
这篇文章完整记录了一位后端开发者,基于GitHub上的开源项目minimind,从零开始手写一个小型GPT模型的实践过程。说实话,在AI技术飞速发展的当下,每天调调大模型的接口、写写提示词,总觉得有点不过瘾。那个驱动一切的语言模型,究竟是如何打造出来的?出于好奇,我决定亲手尝试一番。这篇笔记,就是从小白视角出发,到成功训练出一个miniGPT的完整记录。
2. 机器学习
2.1. 为什么需要机器学习
作为后端开发,日常工作流程通常是:明确需求、设计架构、开发测试、运维上线。本质上,就是把复杂的业务逻辑翻译成代码。但当规则的复杂度超出人力所能清晰梳理的范围时,机器学习便应运而生。简单来说,规则难以写清楚时,就让模型从数据中自己“学习”出来。
2.2. 机器学习究竟是什么?
机器学习的核心公式可以简化为 y = wx + b(当然实际场景远不止如此)。它的本质是利用已知数据去预测未知数据。从这个基础公式也能看出,大语言模型之所以火爆全球——过去每个新领域都需要重新训练 w 和 b,但大模型将这两个参数优化到了极致,用户只需改变输入 x,就能得到期望的输出 y。某种程度上,这确实是“力大砖飞”的体现。
2.3. 一个最小化的机器学习示例
前面铺垫完背景,现在来看一个最简洁的机器学习流程。这里我们使用机器学习来拟合一个多项式函数 y = a*x³ + b*x² + 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结构,仅使用了Transformer的Decoder部分。而Decoder的核心,正是多头自注意力机制。
3.2. Multi-Head Self-Attention
多头自注意力机制究竟解决了什么问题?先来探讨为什么需要“自注意力”,以及为何要“多头”。在更早期的语言模型中,RNN是主流方案,但它有一个致命缺陷——难以处理长距离依赖。举例来说:“狗没有穿过街道,因为它很累。”句中的“它”显然指的是狗,但用RNN处理这种远距离指代关系,效果往往不理想。于是,自注意力机制被引入——通过一个二维矩阵,记录每两个token之间的相关性。公式如下:
Q*K^T图示:
而“多头”的意义在于,让模型能够从不同角度理解句子——比如同时关注语法结构、近义词关系、句法成分等多个子空间的信息。单个注意力头只能捕捉一种关系,而多头机制让模型学会了多维度思考。
GPT这类模型的核心逻辑是“根据上文预测下文”,因此需要一种机制,让当前token只能看到自己和它前面的内容,不能提前“偷看”后面的信息。这就是Attention Mask的作用。
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"]def __len__(self):return len(self.data)def __getitem__(self, idx):chunk = self.data[idx]return chunk[:-1], chunk[1:]
然后将输入喂给模型,计算输出与目标的损失,再更新梯度:
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,选择概率最高的token。
- Temperature(温度采样):通过temperature参数控制分布的平滑度,温度越低,输出越保守。
- Top-K:从概率最高的K个token中随机采样。
- 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,)if next_token.item() == 151643:breakgenerated = torch.cat([generated, next_token.unsqueeze(0)], dim=1)# (1, T+1)return generated[0]# (total_len,)
预训练的效果如下:
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|>
与预训练不同,SFT阶段的损失计算只针对回答部分:
再见!<|im_end|><|endoftext|>
在预训练流程基础上,需要做以下调整:
- 文本处理:原来只在每句话结尾加<|endoftext|>,现在改为添加<|im_start|>和<|im_end|>标签。预训练时文本可以切片,但SFT需要完整对话进行训练,同时保持与预训练相同的长度。
- 损失计算:只计算回答部分的损失。
- 学习率:预训练时设为1e-4,SFT阶段适当降低,比如调整到1e-5,避免覆盖预训练阶段学到的知识。
- 推理测试:按照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是这样,练电吉他也是如此。
后来慢慢学会做减法,少做一些无关紧要的事,把时间留给真正想做的事。
嗯,终于松了一口气。
