从零构建大语言模型：Transformer架构、训练技巧与实战指南

1. 项目概述从零构建你自己的大语言模型最近几年大语言模型LLM的热度居高不下从ChatGPT到Claude再到国内外的各种开源模型它们展现出的理解和生成能力让人惊叹。但你是否也和我一样在惊叹之余总感觉这些“庞然大物”离自己很远它们像是运行在遥远云端服务器里的黑箱我们只能通过API调用却对其内部构造、训练过程乃至微调细节知之甚少。这种距离感对于想深入理解AI技术本质的开发者或研究者来说无疑是一种遗憾。“datawhalechina/diy-llm”这个项目恰好击中了这个痛点。它不是一个现成的、拿来即用的模型而是一份详尽的“建造手册”。其核心目标非常明确引导你从零开始亲手搭建、训练并理解一个属于你自己的、可运行的大语言模型。它剥离了商业产品的神秘面纱将LLM的构建过程拆解成一个个清晰、可执行的步骤让你不再是API的调用者而是模型的创造者。这个项目适合谁我认为有三类人最应该关注AI/机器学习领域的在校学生或初级从业者教科书上的理论总是抽象的而亲手实现一遍是理解Transformer架构、注意力机制、词嵌入等核心概念最有效的方式。希望将LLM能力深度集成到自身业务中的开发者当你需要针对特定领域如法律、医疗、金融定制模型时仅仅调用通用API往往不够。理解模型底层你才能更好地进行数据准备、模型微调和效果优化。任何对AI技术有强烈好奇心并愿意动手实践的爱好者即使你并非科班出身只要具备一定的编程基础尤其是Python和学习的耐心这个项目也能带你领略AI创造的魅力。项目的价值远不止于得到一个能运行的模型。通过这个过程你将深刻理解数据如何转化为模型的“知识”超参数如何影响模型的行为以及训练过程中的种种“坑”与应对之道。这就像学做菜看一百遍菜谱不如亲手炒一次。接下来我将结合自己的实践带你深入拆解这个“自造LLM”的完整旅程。2. 核心思路与架构选型解析在动手之前我们必须想清楚要造一个什么样的“轮子”。完全复现一个千亿参数的GPT-4既不现实也无必要。DIY-LLM项目的设计思路非常务实在有限的资源通常是单张消费级显卡下实现一个架构完整、流程清晰、且具有一定可玩性的小型LLM。这个定位决定了后续所有的技术选型。2.1 模型规模与架构的权衡首先面临的是模型规模的选择。参数量直接决定了模型的能力上限和训练成本。对于DIY项目目标通常是一个参数量在百万到十亿级别的小模型。例如一个类似GPT-2 Small约1.2亿参数或更小的模型是常见选择。为什么不是更大的计算资源限制训练一个百亿参数模型需要庞大的GPU集群和数周时间个人开发者难以承受。调试与理解模型越小前向传播和反向传播的计算图越容易跟踪便于调试和理解每一层的作用。快速迭代小模型训练周期短允许我们快速尝试不同的架构调整、数据策略和训练技巧积累经验。在架构上Transformer的Decoder-Only结构是当前LLM的事实标准。DIY-LLM项目也必然基于此。我们需要实现的核心模块包括词嵌入层Embedding将离散的文本token转换为连续的向量表示。这里涉及词表Vocabulary的构建是模型理解文本的基础。位置编码Positional Encoding为模型注入序列中token的顺序信息。除了原始Transformer的固定正弦编码更现代的做法是使用可学习的相对位置编码如RoPE, Rotary Positional Encoding这对长文本建模更友好。Transformer Block这是模型的核心。每个Block包含多头自注意力机制Multi-Head Self-Attention让模型能够关注输入序列中所有token之间的关系是理解上下文的关键。前馈神经网络Feed-Forward Network通常是一个两层MLP为每个位置提供非线性变换。层归一化LayerNorm和残差连接Residual Connection这两项技术是训练深层网络稳定的基石必须正确实现。注意在实现注意力机制时要特别注意**因果掩码Causal Mask**的应用。这是生成式模型与BERT等双向模型的核心区别它确保在预测下一个token时模型只能看到当前及之前的token而不能“偷看”未来的信息。这是实现文本自回归生成的关键。2.2 训练目标与数据流水线设计大语言模型本质上是下一个词预测器。其训练目标损失函数是标准的交叉熵损失Cross-Entropy Loss即最大化模型预测的下一个token与真实token的一致性。然而比模型实现更复杂、也更容易出问题的是数据流水线。DIY-LLM项目成功与否一半取决于数据处理。我们需要一个高效的文本加载、分词、批处理和数据增强的管道。原始文本处理网络爬取的文本通常杂乱包含HTML标签、特殊字符、重复内容等。需要经过清洗、标准化如统一全半角、繁简体和分段。分词器Tokenizer选择或训练一个合适的分词器至关重要。对于中文Byte-Pair Encoding (BPE) 或WordPiece是常用方法。你可以使用tiktoken(OpenAI)或sentencepiece(Google)等库。分词器的词表大小是一个重要超参数通常在1万到5万之间需要在模型容量和分词粒度间取得平衡。数据格式与批处理通常我们会将文本处理成固定的长度序列如1024个token。对于不足长度的进行填充Padding过长的进行截断或滑动窗口分割。批处理时需要使用DataLoader并确保同一批次内的序列长度一致动态填充以减少计算浪费。2.3 框架与工具链选择虽然从零实现所有数学运算使用纯NumPy是终极学习方式但为了效率和实用性我们通常会基于深度学习框架。PyTorch是当前研究和DIY项目的首选因其动态图特性非常灵活调试直观。工具链的典型选择如下深度学习框架PyTorch必选。分布式训练可选如果有多张GPU可以考虑使用PyTorch的DistributedDataParallel(DDP) 或 Fully Sharded Data Parallel (FSDP) 来加速。对于单卡DIY可以暂不考虑。实验跟踪使用Weights Biases (WB)或TensorBoard来记录损失曲线、评估指标和模型权重这对于分析训练过程不可或缺。性能分析使用PyTorch Profiler或nvprofNVIDIA来查找训练瓶颈是优化代码、提高GPU利用率的关键。这个选型思路的核心是平衡学习深度与实践可行性。我们不会使用过度封装的高级API如直接调用transformers库的GPT2LMHeadModel而是会亲手搭建关键组件同时利用成熟框架的自动求导和GPU加速功能确保项目能够在一个可接受的时间内跑通并看到结果。3. 从零开始的模型实现关键细节理解了整体蓝图后我们进入具体的建造环节。这里我将分享几个在实现模型核心组件时容易踩坑的关键细节和实操心得。3.1 词嵌入与位置编码的“耦合”问题词嵌入层很简单就是一个nn.Embedding(vocab_size, hidden_dim)。位置编码层也很简单无论是计算正弦余弦向量还是初始化一个可学习的嵌入矩阵。但问题在于如何将两者结合起来输入到第一个Transformer Block常见的做法有两种相加Additioninput token_embedding position_embedding。这是原始Transformer论文的做法。这里有一个细微但重要的点你需要确保position_embedding的维度与token_embedding完全一致并且对于序列中[PAD]token的位置其位置编码也应该是零向量或一个特定的掩码值避免引入噪声。拼接后投影Concatenation Projection将token embedding和position embedding在特征维度拼接然后通过一个线性层投影到hidden_dim。这种方式给了模型更多灵活性去学习如何融合两种信息但增加了少量参数。我个人的经验是对于中小模型直接相加简单有效且是大多数开源模型采用的方式。在实现时务必写一个简单的测试用例验证对于不同长度的输入序列你的位置编码生成函数是否能正确工作。import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) self.register_buffer(pe, pe.unsqueeze(0)) # (1, max_len, d_model) def forward(self, x): # x: (batch_size, seq_len, d_model) seq_len x.size(1) x x self.pe[:, :seq_len, :] return x3.2 高效且正确的注意力机制实现自注意力机制的计算公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。在PyTorch里几行代码就能实现。但DIY时有三个性能与正确性陷阱需要避开因果掩码的实现掩码需要在softmax之前应用将未来位置的注意力分数设置为一个极大的负数如-1e9这样softmax后其权重就接近0。# attn_scores: (batch, num_heads, seq_len, seq_len) if causal_mask: mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() attn_scores attn_scores.masked_fill(mask, -1e9) attn_weights F.softmax(attn_scores, dim-1)缩放因子的重要性公式中的sqrt(d_k)至关重要。如果没有它当d_k较大时点积QK^T的值会非常大导致softmax梯度极小梯度消失模型将无法学习。Flash Attention的考量标准的注意力实现需要计算并存储一个(seq_len, seq_len)的中间矩阵这在序列很长时内存消耗巨大O(N²)。对于想尝试更长上下文的DIYer可以考虑集成Flash Attention。这是一个高度优化的注意力计算内核能显著降低内存占用并提升速度。不过它的实现较为复杂初期可以先用标准实现后续再作为优化点引入。3.3 层归一化与残差连接的放置顺序Transformer Block中LayerNorm和残差连接的顺序有两种主流变体“Pre-LN”和“Post-LN”或称“原始Transformer”。Post-LN原始输出 LayerNorm(x Sublayer(x))。即先做残差连接再进行层归一化。这种结构在训练初期可能不太稳定。Pre-LN现代主流输出 x Sublayer(LayerNorm(x))。即先对输入进行层归一化再经过子层注意力或FFN最后做残差连接。这种结构被证明能使训练更稳定、收敛更快。目前绝大多数新模型如GPT-3、LLaMA都采用Pre-LN结构。在DIY时我强烈建议你也使用Pre-LN它能减少很多训练调试的麻烦。你的Transformer Block结构看起来会是这样输入x | v LayerNorm1 | v MultiHeadAttention | v Dropout (可选) | v x (上述结果) // 第一次残差连接 | v LayerNorm2 | v FeedForward | v Dropout (可选) | v (上一步结果) (上述结果) // 第二次残差连接 | v 输出3.4 损失函数与标签对齐的“错位”问题这是一个新手极易出错的地方。我们训练的是自回归语言模型给定序列[token_1, token_2, ..., token_n]模型的任务是预测下一个token。因此对于输入inputs我们的标签labels应该是inputs向右移动一位。输入模型看到的[token_1, token_2, token_3, ..., token_{n-1}]标签我们要预测的[token_2, token_3, token_4, ..., token_n]在计算损失时我们取模型最后一层输出的所有位置的logits但只计算对应标签位置上的交叉熵损失。通常我们会忽略对[PAD]token的损失计算通过一个ignore_index参数。在PyTorch中可以这样操作criterion nn.CrossEntropyLoss(ignore_indexpad_token_id) # logits shape: (batch_size, seq_len, vocab_size) # labels shape: (batch_size, seq_len) loss criterion(logits.view(-1, vocab_size), labels.view(-1))确保你的数据加载器正确生成了这种错位对齐的inputs和labels这是模型能够学会生成文本的前提。4. 训练流程、技巧与超参数调优模型搭建完毕只是万里长征第一步。训练一个LLM更像一门艺术需要精心调配各种“原料”超参数并控制“火候”训练动态。4.1 优化器与学习率调度策略对于LLM训练AdamW优化器是绝对的主流选择。它相比原始Adam引入了权重衰减的正则化能更好地防止过拟合。关键参数是学习率lr和权重衰减weight_decay。学习率Learning Rate这是最重要的超参数之一。对于亿级参数模型初始学习率通常在1e-4到5e-4之间。一个必须使用的技巧是学习率预热Warmup。在训练开始时模型参数是随机初始化的直接使用较大的学习率可能导致训练不稳定。Warmup策略会在前几千个step内将学习率从0线性或余弦增加到预设的峰值。学习率调度Scheduler在Warmup之后通常采用**余弦退火Cosine Annealing**来逐渐降低学习率。这模拟了模拟退火过程有助于模型在训练后期收敛到更优的局部最优点。一个典型的组合是AdamWLinear WarmupCosine Annealing。你可以使用PyTorch的optim.lr_scheduler模块或更灵活的transformers库中的get_linear_schedule_with_warmup和get_cosine_schedule_with_warmup函数即使你不用其模型也可以单独用调度器。4.2 批次大小与梯度累积在GPU内存有限的情况下我们可能无法放下理想的大批次Batch Size。这时梯度累积Gradient Accumulation是救命稻草。其原理是在多个小批次micro-batch上前向传播并计算梯度但不立即更新参数optimizer.step()而是将梯度累加。在累积了足够多的小批次accumulation steps后再用累积的总梯度进行一次参数更新。例如你希望有效批次大小为64但GPU只能放下批次大小为16的数据。你可以设置accumulation_steps4。代码逻辑如下optimizer.zero_grad() for step in range(accumulation_steps): inputs, labels get_micro_batch() loss model(inputs, labels) loss loss / accumulation_steps # 对损失进行缩放使梯度数值稳定 loss.backward() # 梯度被累加到.grad属性中 optimizer.step() # 用累积的梯度更新一次参数 scheduler.step() # 更新学习率 optimizer.zero_grad() # 清空梯度准备下一轮累积这样做既模拟了大批次训练的效果梯度估计更准确又突破了单卡内存的限制。注意在计算损失时进行了缩放这是为了保持梯度累加前后参数更新的总幅度一致。4.3 模型评估与保存策略训练不能只看训练损失Training Loss不断下降还必须关注验证损失Validation Loss或困惑度Perplexity, PPL。困惑度是衡量语言模型好坏的核心指标其值越低越好。计算公式为PPL exp(loss)。你需要定期如每1000个训练step在预留的验证集上跑一次评估计算验证损失和困惑度。当验证损失开始上升而训练损失继续下降时就意味着过拟合Overfitting发生了。模型保存策略至关重要最佳模型保存始终保存在验证集上困惑度最低的模型 checkpoint。定期保存每隔一定步数或epoch保存一次以防训练中途中断。保存完整状态保存的checkpoint不应只包含模型参数model.state_dict()还应包括优化器状态optimizer.state_dict()、当前步数、学习率调度器状态等。这样可以从中断点精确恢复训练。使用WB或TensorBoard将这些损失曲线、学习率变化、模型权重分布直方图可视化是诊断训练问题的强大工具。4.4 一些实用的训练技巧Tricks梯度裁剪Gradient Clipping在调用optimizer.step()之前使用torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)来裁剪梯度范数。这可以防止梯度爆炸稳定训练过程尤其是在训练初期。max_norm通常设为1.0或0.5。混合精度训练Mixed Precision Training使用torch.cuda.amp进行自动混合精度训练。这能显著减少GPU内存占用因为部分计算使用FP16并可能加快训练速度。对于DIY项目这几乎是必选项能让你在同等显存下使用更大的批次或模型。scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss model(inputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()激活检查点Activation Checkpointing这是一种用计算换内存的技术。它在前向传播时不保存某些中间激活值在反向传播需要时重新计算它们。对于层数很深的模型这能大幅降低内存消耗。PyTorch中可以通过torch.utils.checkpoint.checkpoint函数实现。5. 文本生成、评估与常见问题排查模型训练完成后最激动人心的时刻就是让它“开口说话”。文本生成是推理阶段的核心同样有很多门道。5.1 自回归文本生成算法最基本的生成方式是贪心搜索Greedy Search每一步都选择概率最高的token作为下一个输出。这种方法简单高效但容易导致重复、乏味的文本。更常用的是核采样Top-p Sampling 或称Nucleus Sampling和Top-k Sampling。Top-k Sampling每一步只从概率最高的k个token中采样。这避免了选择那些概率极低的奇怪token。Top-p Sampling每一步从累积概率超过p的最小token集合中采样。例如设p0.9我们将所有token按概率从高到低排序然后依次累加概率直到总和超过0.9然后只从这个集合中采样。这种方法能动态调整候选集的大小通常比固定的Top-k更灵活。在实际应用中结合Top-p和温度Temperature是最佳实践。温度参数T用于调整softmax分布的平滑程度。T1使用原始分布T1会使分布更尖锐模型更自信输出更确定T1会使分布更平滑模型更不确定输出更多样、更有创造性。def generate_text(model, prompt, max_len50, temperature0.8, top_p0.9): model.eval() tokens tokenizer.encode(prompt) for _ in range(max_len): with torch.no_grad(): logits model(torch.tensor([tokens])) # 取最后一个位置的logits next_token_logits logits[0, -1, :] / temperature # Top-p采样 sorted_logits, sorted_indices torch.sort(next_token_logits, descendingTrue) cumulative_probs torch.cumsum(F.softmax(sorted_logits, dim-1), dim-1) # 移除累积概率大于top_p的部分 sorted_indices_to_remove cumulative_probs top_p # 保留第一个超过阈值的token所以将其设为False sorted_indices_to_remove[1:] sorted_indices_to_remove[:-1].clone() sorted_indices_to_remove[0] False indices_to_remove sorted_indices[sorted_indices_to_remove] next_token_logits[indices_to_remove] -float(Inf) # 从剩余token中采样 probs F.softmax(next_token_logits, dim-1) next_token torch.multinomial(probs, num_samples1).item() tokens.append(next_token) if next_token eos_token_id: # 遇到结束符则停止 break return tokenizer.decode(tokens)5.2 如何评估你的DIY模型训练损失和验证困惑度是内部指标。我们还需要一些外部指标来评估生成文本的质量。对于DIY的小模型不要指望它能达到ChatGPT的水平。合理的评估方向包括连贯性Coherence生成的句子在语法和基本逻辑上是否通顺可以人工阅读判断。相关性Relevance生成的文本是否与给定的提示Prompt相关多样性Diversity多次生成的结果是否丰富多样而不是千篇一律特定任务测试如果你用特定领域数据微调过可以设计一些领域内的问答或补全任务进行测试。一个实用的方法是构建一个小型测试集包含各种类型的提示如开放式问题、文章续写、对话开头等然后让模型生成并邀请几个人进行主观评分例如1-5分。虽然主观但对于小规模项目非常有效。5.3 训练与生成过程中的典型问题与排查在DIY过程中你几乎一定会遇到下面这些问题。这里是我的排查清单问题1训练损失Loss不下降或者下降非常缓慢。检查数据首先确认你的数据流水线是正确的。打印几个batch的inputs和labels看它们是否错位对齐标签里是否包含大量[PAD]数据本身质量是否太差乱码、无意义文本检查模型初始化糟糕的参数初始化可能导致梯度消失或爆炸。尝试使用标准的初始化方法如nn.init.xavier_uniform_或nn.init.normal_(weight, mean0.0, std0.02)后者是GPT系列常用的。检查学习率学习率可能太小了。尝试增大一个数量级如从1e-4到1e-3并确保Warmup步骤设置正确。检查梯度在训练初期打印出模型各层的梯度范数。如果梯度普遍接近0可能是初始化或激活函数如ReLU导致“神经元死亡”。如果梯度非常大可能需要梯度裁剪。问题2模型输出全是乱码、重复词或无意义的常见token如“the”, “.”。这是“模式坍塌Mode Collapse”的典型症状。模型找到了一个能轻易降低训练损失比如总是预测最常见的token的简单模式。降低学习率过高的学习率可能导致模型无法学习复杂模式直接坍缩到简单解。调整采样参数在生成时如果温度T设得太低如0.1会导致确定性太强容易重复。尝试提高温度如0.7-1.0或使用Top-p采样。检查损失函数确认你在计算损失时正确忽略了[PAD]token。如果没有忽略模型可能会学会专注于预测Padding而不是真正的文本。问题3训练后期验证损失开始上升过拟合。增加正则化可以尝试增大weight_decayAdamW优化器的参数或者在FFN层后增加Dropout。获取更多数据对于语言模型数据量几乎永远不嫌多。如果可能收集更多样、更高质量的文本数据。早停Early Stopping这是最直接的方法。一旦验证损失在连续多个评估点不再下降甚至上升就停止训练并回滚到验证损失最低的checkpoint。问题4生成速度非常慢。确认处于推理模式在生成前调用model.eval()这会禁用Dropout等训练专用层。使用缓存Key-Value Cache这是推理加速的关键技术。在自回归生成中对于已经生成的token其Key和Value向量在计算后续token的注意力时是可以重复使用的。实现KV Cache可以避免大量重复计算将生成复杂度从O(n²)降低到O(n)。虽然实现稍复杂但对于提高交互体验至关重要。批量生成如果硬件允许可以一次处理多个生成请求批量推理能更好地利用GPU的并行能力。6. 项目扩展与未来方向思考完成一个基础版本的DIY-LLM后你获得的不仅仅是一个模型更是一套完整的LLM构建方法论。以此为起点你可以向多个方向深入探索让这个项目持续产生价值。6.1 方向一模型架构改进你的第一个模型很可能是一个标准的Transformer Decoder。你可以尝试集成最新的研究成果来提升其能力或效率替换注意力机制尝试MQAMulti-Query Attention或GQAGrouped-Query Attention。它们在保证效果相近的前提下显著减少了推理时的KV缓存大小从而提升了生成速度并降低了内存占用。这对于在资源受限环境下部署模型非常有用。尝试不同的归一化层除了LayerNorm可以试试RMSNormRoot Mean Square Normalization它被用于LLaMA等模型计算更简单且据说训练更稳定。引入滑动窗口注意力如果想让模型处理更长的文本可以引入类似Longformer或FlashAttention中使用的滑动窗口注意力将计算复杂度从序列长度的平方级降低到线性级。6.2 方向二指令微调与对齐基础语言模型只是“续写大师”要让它成为“对话助手”或“任务执行者”需要进行指令微调Instruction Tuning和人类反馈强化学习RLHF。指令微调收集或生成大量的(指令, 期望输出)配对数据。例如“将这句话翻译成英文今天天气真好。” - “The weather is really nice today.” 然后用这些数据在你的基础模型上进行有监督微调。这能教会模型理解并遵循人类指令。对齐技术RLHF是让模型输出更符合人类偏好的关键技术但其实现非常复杂。一个更易入手的替代方案是直接偏好优化DPO。它不需要训练一个复杂的奖励模型只需要一个偏好数据集即对于同一个提示有两个模型输出标注哪个更好就能直接优化模型使其输出更受偏好。这是当前个人研究者进行模型对齐的热门选择。6.3 方向三模型压缩与高效部署训练出的模型如何应用到实际场景模型压缩和高效部署是必经之路。量化Quantization将模型权重从FP32转换为更低精度如INT8、INT4可以大幅减少模型体积和内存占用并提升推理速度。PyTorch提供了torch.quantization模块社区也有GPTQ、AWQ等针对LLM的后训练量化算法。模型剪枝Pruning移除模型中不重要的权重例如那些接近0的权重得到一个更稀疏、更小的模型。使用推理引擎将PyTorch模型导出为ONNX格式然后使用专门的推理引擎如TensorRT,ONNX Runtime进行部署可以获得比原生PyTorch更好的推理性能。6.4 方向四构建垂直领域专家这是最具实用价值的方向。利用你的DIY框架在某个垂直领域如医学文献、法律条文、编程代码、小说创作收集高质量数据从头预训练或从你的基础模型进行领域自适应继续预训练。你将得到一个在该领域表现远超通用模型的“专家”。这个过程会让你深入理解数据质量、领域词汇、评估指标等实际问题。整个DIY-LLM项目就像学习木工。一开始你照着图纸做一把简单的椅子可能会歪歪扭扭。但在这个过程中你熟悉了锯子、刨子、凿子的用法理解了榫卯结构的原理。之后你就能自己设计并制作出更复杂、更精美的家具。从理解Transformer的每一行公式开始到最终让模型流畅地生成文本、回答问题这种亲手创造的成就感和深刻理解是任何API调用都无法给予的。这个项目最大的收获不是终点那个模型文件而是你脑中构建起来的那套完整、清晰、可扩展的LLM知识体系。它让你在面对未来更复杂的AI模型时不再是一个被动的使用者而是一个自信的探索者和创造者。