最近在跟一门以“从零开始构建语言模型”为核心的研究生课程课程代号就不报了反正圈内人一听就知道是哪个。整个学期的作业设计得很硬核尤其是Assignment 1要求我们不依赖现成的深度学习框架封装把一个小型GPT模型的训练全链路手工实现出来。我花了两周时间把实验部分完整跑完过程中踩了不少坑也理顺了很多平时用框架时根本没注意过的细节。这篇笔记主要是记录我在Assignment 1的Experiments环节里做了什么、怎么做的、遇到了什么问题以及最后怎么定位和解决的。如果你是准备做类似课程作业的学生或者自己折腾过小型语言模型训练这篇应该能帮你省掉不少摸排时间。我会尽量把“为什么这样做”也讲清楚不只是贴结论。1. Assignment 1到底在考核什么1.1 作业拆解数据、模型、训练、实验四块先把作业结构捋一下。Assignment 1整体分两大块代码实现和实验报告。代码实现又包含四个子任务数据加载与预处理、模型前向计算、训练循环含梯度更新、以及评估与生成。实验部分要求我们在实现完成后基于一个小型数据集完成若干组对比实验并回答一些关于模型行为的问题。从标题里的“Experiments Implement”可以看出这门课的侧重点并不仅仅是把代码跑通而是强调“能设计实验、能解读实验结果”。很多同学会把精力全放在模型实现上结果到了写实验报告时才发现自己的实验设计漏洞百出变量没控制住对照也不成立。所以这里先提醒一下实验设计和代码实现要同步推进甚至在写代码前就要想好你要验证什么。作业的模型规模被刻意设计得比较小embedding维度通常在128到256之间transformer层数2到4层训练语料是几十MB级别的公开数据子集。这意味着在一张消费级显卡上就能完成训练但也意味着模型的容量和数据的多样性都有限实验结果不会特别“漂亮”需要靠合理的分析和可视化来支撑论点。1.2 为什么这类“从零实现”作业值得认真做我见过不少同学直接用PyTorch的nn.Transformer模块搭一个模型几行代码就能跑训练然后写报告。但这么做基本是浪费时间——作业考察的本来就不是“你会调用现成API”而是“你对每个模块内部机制的理解程度”。举个例子注意力机制里的scale操作在框架里就是一个除法参数。但如果你从零实现就必须理解为什么需要除以sqrt(d_k)因为点积的方差会随维度增大而增大如果不scalesoftmax输入的绝对值过大梯度会变得非常小训练几乎停滞。再比如层归一化的可学习参数gamma和beta你不仅要实现它还得理解它解决的是什么问题——当输入分布的均值方差在训练中不断变化时网络层的输出分布会漂移这会让训练不稳定。这些细节只有在手写时才会真正刻进脑子里。另一个实际收益是调试能力。当你的模型loss出现nan或者生成结果全是重复token时你面对的是自己的代码没有任何框架封装来兜底。你必须学会通过检查中间量的分布、梯度范数、学习率曲线等来定位问题。这种排查能力在现代AI工程中非常值钱因为在真实业务里模型训练出现诡异问题时没有现成答案可查。2. 从原始语料到训练batch数据管线的那些细节2.1 Tokenizer的选择直接决定后期的坑多不多Assignment 1一般不会要求你从头训练BPE分词器通常会给一个预训练好的tokenizer文件或者要求用现成的tokenizer库来做。但即便如此数据预处理的环节依然有不少细节值得记录。我这次使用的是BPEByte Pair Encoding分词方式。BPE的核心逻辑是从字符级别开始反复统计最常出现的相邻符号对把它们合并成一个新符号迭代执行预先设定的合并次数。最终得到的词表既能覆盖常见单词也能通过子词组合处理未登录词。比如“playing”可能被分成“play”和“ing”这样即使测试时遇到了“played”也能通过子词拼接来近似表达。我在处理原始语料时遇到的一个问题是语料里包含大量HTML标签和特殊符号。如果直接清洗掉所有非字母数字字符句子边界信息会丢失如果保留太多分词后会有大量低频的“噪音token”挤占词表空间。我的处理方案是分两步先用正则把URL、HTML实体、连续空白符做标准化再依赖BPE分词器自然地处理剩余特殊符号。实测下来词表里的特殊token数量从几千降到了几百训练速度也提升明显。2.2 定长切分与Attention Mask的设计逻辑Transformer模型要求输入是固定长度的token序列不考虑变长batch的特殊实现。作业里通常使用的是“将整个语料直接切分成固定长度块”的做法而不是按句子或段落切分原因是后者会造成大量样本被padding浪费计算资源。而前者每个训练样本都是满长的GPU利用率更高。但直接切分会带来一个问题序列边界处可能会切断一个句子导致模型要基于不完整的上下文做预测。这对于语言模型来说还不算致命伤因为语言模型本质上就是在做“给定前面的token预测下一个token”的任务只要上下文长度足够切断的句子不会对训练目标造成太大破坏。Attention Mask这个东西我在实现时纠结了一会儿。标准做法是训练时因为所有序列都是有效内容mask只需要一个下三角掩码来防止“看到未来”。但有些实现会把“PAD位置不参与注意力计算”的mask也一并考虑。由于我们用的是定长切分而不是padding所以不需要PAD mask只需要因果掩码。这里要提醒的是如果以后换成按句子batch的padding模式千万别忘了把PAD位置在注意力计算中屏蔽掉否则模型会学到把padding当作有效内容。2.3 Dataloader的预取与混洗策略在单卡训练场景下dataloader的瓶颈不太明显但几个小优化依然值得做。我使用的是DataLoader的num_workers设置和prefetch_factor参数让数据加载和GPU计算在时间上交叠。比较关键的一点是epoch内样本的混洗策略。如果语料被切成了几千个等长块那么同一批相邻的样本原本在语料中也是相邻的它们之间的梯度方向可能高度相关。如果训练时按顺序遍历模型会对语料某一部分过拟合然后突然切换到另一段完全不同的风格loss曲线会呈现周期性的尖刺。混洗之后这种波动会显著缓解。实测同一个超参数配置下混洗相比顺序遍历的收敛速度大约提升了一到两成。另一个隐藏点是num_workers过多时主进程和子进程之间的数据序列化开销会抵消预取的收益。在小数据集上2到4个worker通常就够用了不要盲目追求高并发。3. 手写Transformer核心组件时容易忽略的数学细节3.1 Embedding层与位置编码的组合方式从零实现的第一步是词嵌入层。嵌入矩阵的形状是[vocab_size, embed_dim]查表操作其实就是一个带权重矩阵的线性层只不过输入是token索引而不是one-hot向量。PyTorch里可以用nn.Embedding也可以自己写一个nn.Linear加one-hot的输入但这样太浪费内存不推荐。真正值得注意的是位置编码和词嵌入的组合方式。原版Transformer论文用的是固定三角函数位置编码每个位置的编码向量由sin和cos函数交替生成。它的优点是不同频率的正弦基函数可以提供不同分辨率的相对位置信息而且即使序列长度超出训练时的长度也能用公式外推生成新的位置编码。但如果你用的是可学习位置编码就必须在训练前确定好最大序列长度一旦推理时序列更长就尴尬了。我在这份作业里选择了可学习位置编码因为实现简单且在固定长度的小模型上效果通常不弱于三角函数。但组合时有个细节是直接将词嵌入和位置编码相加还是沿特征维度拼接两者计算量不同信息编码方式也不同。相加是主流做法因为位置信息被调制在词向量的各个维度上而拼接则相当于给每个token额外增加了一倍的维度模型容量需求更大。相加时要注意初始化位置编码的方差不能太大否则会淹没词嵌入的信号导致训练初期loss下降缓慢。我把位置编码初始化为Normal(0, 0.02)效果比较稳定。3.2 注意力机制的循环实现与向量化实现从零实现注意力机制有两种层次一种是写成显式的for循环对每个query遍历所有key清晰展示“query和key做点积→缩放→softmax→加权求和value”的流程另一种是用矩阵乘法一次算完所有位置的注意力权重。作业一般要求先写循环版本并用于梯度检查再改成矩阵版本提高效率。我一开始先写了循环版本用随机输入和数值梯度做了梯度check然后再跑向量化版本确保两者输出一致。这个流程强烈建议保留因为很多实现错误出在维度permute上只有对比才能快速发现。向量化实现中最容易出错的是注意力得分矩阵的转置方向。Q和K的shape都是[batch, num_heads, seq_len, head_dim]计算Q乘以K的转置时需要在最后两个维度上转置也就是permute(0,1,3,2)。有人会不小心写成转置batch维度得到的结果完全错误。这个错在loss上不会立刻爆掉但训练曲线会明显异常。还有一点是Flash Attention的使用场景。作业如果允许使用torch的scaled_dot_product_attention建议直接用因为它内部做了分块计算在长序列下显存占用大幅降低。但为了满足“从零实现”的作业要求最好还是自己先写一版标准版再切换。我在实验环节就是标准实现跑小模型对比实验时再用scaled_dot_product_attention加速。3.3 层归一化与残差连接的前向实现层归一化在Transformer中的位置有讲究。原版架构是layer norm在残差加法之后Pre-LN架构则是把layer norm放在残差分支的起始位置。这篇作业使用的是Post-LN也就是“attention→residual→layer norm”和“FFN→residual→layer norm”的顺序。实验结果普遍表明Pre-LN在深层网络中训练更稳定但Post-LN在浅层小模型上也够用。实现层归一化时核心操作是对最后一维做均值方差归一化然后再做仿射变换。容易出错的是方差计算时是否使用 unbiased 的贝塞尔校正。在归一化中我们使用的是总体方差除以N而不是样本方差除以N-1因为归一化对象是某一层在某个batch内所有特征的分布不涉及对总体的估计使用总体方差更符合这个场景的定义结果也更具可重复性。残差连接的一个坑是“缩放残差”的问题。为了训练稳定性有些实现会给残差分支乘以一个小于1的缩放因子比如0.1这在深层模型里很常见但在作业这种浅层模型里反而会拖慢收敛。我建议初期不用加缩放等碰到稳定性问题时再考虑。前馈网络部分就是一个两层的MLP先把embed_dim升到4倍或某个隐藏维度经过GELU激活再降回embed_dim。GELU函数的实现也有讲究它定义为x乘以标准正态分布的累积分布函数。精确形式里包含erf函数计算开销较大近似形式用tanh函数拟合速度更快在PyTorch里默认用的就是tanh近似。代码实现上就是两步线性加一个激活没有太多陷阱但要注意dropout的位置——通常放在FFN的第二个线性层之前而不是激活之前。4. 训练循环里的那些隐藏细节4.1 学习率warmup与余弦退火的配合训练GPT小模型时学习率调度策略对收敛速度影响非常明显。我使用的是“warmup 余弦退火”的组合前500步学习率从一个极小值线性增长到峰值之后按余弦曲线衰减到峰值的1/10左右。这个策略几乎已经成为小型语言模型训练的默认配置。为什么要warmup因为训练初期模型的参数是随机初始化的如果直接用较大的学习率梯度信号虽然存在但loss landscape在不同方向上差异很大容易被推到某个过于陡峭的区域导致后续训练不稳定甚至发散。先小步走几步让参数落在一个相对平滑的区域再逐步加大步子能明显减少早期震荡。warmup步数的设定有讲究。太短比如50步没有热身效果太长比如2000步则浪费训练时长。我这次的做法是把warmup步数设为总训练步数的1%到2%。训练步数约3万步warmup用了600步峰值学习率设为3e-4效果不错。余弦退火的优势在于训练后期学习率缓慢降低有助于模型在损失曲面的最小值附近做精细搜索。注意峰值学习率和warmup比例不能单独调优它们高度耦合。先固定warmup步数再扫峰值学习率得到的结果会比同时扫两个参数更稳定。4.2 交叉熵损失的计算注意点语言模型的标准训练目标是下一个token预测因此每个样本的标签就是输入序列自身右移一位后的序列。比如输入是token列表[t1, t2, ..., tN]那么标签就是[t2, t3, ..., tN1]其中tN1通常被认为是序列结尾后的第一个token但在训练中我们会把最后一个位置的loss丢弃因为不存在对应的“下一个token”。实现中的常见做法是模型对输入序列的每个位置输出一个logits向量shape为[batch, seq_len, vocab_size]。计算loss时把logits的seq_len维度与labels对齐忽略最后一个位置的输出。在代码上将labels切片为inputs[:, 1:]同时模型输出取前seq_len-1的logits两者对齐后求交叉熵。这里还有一个性能细节在自回归任务中每个位置的logits都对应一个vocab_size维的类别概率标准交叉熵会对所有位置求和再平均。但PyTorch的CrossEntropyLoss默认是基于每个位置的损失平均值而不是基于batch内所有token的损失平均值。在大batch下两者数值上几乎一致在小batch下可能因为样本数的微小差异导致评价指标看起来不同。建议统一以“per_token平均”为准这样在改变batch size时评估指标更一致。这里需要说明的是我沿用了自动化评测脚本和课程基准中常见的per_token平均口径以此保证不同实验设置之间的结果可比性。4.3 GPU显存与Batch Size的平衡术实验阶段模型的参数量大约在几千万级并不算大。真正吃显存的不是参数本身而是前向传播过程中保存的中间激活。比如一个4层、embedding维度256、序列长度512、batch size 32的模型其各层中间激活总量可以轻松超过参数总量好几倍。我一开始尝试batch size 64序列长度512结果在8GB显存的显卡上直接OOM报错。解决思路有两条一是降低batch size但这样理论上会增大梯度噪声削弱训练稳定性二是用梯度累积把64的batch分成4个16的micro-batch每步反向传播后累积梯度攒够4步再更新参数。梯度累积在等效果batch size的前提下训练效果和直接大batch几乎一致。要注意的是累积梯度时学习率需要根据实际累积的步数来做相应调整因为每次参数更新的有效batch变大了梯度方向平均化之后方差减小通常可以略微增大学习率。不过我实测下来在batch size从64降到16、累积4步的情况下保持原学习率依然能收敛只是初期loss下降稍慢。梯度累积的实现还有一个细节如果代码里直接在loss.backward()后调用optimizer.step()那么每步都会更新参数累积逻辑就没有生效。正确做法是只在累积步数达到阈值时才调用step()同时把optimizer.zero_grad()的调用放在参数更新之后、下一个累积周期开始时。另外梯度裁剪( clip grad norm )是个保险项。我一般把全局梯度范数上限设为1.0。这在原生实现里尤其重要因为手写代码里很少会自动处理梯度爆炸的问题而小模型的embedding梯度在训练早期确实容易出现较大的范数。5. 实验设计与结果解读不是跑完就完了5.1 Baseline与变量控制实验部分我的做法是先把一个“标准配置”作为baseline跑通然后每次只改动一个变量。比如第一组实验修改模型层数2层 vs 4层第二组实验修改embedding维度128 vs 256第三组实验修改学习率1e-4 vs 3e-4 vs 1e-3。固定种子是我非常强调的不同配置之间对比时才不会把随机性当成系统差异。如果每次实验不固定随机种子即使使用完全相同的超参数两次训练的最终loss也可能相差几个百分点。对比实验时如果不固定种子误差可能掩盖真实差异导致判断失误。我的习惯是给每个配置跑至少两次取平均曲线如果两次差异过大就说明这个配置本身对种子敏感需要更多重复实验。实验中我发现一个现象模型从2层增加到4层时如果不相应调整学习率loss反而会上升。原因是更深模型的反向传播路径更长梯度尺度变化更大3e-4的峰值学习率在4层模型上会让训练失去稳定性。把峰值学习率降到2e-4后4层模型顺利收敛。这个现象在报告里也可以成为一个很好的分析点。5.2 评估指标不只是loss还有Perplexity和生成质量评估语言模型的经典指标是perplexity困惑度定义为交叉熵损失取指数。perplexity可以直观理解为“模型在每个位置上平均有几种可能的候选token”。如果perplexity是50就相当于模型每次预测的候选集合大约有50个词对于一个小词表来说这个数值说明模型已经学到了一些语言规律。在做生成质量评估时作业要求我们使用greedy decoding和temperature sampling两种方式各生成几段文本。我观察到greedy decoding输出往往陷入重复的短句因为一旦某个路径被选中模型就会沿着高概率的方向继续走很难跳出循环。而temperature较高时采样多样性上升但句子的连贯性会下降。把生成的文本打印出来后我发现模型已经学会了一些局部语法搭配和标点习惯这是一个很有说服力的证据说明模型确实在训练中捕获了语料的结构特征。5.3 从数据角度解释“为什么效果看起来不够好”很多同学跑完实验会觉得困惑loss一直在下降但生成的句子读起来还是语法不通顺怎么解释这里有一个很关键的认知语言模型的训练目标是预测下一个token的概率分布对短文本的局部连贯性建模相对容易但对长程依赖的建模难得多。在小模型、小数据上长程依赖信号非常微弱模型很可能只能记住语料中频繁出现的固定搭配而无法理解上下文中的复杂逻辑关系。在写实验报告时我建议把“训练loss曲线”和“生成样本”并排放置分析两者之间的关系。比如训练loss持续下降但生成文本在词层面仍然不通顺可能是因为训练语料中包含大量非正式文本、噪声数据模型把这些噪声也学进去了。这本身就是一个有价值的发现说明数据质量对生成质量的影响甚至大于模型容量。如果时间允许还可以做一组小规模“干净语料 vs 原始语料”的对比实验来看数据清洗对最终loss的影响幅度。6. 排坑实录现象、根因、修复对照表做Assignment 1时遇到的坑不少很多是有共性的。我把自己踩过的、以及在课上看到别人踩的问题整理成一张表每个问题都给出现象和排查方向供参考。问题现象可能根因排查与修复建议训练loss为nan学习率过大 / Embedding或logits层数值过界先看初始loss是否合理打印梯度的最大值和norm降低学习率或对梯度裁剪loss正常但速度很慢矩阵维度顺序导致大量非连续内存访问确认QKV变换后head维是否在最前面用torch.profiler分析耗时热点生成文本全为重复词模型欠拟合 / 解码策略过于贪心观察验证集loss是否仍在下降改用temperature采样并调高到0.8以上不同batch的训练loss波动大数据混洗不充分或batch内语料风格跳变增大shuffle粒度检查是否误用顺序遍历loss在warmup结束后突然上涨学习率峰值过高将峰值学习率降低约一半检查grad norm是否随训练快速增大验证集loss比训练集低训练集有大量重复句子或验证集混入训练数据检查数据切分的item ID是否严格按811划分样本级去重变更batch size后曲线变化明显未同步调整学习率或评估口径保持per_token平均的loss口径大batch时试探性降低学习率我印象最深的一个问题模型训练到第4000步时梯度范数突然从0.9涨到8.5紧接着loss就变成了nan。排查过程是先dump了最近的输入样本发现里面有一段长度异常且充满连续重复字符的文本然后再看梯度的分布发现embedding层和lm_head层的梯度特别大。解决办法是双重保险加上梯度裁剪同时把训练数据里长度小于10的短文本全部过滤掉。因为短文本在定长切分中会变成一个重复堆叠的模式容易产生极端子词序列这在BPE词表下会合成低频token导致对应行Embedding的梯度异常大。这个现象在作业文档里没有明说属于典型的“实践出真知”。7. 个人体会与一个值得一试的扩展做完Assignment 1的实验部分我最大的感受是语言模型训练中那些所谓的“常识”其实都是被大工程实践逼出来的经验总结。如果没有亲手从零实现并跑完完整实验我永远不会真正理解为什么学习率要warmup、为什么要用余弦退火、为什么batch size和学习率要协同调整。最后分享一个值得一试的扩展在完成标准baseline之后可以尝试把语言的排序改为“从左到右”与“从右到左”各训练一个模型并在相同的token级别评估集上对比perplexity。这个实验不复杂只需要在数据切分时把每个样本反转即可。由于自然语言本身的方向性特点两个方向的模型表现可能会出现可观测的差异在适当训练步数下多观察几条两者生成的文本基本就能直观体会到“语言先验”在训练数据里到底扮演了多大的角色。这个方向性对比在后续进阶的Masked Language Model预训练作业里也能继续派上用场。