从零手撸小型LLM：Transformer原理与训练全流程实战

从零手撸一个LLM到底难不难说实话如果目标是复现某大厂动辄几千亿参数的那种模型那我劝你趁早放弃。但如果你只是想彻底搞懂大语言模型内部的运作原理亲手训练一个几千万参数、能跟你对话的小型模型这条路完全走得通而且收获远比刷十篇论文、看一百个视频都大。这篇不是论文复现笔记而是我自己从环境搭建、数据处理到模型训练、效果调优全流程走下来的一份完整实操记录。内容会一直覆盖到出现loss不降显存爆炸生成质量稀碎这些典型坑时怎么排查。无论你是刚入门深度学习的学生还是想在公司内部做内部工具的开发工程师这篇文章都能让你少走几个月的弯路。我尽量不用教科书式语言见到公式就解释见到代码就说为什么这么写。1. 整体设计思路你的第一个LLM到底该做成什么样1.1 先明确目标不是复刻大模型是复刻原理很多人一听到构建自己的LLM脑子里想的是对标那些超大模型。这是最大的误区。真正适合从零入门的目标是训练一个参数规模在几千万级别的decoder-only语言模型让它能够在给定上下文的情况下生成逻辑相对连贯、语法基本正确的文本。我给自己定的验收标准有两条模型能根据今天天气真不错我们去这样的开头续写出通顺的句子模型在简单的问答式任务上能够泛化出一定的记忆推理能力而不是单纯背诵训练数据把目标定成这样硬件、数据、代码的实现难度都会掉到一个可控的范围。整体技术路线也非常清晰Transformer架构 自回归语言建模 亚词级分词。这个路线跟当前主流LLM完全同源只是规模小了无数倍。1.2 选型决策为什么选decoder-only为什么不从零写CUDA做架构选型时我其实纠结过一阵子。最初的方案有三个encoder-decoder架构、纯decoder架构、以及直接用现成大模型做微调。后来我全部推翻了。先说为什么不做微调。微调虽然省事但你学不到tokenizer怎么构建、attention怎么张量化、训练时的梯度问题怎么处理这些东西才是理解LLM的关键。再说为什么选decoder-only而非encoder-decoder。语言模型生成时的核心机制是自回归即每一步只用已经生成的token来预测下一个tokendecoder架构天然契合这个逻辑实现起来也最简单。BERT那种encoder架构适合做理解任务不是我们要的生成式LLM。至于不自己写CUDA算子纯粹是性价比问题。PyTorch底层已经封装了高效的attention实现对初学者来说先学会用这些原语搭建模型把注意力放到逻辑正确上等真正需要极致性能了再去啃底层。毕竟我们的目标是搞懂原理不是写一个能跑benchmark的推理引擎。2. 环境准备与数据工程地基不牢后面全白干2.1 硬件、框架与依赖清单训练几千万参数的模型其实不需要多吓人的硬件配置。我实际跑通的这套方案使用单张消费级显卡也能完成显存在8GB以上就行。如果你的机器连独立显卡都没有也可以退一步用云平台或者干脆把模型参数量再砍一半。我的建议配置是GPU8GB显存起步最好是16GB比如消费级的中端卡就完全够跑内存16GB以上因为数据预处理和词汇表构建环节挺吃内存的框架PyTorch 2.x配合HuggingFace生态的tokenizers库关键Python包transformers只用它加载tokenizer和数据集、numpy、tqdm提示不要一开始就在环境上死磕。Python版本用3.9~3.11之间都行重点是torch的CUDA版本跟你显卡驱动要匹配否则装好了也用不上GPU。整个环境的搭建时间控制在半天内超过这个时间就要检查是不是驱动和版本的问题。我见过太多人在conda里折腾三天装环境最后项目还没开始就放弃了。2.2 语料获取与清洗的三个核心原则训练LLM的语料质量和数量同等重要。我第一次做的时候直接找了个开源中文语料就往里灌结果模型的输出全是乱码和重复废话。后来反复清洗才明白语料处理有黄金三原则第一去噪比去重更重要。HTML标签、特殊符号、超链接、无意义的重复标点这些噪音对模型的影响极大。模型会学到输出这样的垃圾。清洗后的语料应该看起来就像一篇干净的网页文章纯文本、无杂质。第二混入多领域数据。代码片段、对话记录、百科条目、新闻内容不同领域的数据可以让模型的泛化能力更强。我当时混合了三个公开数据源比例上以百科和新闻为主这样模型生成的句子既有知识性又贴近日常表达。第三长度过滤。太短的句子比如少于20个字说明含金量低太长的段落比如超过1000字可能是文档拼接都建议过滤掉。处理后的理想样本长度在200到800字之间。数据预处理这一步值得多说两句。我用的是内存流式处理先把语料全部读进内存遍历一遍做清洗然后按照长度切分句子最后打乱顺序。这个流程看起来简单但一旦数据量上来代码里的一个低级错误比如忘了处理空行就能让你的训练集莫名其妙多出一堆无效样本。3. 模型架构实现Transformer核心模块逐行拆解3.1 Embedding与位置编码让模型理解词语和顺序模型的输入是一串token id首先要把它们映射成向量这就是nn.Embedding做的事。这个环节有两个关键参数vocab_size词汇表大小和d_model向量维度。对千万级参数的小模型来说词汇表建议控制在5000到一万之间向量维度采用256或384都可以。映射到向量之后模型还需要知道词语在句子里的位置。Transformer本身没有时序概念如果不加位置信息模型会把猫追狗和狗追猫当成完全一样的输入。早期的绝对位置编码用的是sin/cos函数那套方案数学上很漂亮但实际上训练起来不太友好我建议直接用可学习的nn.Embedding来做位置编码实现简单效果也够用。import torch import torch.nn as nn class GPTConfig: vocab_size 8000 block_size 128 # 最大上下文长度 n_layer 6 n_head 8 d_model 384 dropout 0.1 class TokenAndPositionalEmbedding(nn.Module): def __init__(self, cfg): super().__init__() self.tok_emb nn.Embedding(cfg.vocab_size, cfg.d_model) self.pos_emb nn.Embedding(cfg.block_size, cfg.d_model) self.dropout nn.Dropout(cfg.dropout) def forward(self, idx): _, T idx.shape pos torch.arange(T, deviceidx.device).unsqueeze(0) return self.dropout(self.tok_emb(idx) self.pos_emb(pos))注意这里有个细节位置编码的张量形状是(1, T, d_model)在广播机制下每一个batch里的序列都会拿到相同的位置向量。这个相加操作就是在告诉模型你现在处理的这个token它所在的位置是第几个。3.2 多头注意力机制核心中的核心Attention可以说是整个LLM的发动机。它的核心思路就一句话每个token在生成时应该把上下文里其他token的信息按重要程度加权融合进来。公式不复杂[ Attention(Q,K,V) softmax(\frac{QK^T}{\sqrt{d_k}})V ]这里Q、K、V分别代表查询、键、值矩阵。生活化地理解Q像是在问我需要什么信息K像是在说我手里有什么信息两者做点乘算出每个token跟当前token的相关性分数经过softmax归一化变成权重最后用权重去加权V得到融合后的输出。除以根号d_k是为了防止点积结果过大导致softmax梯度消失。多头机制就是把Q、K、V切分成多组让不同组关注不同维度的关系。比如一个头负责语法关系另一个头负责语义相似度最后拼接在一起。实现时有一个关键易错点切分后要把(B, T, n_head, head_dim)先转成(B, n_head, T, head_dim)让注意力计算在最后的两个维度上做矩阵乘法。class MultiHeadAttention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head cfg.n_head self.d_model cfg.d_model self.q nn.Linear(cfg.d_model, cfg.d_model) self.k nn.Linear(cfg.d_model, cfg.d_model) self.v nn.Linear(cfg.d_model, cfg.d_model) self.out nn.Linear(cfg.d_model, cfg.d_model) def forward(self, x): B, T, C x.shape q self.q(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) k self.k(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) v self.v(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) att q k.transpose(-2, -1) / (C // self.n_head) ** 0.5 att att.softmax(dim-1) y (att v).transpose(1, 2).reshape(B, T, C) return self.out(y)做decoder模型时还必须加一个下三角掩码确保每个位置只能看到自己之前的信息。这一步常被忽略如果漏了模型在训练时就偷看了未来的答案生成时性能会暴跌。3.3 前馈网络与层归一化稳定训练的隐形功臣Attention的每个输出后面都会接一个两层的前馈网络FFN。这个结构用非线性变换对融合后的信息做进一步加工通常中间维度会扩大到四倍再加一个ReLU或GELU激活函数。没有FFN多层attention堆叠后模型的表达能力提升非常有限。层归一化放在每层的Attention和FFN之前也就是Pre-LN结构。这跟最早的Transformer论文用的Post-LN不同Pre-LN在训练时更稳定梯度更健康也是现在主流大模型的标配。所有归一化操作都作用在d_model维度上而不是batch维度因为序列长度和batch大小常常变化基于特征维度的归一化更扛折腾。最后一层是线性输出层把d_model维度映射回vocab_size再经过softmax得到每个token的概率分布。这里有个共享权重的细节三个全连接层的权重其实是同一个张量共享后参数量能省下不少。4. 完整训练流程实操从数据加载到模型收敛4.1 Tokenizer构建与数据加载管线的落地Tokenization是把文本切成小块再映射成数字ID的过程。我个人强烈建议用BPE字节对编码方案它既能覆盖绝大多数常用词又没有字母都会被切开的低效率。直接用tokenizers库里的BPE类训练一个自己的分词器几分钟就能搞定。训练tokenizer时喂进去的语料最好是原始文本本身不要提前截断。BPE会先从字符级别开始逐步合并出现频率最高的相邻字节对。我在一个800MB规模的中文语料上训练设置vocab_size8000训练完看一眼词汇表会发现常用汉字、人工智能这类高频词都成了一个独立的token。数据加载我选择了最简单的思路把清洗好的语料切分成固定长度我用的序列长度是128的样本每64个样本打包成一个batch。这样做的好处是batch内序列长度一致方便张量运算。加载管线大概长这样from torch.utils.data import Dataset, DataLoader class TextDataset(Dataset): def __init__(self, tokens, block_size): self.tokens tokens self.block_size block_size def __len__(self): return len(self.tokens) - self.block_size def __getitem__(self, i): x torch.tensor(self.tokens[i:iself.block_size], dtypetorch.long) y torch.tensor(self.tokens[i1:iself.block_size1], dtypetorch.long) return x, y这个数据集很有意思的特征是输入x和标签y错开了一个位置。模型看到前128个token要预测第129个。每一步预测的都是下一个token这就是自回归语言建模的数据形态。4.2 训练超参数与优化器选择训练超参数这块我直接给出实测下来比较靠谱的组合你照抄就能跑优化器AdamW学习率初始为3e-4学习率调度预热(warmup) 2000步后做余弦退火训练批次batch size为64总共训练约2万步梯度裁剪max_grad_norm设为1.0正则化Dropout设为0.1权重衰减设为0.1AdamW最关键的设置在权重衰减上。原版Adam的权重衰减是直接在梯度里减的效果不理想。AdamW的做法是把权重衰减跟梯度的自适应学习率解耦直接写进参数更新的最终值里。用PyTorch内置实现即可注意不要把bias和归一化层的参数也加了权重衰减需要单独分出来。关于训练时长的预期管理在单张消费级显卡上一个3800万参数的模型batch size为64序列长度128大概每小时能跑1500步。2万步大约需要十几个小时。首次训练不用跑满让我看一下1000步之后的loss变化趋势来判断是否值得继续。optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay0.1) def get_lr(step, warmup_steps2000, max_steps20000, max_lr3e-4): if step warmup_steps: return max_lr * (step 1) / warmup_steps progress (step - warmup_steps) / (max_steps - warmup_steps) return max_lr * 0.5 * (1 (1 - progress) * progress)学习率从零开始线性爬升到峰值再缓慢衰减到接近零。预热阶段防止模型在初始几万步内因为过大的梯度而震荡之后让学习率逐渐变小模型才能精细地逼近最优解。4.3 生成策略与采样参数让模型的输出更接近真人训练完成之后你会拿到一个能预测下一个token的模型。但怎么从概率分布里挑一个token输出里面的门道也不少。最简单的做法是贪心解码每次选概率最高的token。这样生成的文本稳定、逻辑连贯但是非常平庸经常出现重复。更好的办法是带温度的采样加top-k约束。温度T的作用是把token概率分布拉平或者压陡。T大于1时分布更平滑罕见词也有机会被选中输出多样性增加。top-k的做法是只从概率最高的前k个token里抽样防止选了太离谱的词。这两个手段组合起来是我实测中最有效的配置def generate(model, idx, max_new_tokens, temperature0.8, top_k50): for _ in range(max_new_tokens): idx_cond idx[:, -cfg.block_size:] logits model(idx_cond)[:, -1, :] / temperature if top_k is not None: v, _ logits.topk(top_k) logits[logits v[:, -1:]] -float(Inf) probs torch.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat([idx, idx_next], dim1) return idx初次训练完模型生成的句子可能还是有一堆语法错误。不用气馁这说明训练数据量或者迭代步数还不够。我记得自己第一次跑完2万步时模型已经能生成这本书的内容非常丰富值得每一个人认真阅读这种接近自然语言的句子那一刻还是有点小兴奋的。5. 常见问题与排查技巧实录5.1 Loss不下降先从这三处查起训练到一定阶段loss断断续续地卡住不动这是最常见的现象。我的排查顺序是第一看数据管线。先跑一次next(iter(dataloader))把样本打印出来确认x和y确实错位了一个token。很多诡异的loss不降问题源头都是标签和输入完全对齐了模型学到的只是恒等映射。第二检查学习率。如果loss从一开始就纹丝不动很可能是学习率太大导致模型发散或者是太小导致学习缓慢。用一个简单方法验证把学习率设成一个极大值比如0.01跑100步。如果loss直接NaN或者暴涨说明代码本身能更新问题在于学习率调度如果loss仍然不变化问题大概率出在模型前向传播上。第三回归验证模型初始输出。模型初始化后如果输入的logits分布非常不平滑softmax之后某个token的概率极大训练一开始就陷入局部最优。这个情况可以把初始化方式换成更小方差的方案或者把最后一层线性层的初始化scale调小一点。5.2 显存溢出参数量没变为什么内存爆了有一个我踩过很深的坑模型参数量只有几千万按理说显存完全够但训练时还是OOM。原因在于显存消耗的大头是激活值activations不是模型参数。序列长度128还好一旦把上下文长度加长到512中间特征的形状会成倍膨胀显存占用跟着飙升。降低显存占用的优先级如下缩小batch size降到16或8这是最直接的办法开启梯度累积用小batch累积梯度再更新效果等价于大切size参考PyTorch的torch.utils.checkpoint做激活检查点用计算换显存实在不行降低序列长度或减少层数梯度累积的写法是用loss.backward()后不立即step()等累积几步后再更新。注意loss除以累积步数否则梯度会放大好几倍。这个细节我漏过一回模型直接飘了。5.3 生成内容空洞且重复的调优心得模型能跑了但生成的全是是的是的下面是下面是下面这种复读机话术这属于典型的解码退化问题。通常不是模型不行而是采样策略没配好。我从实践中总结出来的生效组合是temperature0.7、top_k50、top_p0.95三管齐下。如果调整采样后还是重复那就要考虑训练层面的问题比如训练数据里本身有很多重复文本或者模型容量不够。我曾经在清洗语料时偷懒没做去重结果模型有个坏毛病遇到通用词汇时会反复以最高概率输出同一个token整个句子就卡在循环上。把语料里明显重复的段落删掉重新训练一些步数之后这个问题缓解了很多。还有一个容易忽略的陷阱生成时上下文长度超过了训练时的block_size。模型看到的位置编码是它从来没见过的位置ID输出会变得很奇怪。生成时一定要先截断输入序列保证长度不超过训练时的窗口。6. 项目扩展方向你的第一个LLM还能往哪里走模型训练完成、能流畅生成句子之后这个项目的价值还远没有榨干。我尝试过几个扩展方向按性价比排序供你参考第一个方向是指令微调SFT。准备几千条问题-回答对数据在预训练模型基础上用较小的学习率继续训练模型就会从文本续写师变成对话助手。这个阶段的数据质量远大于数量我当时手工攒了3000条高质量问答效果比直接用10万条网上爬来的数据好得多。第二个方向是加一层LoRA低秩适配。如果你后续想针对特定领域比如法律、医学、编程继续优化用LoRA只需要训练原模型参数量的2%~3%成本极低。LoRA相当于在原有权重旁边加一个低秩的旁路矩阵训练时只更新旁路原始权重保持不动。第三个方向是从工程角度打磨把模型导出成ONNX格式写一个简单的服务接口部署成一个本地小工具。虽然参数量小但应对一些固定的文本生成任务体验其实很不错。我自己后来做了一个简单的文档自动摘要服务运行成本和响应速度都让我很满意。如果你想让效果再上一个台阶可以尝试把模型的上下文长度从128逐步扩到512。这个过程并不用重新训练只需把位置编码层重新初始化然后用一段长文本数据做少量步数的二次训练。模型的长期依赖推理能力会明显提升生成的文本在逻辑连贯性上的进步非常直观。最后再分享一点个人体会构建自己的LLM收获最大的不是模型本身而是你在踩坑过程中建立起来的直觉。什么时候该调学习率、什么时候该查数据、生成效果差到底是模型的问题还是解码策略的问题这些判断力无法从任何教程里直接学到非得亲手把一套流程走完。朋友圈里晒两张loss曲线很容易真正有价值的是你能在深夜盯着一屏报错信息冷静地说出哦原来是batch维度写反了的那一刻。