从零构建Transformer语言模型：数据处理、训练与推理全流程实战

语言模型这个方向我从早期的n-gram一路跟到RNN、LSTM再到现在的Transformer架构踩过的坑不算少。今天想聊的是怎么从零用Transformer搭一个能跑通、能训练、能出效果的语言模型。这不是那种调个库三行代码搞定的教程而是把整个链路拆开讲清楚每一步为什么这么做、哪里容易翻车、实际训练时有哪些参数需要盯着调。如果你已经了解基本的深度学习概念知道什么是张量、梯度、损失函数但还没亲手搭过一个完整的Transformer语言模型那这篇内容就是写给你的。如果你已经搭过但训练效果不理想也可以看看我在数据处理、学习率调度、梯度裁剪这些环节的经验说不定能找到问题所在。整篇内容会围绕一个核心目标展开用Transformer架构构建一个能进行下一词预测的语言模型从数据准备、模型定义、训练循环到推理生成每一步都给出可复现的操作和背后的逻辑分析。1. 为什么语言模型的第一道坎永远是数据处理很多人一上来就急着定义模型结构结果训练半天loss不降回头查才发现是数据管道出了问题。我见过太多这种情况了包括我自己早期也犯过。语言模型的数据处理不像图像分类那样直观它涉及到分词、词表构建、序列截断、批次组装等多个环节每个环节都有坑。1.1 分词策略的选择字符级、词级还是子词级分词是语言模型的第一步也是影响后续所有环节的关键决策。常见的方案有三种字符级分词把每个字符当作一个token。优点是词表极小中文常用字几千个英文就百来个不存在OOV未登录词问题。缺点是序列变得很长模型需要从字符层面学到语义组合训练成本高。词级分词按空格或分词工具切分成词。优点是序列短、语义单元清晰。缺点是词表爆炸动辄几十万OOV问题严重遇到新词就抓瞎。子词级分词介于两者之间用BPEByte Pair Encoding、WordPiece或Unigram等算法把常见词保留为整体罕见词拆成子词。这是目前主流方案GPT系列、BERT系列都用的这种。我实测下来的感受是如果你要做中文语言模型子词级分词几乎是必选项。中文不像英文有天然空格分隔词级分词依赖分词工具的质量而分词工具本身就有误差。用BPE直接在字符层面学习合并规则能自动发现高频字符组合效果稳定很多。具体操作上可以用sentencepiece或tokenizers库来训练BPE模型。以sentencepiece为例import sentencepiece as spm spm.SentencePieceTrainer.train( inputcorpus.txt, model_prefixtokenizer, vocab_size32000, model_typebpe, character_coverage0.9995, max_sentence_length8192, num_threads8 )这里有几个参数值得注意。vocab_size设多少太小会导致序列过长太大则embedding矩阵参数量暴涨。32000是一个比较通用的选择GPT-2用的50257BERT中文用的21128。我一般根据语料规模来定语料小于1GB用16000-240001-10GB用3200010GB以上可以考虑50000。character_coverage设为0.9995意味着覆盖99.95%的字符剩下的当作UNK处理这个值不要设成1.0否则遇到生僻字会报错。1.2 序列长度与批次组装的工程细节语言模型的训练数据是一条条不等长的文本。要把它们组装成固定形状的批次有两种常见做法做法一固定长度截断填充。设定一个最大长度比如512超过的截断不足的用PAD token补齐。这种做法简单直接但浪费计算资源——如果一批数据里大部分序列只有100个token你却在512的长度上做注意力计算大量算力花在了PAD位置上。做法二动态长度分批。把长度相近的序列分到同一个批次里每个批次的最大长度根据实际数据动态确定。这种做法能显著提升训练效率但实现起来稍复杂。我通常用第二种方案的简化版先按长度排序然后按批次大小切分每个批次内padding到该批次的最大长度。这样既避免了全局最大长度的浪费又不需要复杂的动态调度。def collate_batch(batch, pad_id0): # batch是(token_ids, label_ids)的列表 max_len max(len(item[0]) for item in batch) input_ids [] labels [] attention_mask [] for input_seq, label_seq in batch: pad_len max_len - len(input_seq) input_ids.append(input_seq [pad_id] * pad_len) labels.append(label_seq [-100] * pad_len) # -100在CrossEntropyLoss中被忽略 attention_mask.append([1] * len(input_seq) [0] * pad_len) return torch.tensor(input_ids), torch.tensor(labels), torch.tensor(attention_mask)注意labels的padding用-100而不是pad_id因为PyTorch的CrossEntropyLoss默认ignore_index-100。如果用pad_id作为忽略索引需要显式设置ignore_index参数。这个细节看起来小但搞错了会导致模型在PAD位置也计算损失影响收敛。1.3 数据管道中容易被忽视的坑说几个我在实际项目中踩过的坑第一个坑词表大小和embedding层不匹配。训练完tokenizer后词表大小是32000但模型定义时embedding层写成了30000训练时直接报index out of range。这个错误很蠢但很常见尤其是tokenizer和模型分开写的时候。我的做法是在模型初始化时动态读取tokenizer的vocab_size而不是硬编码。第二个坑特殊token的ID冲突。BPE训练出来的词表中PAD、UNK、BOS、EOS这些特殊token的ID需要明确指定。如果不指定sentencepiece会自动分配可能和你代码里假设的不一致。建议在训练tokenizer时就通过pad_id、unk_id、bos_id、eos_id参数固定下来。第三个坑数据泄露。构建训练集和验证集时如果按行随机划分同一个文档的相邻段落可能同时出现在训练集和验证集中导致验证loss虚低。正确做法是按文档级别划分确保同一文档的所有内容只出现在一个集合中。2. Transformer语言模型的骨架拆解数据处理搞定之后接下来是模型本身。Transformer的原始论文给出了完整架构但直接照搬不一定适合语言模型任务。这里我按语言模型的需求把关键模块拆开讲。2.1 从原始Transformer到Decoder-only架构原始Transformer是为机器翻译设计的包含Encoder和Decoder两部分。但语言模型任务下一词预测只需要Decoder部分因为它的本质是根据前面的token预测下一个token不需要编码器来理解源语言。Decoder-only架构的核心特点是因果掩码Causal Mask。在计算自注意力时每个位置只能看到自己和之前的位置不能看到后面的token。这是语言模型和BERT这类双向模型最本质的区别。def generate_causal_mask(seq_len, device): # 下三角矩阵1表示可见0表示不可见 mask torch.tril(torch.ones(seq_len, seq_len, devicedevice)) # 扩展维度以适配注意力分数的形状(batch, heads, seq_len, seq_len) mask mask.view(1, 1, seq_len, seq_len) return mask这个掩码的作用用生活化的类比来说就像考试时你只能看已经做过的题目不能翻到后面看还没做的题。模型在预测第5个词时只能利用第1到第4个词的信息。2.2 多头注意力的参数量与计算量估算多头自注意力是Transformer的核心组件。假设模型维度为d_model头数为n_heads每个头的维度就是d_model / n_heads。以d_model512、n_heads8为例每个头的维度是64。单个注意力头的计算过程将输入投影到Q、K、V三个空间每个投影矩阵形状为(d_model, d_k)计算注意力分数Q K^T / sqrt(d_k)应用softmax和因果掩码加权求和attention_weights V参数量方面Q、K、V的投影矩阵各有d_model * d_k个参数输出投影矩阵有d_model * d_model个参数。以d_model512、n_heads8为例每个头的QKV参数量3 * 512 * 64 983048个头的QKV总参数量8 * 98304 786432输出投影参数量512 * 512 262144多头注意力总参数量约104万这只是一个注意力层的参数量。如果堆叠12层光注意力部分就有约1250万参数。再加上前馈网络通常维度是d_model的4倍即2048每层前馈网络参数量为512 * 2048 * 2 209715212层就是约2500万。整个模型参数量在4000万左右属于中小规模。实操建议如果你显存有限优先减小d_model和层数而不是减小头数。头数太少比如小于4会导致每个头的维度太小表达能力受限。一般保持每个头维度在64-128之间比较合适。2.3 位置编码正弦编码还是可学习编码Transformer本身没有循环结构也没有卷积结构如果不加位置编码它无法区分猫追狗和狗追猫。位置编码就是给每个位置一个独特的向量表示。两种主流方案正弦位置编码用不同频率的正弦和余弦函数生成固定编码不需要训练。优点是理论上可以外推到训练时没见过的长度。缺点是在实际使用中外推效果往往不理想。可学习位置编码为每个位置初始化一个可训练的向量随模型一起训练。优点是灵活能适应具体任务。缺点是最大长度在训练时就固定了无法处理更长的序列。我个人的选择是如果序列长度固定且不超过训练长度用可学习编码如果需要处理变长序列或可能超出训练长度用正弦编码或者RoPE旋转位置编码。RoPE是近几年的主流方案它通过旋转矩阵将位置信息注入到Q和K中外推能力比正弦编码好很多。class LearnedPositionalEncoding(nn.Module): def __init__(self, d_model, max_len512): super().__init__() self.embedding nn.Embedding(max_len, d_model) def forward(self, x): seq_len x.size(1) positions torch.arange(seq_len, devicex.device).unsqueeze(0) return x self.embedding(positions)2.4 前馈网络与激活函数的选择Transformer中的前馈网络是一个两层全连接网络中间夹一个激活函数。原始论文用的是ReLU后来GPT系列改用了GELU现在很多模型用SwiGLU。这三种激活函数的对比激活函数计算复杂度效果使用场景ReLU低一般早期模型GELU中较好GPT-2/3、BERTSwiGLU高最好LLaMA、PaLMSwiGLU的效果最好但参数量会增加——它需要三个投影矩阵而不是两个。如果你的计算资源有限GELU是性价比最高的选择。前馈网络的中间维度通常是d_model的4倍。但这个比例不是固定的LLaMA用的是8/3倍然后取整到最近的64的倍数。我一般用4倍简单省事效果也不差。3. 训练循环里的那些关键决策模型定义好了接下来是训练。训练循环看起来简单——前向传播、计算损失、反向传播、更新参数——但里面的细节决定了模型能不能收敛、收敛得多快、最终效果好不好。3.1 学习率调度预热与余弦退火学习率是训练中最重要的超参数没有之一。太大导致loss震荡甚至发散太小导致收敛极慢。Transformer类模型的标准做法是预热衰减前若干步线性增加学习率预热然后按余弦函数或线性函数衰减。def get_lr(step, d_model, warmup_steps, max_steps): # 预热阶段线性增加 if step warmup_steps: return d_model ** (-0.5) * step / warmup_steps # 衰减阶段余弦退火 progress (step - warmup_steps) / (max_steps - warmup_steps) return d_model ** (-0.5) * 0.5 * (1 math.cos(math.pi * progress))为什么需要预热因为训练初期模型参数是随机初始化的梯度方向不稳定。如果一开始就用大学习率参数更新幅度过大可能把模型推到一个很难恢复的坏区域。预热让模型先小步走等梯度方向稳定了再加速。预热步数设多少原始论文用的是4000步但这个和批次大小、数据规模有关。我的经验是预热步数约为总训练步数的5%-10%。比如你计划训练100000步预热5000-10000步比较合适。3.2 梯度裁剪防止训练爆炸的安全阀Transformer模型在训练过程中容易出现梯度爆炸尤其是深层模型。梯度裁剪就是把梯度的范数限制在一个阈值内超过就按比例缩放。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)max_norm设多少1.0是常用值GPT-2用的是1.0很多开源实现也用1.0。如果训练稳定可以适当放宽到2.0或5.0。如果经常爆炸降到0.5试试。注意梯度裁剪要在反向传播之后、优化器更新之前调用。顺序搞反了等于没裁剪。3.3 批次大小与梯度累积的权衡批次大小影响训练稳定性和效率。大批次训练更稳定梯度估计更准确但显存占用大。小批次显存友好但梯度噪声大训练可能不稳定。当显存不足以支撑理想批次大小时可以用梯度累积多次前向反向传播累积梯度然后一次性更新参数。效果等价于大批次训练。accumulation_steps 4 optimizer.zero_grad() for i, batch in enumerate(dataloader): outputs model(batch) loss criterion(outputs, batch.labels) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad()注意loss要除以accumulation_steps否则梯度会累积成原来的N倍。3.4 混合精度训练省显存提速的实用技巧混合精度训练用float16或bfloat16做前向和反向传播用float32维护一份参数副本。这样能省30%-50%显存速度也能提升。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for batch in dataloader: with autocast(dtypetorch.bfloat16): outputs model(batch) loss criterion(outputs, batch.labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad()用float16时需要GradScaler来防止梯度下溢用bfloat16则不需要因为bfloat16的指数范围和float32一样。如果你的显卡支持bfloat16Ampere架构及以上优先用bfloat16省心。我实测下来混合精度训练在保持模型效果基本不变的前提下训练速度能提升约40%显存占用减少约35%。这个收益非常可观。4. 推理生成从贪心搜索到核采样模型训练完之后怎么用它生成文本这里面的策略选择直接影响生成质量。4.1 贪心搜索和束搜索的局限性贪心搜索每一步都选概率最高的token。实现简单但容易生成重复、单调的文本。因为一旦模型进入某个循环模式它就会一直重复下去。束搜索维护k个候选序列每步扩展所有候选并保留概率最高的k个。能缓解贪心搜索的局部最优问题但计算成本是贪心的k倍。而且在开放式文本生成中束搜索倾向于生成安全但无聊的文本因为它在优化整体概率而高概率的文本往往是那些最常见的、缺乏创意的表达。4.2 温度、Top-k与核采样的组合策略实际使用中更常用的是随机采样策略温度缩放在softmax之前将logits除以温度参数T。T1使分布更尖锐更确定T1使分布更平坦更随机。T1就是原始分布。logits logits / temperature probs torch.softmax(logits, dim-1)Top-k采样只从概率最高的k个token中采样其余置零。k1等价于贪心搜索kvocab_size等价于纯随机采样。常用k值在20-100之间。核采样Top-p从累积概率达到p的最小token集合中采样。p通常设为0.9或0.95。相比Top-k核采样能动态调整候选集大小——当分布集中时候选集小分布分散时候选集大。def top_p_sampling(logits, p0.9, temperature1.0): logits logits / temperature sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) # 找到累积概率超过p的位置 sorted_indices_to_remove cumulative_probs p # 保留第一个超过p的token sorted_indices_to_remove[..., 1:] sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] 0 # 将需要移除的token概率置为负无穷 indices_to_remove sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] float(-inf) probs torch.softmax(logits, dim-1) return torch.multinomial(probs, num_samples1)我通常的组合是温度0.8Top-p 0.92同时限制Top-k不超过50。这个组合在生成多样性和连贯性之间取得了不错的平衡。4.3 重复惩罚与生成质量控制即使有了采样策略模型有时还是会陷入重复。这时候可以加重复惩罚对已经出现过的token降低其logits。def apply_repetition_penalty(logits, generated_ids, penalty1.2): for token_id in set(generated_ids): if logits[token_id] 0: logits[token_id] / penalty else: logits[token_id] * penalty return logitspenalty设为1.0表示不惩罚1.2表示中等惩罚1.5以上可能影响生成质量。我一般用1.1-1.3之间。还有一个技巧是no-repeat n-gram禁止生成与前面重复的n-gram。比如设置no_repeat_ngram_size3模型就不会生成任何重复的3-gram。这个对抑制逐字重复很有效但设太大比如5以上可能导致生成不自然的文本。5. 训练不收敛时的排查思路即使按照上面的步骤操作训练过程中还是可能遇到问题。这里分享几个我实际遇到过的场景和排查方法。5.1 Loss震荡或发散的第一反应看到loss震荡先检查三件事第一学习率是不是太大了。把学习率降低10倍试试。如果loss立刻稳定下来那就是学习率的问题。我遇到过用3e-4训练一个12层模型loss一直在2.5到4.0之间跳降到1e-4之后平稳下降到2.0以下。第二梯度裁剪有没有生效。打印梯度范数看看如果经常超过阈值说明梯度确实在爆炸。可以尝试降低max_norm或者检查模型初始化。第三数据里有没有异常样本。比如空序列、超长序列、全是PAD的序列。这些异常样本会产生异常的loss导致训练不稳定。写个脚本统计一下序列长度分布和token频率排除异常。5.2 Loss正常但生成质量差的可能原因有时候loss降得很好但生成出来的文本一塌糊涂。可能的原因暴露偏差Exposure Bias训练时模型看到的是真实的前缀推理时看到的是自己生成的前缀。一旦生成错一个token后续就会雪崩。缓解方法是 scheduled sampling——训练时以一定概率用模型自己的生成结果作为下一步输入。评估指标和实际质量脱节Perplexity低不代表生成质量好。Perplexity衡量的是模型对真实数据的拟合程度但生成质量还涉及多样性、连贯性、信息量等多个维度。建议在训练过程中定期做人工评估生成一些样本看看实际效果。过拟合如果训练loss持续下降但验证loss开始上升说明模型在记忆训练数据。这时候需要增加dropout、减小模型规模、或者增加训练数据。5.3 显存不够时的取舍策略显存不够是训练语言模型最常见的障碍。按优先级排序可以这样取舍减小批次大小最直接但可能影响训练稳定性。配合梯度累积使用。启用混合精度省30%-50%显存几乎无副作用。使用梯度检查点用计算时间换显存能省50%-70%的激活值显存。代价是训练速度降低约30%。减小模型规模减小d_model或层数。这是最后的选择因为会影响模型能力上限。使用更高效的注意力实现比如FlashAttention能省显存且提速。# 梯度检查点的使用 from torch.utils.checkpoint import checkpoint class TransformerBlock(nn.Module): def forward(self, x, mask): # 用checkpoint包裹前向传播 return checkpoint(self._forward, x, mask) def _forward(self, x, mask): # 实际的注意力前馈计算 ...提示梯度检查点会重新计算前向传播所以dropout等随机操作需要特别注意。如果模型里有随机性重新计算时结果可能不一致。建议在checkpoint内部关闭dropout或者使用确定性的随机种子。6. 从训练到部署的衔接模型训练好之后最终是要用起来的。这里简单说一下从训练到推理的衔接环节。6.1 模型保存与加载的注意事项保存模型时不仅要保存参数还要保存配置信息——层数、d_model、头数、词表大小等。否则加载时不知道用什么结构来实例化模型。# 保存 checkpoint { model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), config: { d_model: 512, n_heads: 8, n_layers: 12, vocab_size: tokenizer.vocab_size(), max_len: 512 }, step: global_step } torch.save(checkpoint, model_checkpoint.pt) # 加载 checkpoint torch.load(model_checkpoint.pt) config checkpoint[config] model TransformerLM(**config) model.load_state_dict(checkpoint[model_state_dict])6.2 推理时的批处理与KV缓存推理时如果每次只生成一个token就重新计算整个序列的注意力效率极低。KV缓存是标准优化缓存之前所有层的Key和Value生成新token时只需要计算当前token的Q然后和缓存的K、V做注意力。class CachedAttention(nn.Module): def forward(self, x, cacheNone): q self.W_q(x) k self.W_k(x) v self.W_v(x) if cache is not None: k torch.cat([cache[k], k], dim1) v torch.cat([cache[v], v], dim1) new_cache {k: k, v: v} # 注意力计算... return output, new_cacheKV缓存能把推理速度提升几倍到几十倍具体取决于序列长度。序列越长收益越大。6.3 量化与模型压缩的初步思路如果要在资源受限的环境部署可以考虑量化——把float32参数转成int8或int4。PyTorch提供了动态量化和静态量化两种方案。动态量化对Transformer类模型效果不错精度损失通常在1%以内模型大小减少约75%。# 动态量化 quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear}, dtypetorch.qint8 )不过量化对生成质量的影响需要实际评估。有些模型量化后困惑度上升明显有些则几乎无感。建议量化后在验证集上跑一遍对比量化前后的perplexity和生成样本。整个流程走下来从数据处理到模型定义、训练、推理、部署每个环节都有需要仔细打磨的细节。我自己的体会是语言模型训练是一个细节决定成败的事情——模型结构大同小异但数据处理的质量、超参数的设置、训练过程的监控这些才是拉开差距的地方。建议在开始大规模训练之前先用小规模数据和小模型跑通整个流程确认每个环节都正确再逐步放大。这样能避免浪费大量时间和算力在调试上。