1. 项目概述从零构建百万参数语言模型的野心与实践最近在GitHub上看到一个挺有意思的项目叫“create-million-parameter-llm-from-scratch”。光看这个标题就足以让任何一个对深度学习感兴趣的人心头一热。它直指当前人工智能领域最炙手可热的核心——大语言模型却加了一个极具挑战性的前缀“从零开始”。这不像是在调用某个现成的API或者微调一个预训练好的模型而是意味着你要亲手从最基础的矩阵乘法开始搭建起一个拥有百万级参数的复杂神经网络并让它真正学会理解和生成人类语言。这个项目的价值远不止于得到一个可运行的模型。对于学习者而言它是一张深入神经网络和Transformer架构腹地的绝佳地图。市面上关于LLM的教程很多但大多停留在概念讲解或调用层面。真正能带你走过数据预处理、模型架构设计、训练循环编写、损失函数调试这一整套“炼丹”流程的少之又少。而这个项目恰恰承诺了这样一次完整的、手把手的旅程。它解决的不仅是“怎么做”的问题更是“为什么这么做”以及“做的时候会遇到什么坑”的问题。无论你是希望夯实基础的学生还是想跨界进入AI领域的开发者亦或是好奇黑箱内部运作的研究者这个项目都提供了一个绝佳的动手实验平台。从技术角度看“百万参数”是一个精妙的定位。它既足够复杂能够体现现代LLM的核心架构如Transformer的注意力机制、层归一化、前馈网络避免了过于玩具化的示例同时又相对轻量使得在个人电脑甚至配置较好的笔记本上进行训练和实验成为可能。你不需要动辄数十张A100的算力集群就能亲眼见证一个模型从随机初始化的“哑巴”状态通过大量文本数据的学习逐渐掌握语言规律的全过程。这种亲手“创造智能”的体验是任何理论课程都无法替代的。2. 核心架构设计理解Transformer的骨架与灵魂要真正从零构建一个语言模型我们必须深入其核心——Transformer架构。虽然如今各种变体层出不穷但万变不离其宗。我们这个百万参数级别的LLM将基于最经典的Transformer解码器结构进行构建这是GPT系列模型成功的基石。2.1 自注意力机制模型理解上下文的钥匙自注意力机制是Transformer的灵魂它让模型能够在处理一个词时动态地关注输入序列中所有其他词的重要性。想象一下你在读一句话“苹果很好吃我喜欢它的口感。”为了理解“它”指代什么你需要回顾前文注意到“苹果”。自注意力机制就是让模型自动完成这种“回顾”和“关联”的计算。其核心是“查询-键-值”计算。对于输入序列中的每个词我们通过可学习的权重矩阵将其转换为三个向量查询向量、键向量和值向量。注意力分数的计算本质上是查询向量与所有键向量的点积用以衡量当前词与其他词的相关性。公式可以简化为Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。这里除以sqrt(d_k)是一个关键技巧称为缩放点积注意力。当键向量的维度d_k较大时点积的结果可能落入softmax函数梯度极小的区域导致训练困难缩放操作可以缓解这个问题。在我们的实现中通常会采用多头注意力。与其只做一次注意力计算不如将模型划分为多个“头”每个头在不同的子空间里学习不同的关注模式。比如一个头可能专注于语法结构另一个头可能专注于指代关系。最后将所有头的输出拼接起来再经过一个线性变换。这种设计极大地增强了模型的表征能力。注意在实现时为了效率我们会对整个序列的Q、K、V矩阵进行批量矩阵运算而不是循环处理每个词。同时在训练语言模型时必须使用因果掩码确保当前位置只能关注到它之前的位置包括自身而不能“偷看”未来的词这是保证模型只能用于序列生成而非“作弊”的关键。2.2 前馈网络与残差连接非线性变换与训练稳定器注意力层的输出会被送入一个前馈神经网络。这个FFN通常是一个两层的全连接网络中间包含一个非线性激活函数如ReLU或GELU。它的作用是对每个位置的表示进行独立且复杂的非线性变换。虽然注意力机制擅长捕捉词与词之间的关系但单个位置的语义深化和特征转换则需要FFN来完成。这里有一个至关重要的设计残差连接。在每个子层如自注意力层、FFN层周围我们都会添加一个残差连接即子层输出 LayerNorm(x Sublayer(x))。其中x是输入Sublayer(x)是子层本身的变换。残差连接最初是为了解决深度神经网络中的梯度消失问题而提出的它允许梯度直接流过网络极大地促进了深层网络的训练。在我们的LLM中它使得即使堆叠了数十层模型也能被有效优化。层归一化是另一个稳定训练的关键组件。它作用于每一个样本的每一个特征维度上将其均值和方差归一化。与批归一化不同LN不依赖于批次大小因此对小批量训练和变长序列更加友好。在Transformer中LN通常被放在残差连接之后即Post-LN结构也有一些变体放在之前Pre-LN。从实践来看Pre-LN结构通常能让训练过程更稳定收敛更快这对于我们从零开始训练尤为重要。2.3 位置编码为序列注入顺序信息自注意力机制本身是对位置不敏感的打乱输入词的顺序其输出的集合是不变的。但语言中词的顺序至关重要。因此我们需要显式地将位置信息注入到模型中。最经典的方法是使用正弦和余弦函数生成的位置编码将其与词嵌入向量相加。正弦位置编码的公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model))和PE(pos, 2i1) cos(pos / 10000^(2i/d_model))。其中pos是位置i是维度索引。这种编码的特点是对于固定的偏移量kPE(posk)可以表示为PE(pos)的线性函数这使得模型能够轻松地学习到相对位置关系。在我们的项目中考虑到百万参数的规模也可以考虑使用可学习的位置嵌入。即为每个位置分配一个可学习的向量与词嵌入相加。这种方法在训练数据足够的情况下可能更灵活但失去了正弦编码的外推性即处理比训练时更长的序列的能力。一个折中的方案是在底层使用正弦编码保证外推同时在模型高层加入可学习的相对位置偏置以增强灵活性。3. 从零开始的实现拆解代码层面的核心模块理论清晰后我们进入实战环节。我们将使用PyTorch框架因为它动态图的特点非常适合研究和实验。下面我将拆解几个最核心的模块实现并附上关键代码和解释。3.1 词嵌入与位置编码层这是模型接触数据的第一站。词嵌入层本质上是一个查找表将词汇表中的每个词索引映射为一个高维稠密向量。import torch import torch.nn as nn import math class Embeddings(nn.Module): def __init__(self, vocab_size, d_model): super(Embeddings, self).__init__() self.lut nn.Embedding(vocab_size, d_model) self.d_model d_model # 模型维度用于后续缩放 def forward(self, x): # x: [batch_size, seq_len] # 乘以 sqrt(d_model) 是Transformer论文中的做法用于在求和前平衡嵌入向量的尺度 return self.lut(x) * math.sqrt(self.d_model) class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout0.1, max_len5000): super(PositionalEncoding, self).__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # [max_len, 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) # 偶数维度用sin pe[:, 1::2] torch.cos(position * div_term) # 奇数维度用cos pe pe.unsqueeze(0) # [1, max_len, d_model] 便于广播 self.register_buffer(pe, pe) # 注册为缓冲区不参与训练但会保存到模型状态 def forward(self, x): # x: [batch_size, seq_len, d_model] x x self.pe[:, :x.size(1)] # 只取前seq_len个位置 return self.dropout(x)实操心得nn.Embedding默认的初始化是均值为0、标准差为1的正态分布。在一些更精细的实现中可能会采用更小的初始化范围如init_range 0.02这有助于训练初期的稳定性。位置编码的dropout是一个有效的正则化手段不要省略。3.2 缩放点积注意力与多头注意力实现这是最核心的计算单元。我们需要高效地实现因果自注意力。def attention(query, key, value, maskNone, dropoutNone): 缩放点积注意力计算 query, key, value: [batch_size, num_heads, seq_len, d_k] mask: 掩码用于遮挡无效位置如padding或未来信息 d_k query.size(-1) scores torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) # [batch_size, num_heads, seq_len, seq_len] if mask is not None: # 将mask中为0的位置需要被遮挡替换为一个极小的负数这样在softmax后权重接近0 scores scores.masked_fill(mask 0, -1e9) p_attn torch.softmax(scores, dim-1) # 在最后一个维度key的序列维度做softmax if dropout is not None: p_attn dropout(p_attn) return torch.matmul(p_attn, value), p_attn class MultiHeadedAttention(nn.Module): def __init__(self, num_heads, d_model, dropout0.1): super(MultiHeadedAttention, self).__init__() assert d_model % num_heads 0 self.d_k d_model // num_heads self.num_heads num_heads self.linears nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(4)]) # Q, K, V, 输出投影 self.attn None # 用于保存注意力权重便于可视化 self.dropout nn.Dropout(pdropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) # 1) 线性投影并分头 query, key, value [ lin(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) for lin, x in zip(self.linears, (query, key, value)) ] # 投影后形状: [batch_size, num_heads, seq_len, d_k] # 2) 计算注意力 x, self.attn attention(query, key, value, maskmask, dropoutself.dropout) # 3) 合并多头并做最终投影 x x.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k) return self.linears[-1](x)关键细节mask的构造至关重要。我们需要两种掩码一是填充掩码用于忽略序列中pad符号二是因果掩码下三角矩阵用于防止解码器“看到”未来信息。在训练时通常需要将两者结合。contiguous()的调用是因为transpose操作可能改变内存布局而view需要连续的内存。3.3 前馈网络与解码器层组装将上述组件与层归一化、残差连接组装起来形成一个完整的解码器层。class PositionwiseFeedForward(nn.Module): 两层全连接前馈网络中间有激活函数和dropout def __init__(self, d_model, d_ff, dropout0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 nn.Linear(d_model, d_ff) self.w_2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.activation nn.GELU() # 相比ReLUGELU更平滑在Transformer中效果通常更好 def forward(self, x): return self.w_2(self.dropout(self.activation(self.w_1(x)))) class DecoderLayer(nn.Module): 单个Transformer解码器层 def __init__(self, d_model, num_heads, d_ff, dropout0.1): super(DecoderLayer, self).__init__() self.self_attn MultiHeadedAttention(num_heads, d_model, dropout) self.feed_forward PositionwiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, mask): x: [batch_size, seq_len, d_model] mask: [batch_size, 1, seq_len, seq_len] 或 [batch_size, seq_len, seq_len] # 子层1: 带掩码的多头自注意力 残差 层归一化 (Pre-LN结构) attn_output self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), mask) x x self.dropout1(attn_output) # 子层2: 前馈网络 残差 层归一化 ff_output self.feed_forward(self.norm2(x)) x x self.dropout2(ff_output) return x架构选择这里采用了Pre-LN结构先归一化再进入子层。与原始Transformer的Post-LN相比Pre-LN通常能让梯度流动更顺畅训练初期更稳定对于从头训练深度模型是更友好的选择。Dropout被应用在残差相加之后这是另一种常见的正则化位置。4. 训练流程与工程实践让模型真正“学会”有了模型架构下一步就是设计训练流程。这是将一堆数学公式和代码转化为一个真正具有语言能力的模型的关键。4.1 数据准备与批处理策略数据是模型的燃料。对于语言模型我们需要大量的纯文本数据。开源数据集如WikiText、OpenWebText都是不错的选择。数据处理流程通常包括清洗去除无关标记、HTML标签、规范化空格等。分词将文本转化为模型能理解的离散符号。对于百万参数模型Byte-Pair Encoding是一个高效且流行的选择。它能在子词级别进行切分平衡词汇表大小与序列长度。我们可以使用Hugging Face的tokenizers库来训练自己的BPE分词器。构建数据集将分词后的长文本切割成固定长度的片段如1024个token。同时需要生成相应的因果掩码和填充掩码。批处理需要特别小心。由于文本序列长度不一直接填充到最大长度会造成大量计算浪费。动态批处理或桶排序批处理是更好的选择将长度相近的样本放在同一个批次尽量减少该批次内的填充。PyTorch的DataLoader配合自定义的collate_fn函数可以实现这一点。from torch.nn.utils.rnn import pad_sequence def collate_fn(batch): batch: 一个列表每个元素是(token_ids,) 返回 padded_token_ids, attention_mask src_list [item[0] for item in batch] # 左侧填充对于因果语言模型左侧填充更常见 src pad_sequence(src_list, batch_firstTrue, padding_valuepad_idx) # 创建注意力掩码0代表padding位置 mask (src ! pad_idx).unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len] # 创建因果掩码下三角矩阵 seq_len src.size(1) causal_mask torch.tril(torch.ones(seq_len, seq_len)).bool().unsqueeze(0).unsqueeze(0) # [1,1,seq_len,seq_len] # 合并两种掩码 combined_mask mask causal_mask # 对于注意力计算需要将False需要被遮挡的位置设为0 combined_mask combined_mask.float() return src, combined_mask4.2 损失函数与优化器配置语言模型训练的标准损失函数是交叉熵损失。给定输入序列x_1, x_2, ..., x_T模型的目标是预测下一个词。因此我们将输入序列向右移动一位作为目标计算每个位置预测下一个词的交叉熵损失并取平均。优化器的选择对训练成功至关重要。AdamW优化器是目前训练Transformer模型的事实标准。它修正了原始Adam中权重衰减的实现能带来更好的泛化性能。学习率调度同样关键。我们通常使用带热启动的余弦退火或线性预热余弦退火策略。预热期让模型在训练初期以较小的学习率稳定更新参数避免震荡。import torch.optim as optim from torch.optim.lr_scheduler import LambdaLR def get_optimizer(model, lr, betas, weight_decay): # 通常不对LayerNorm和Bias参数进行权重衰减 no_decay [bias, LayerNorm.weight] optimizer_grouped_parameters [ { params: [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], weight_decay: weight_decay, }, { params: [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], weight_decay: 0.0, }, ] return optim.AdamW(optimizer_grouped_parameters, lrlr, betasbetas) def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles0.5): 创建带预热的余弦学习率调度器 def lr_lambda(current_step): if current_step num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) progress float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) return max(0.0, 0.5 * (1.0 math.cos(math.pi * float(num_cycles) * 2.0 * progress))) return LambdaLR(optimizer, lr_lambda)4.3 训练循环与梯度累积训练循环的编写需要兼顾效率和清晰度。关键步骤包括前向传播、损失计算、反向传播、梯度裁剪和参数更新。model.train() optimizer.zero_grad() total_loss 0 accumulation_steps 4 # 梯度累积步数模拟更大批次 for step, (batch_input, batch_mask) in enumerate(train_dataloader): batch_input batch_input.to(device) batch_mask batch_mask.to(device) # 前向传播 # 输入是序列目标是序列向右移一位 logits model(batch_input, batch_mask) # [batch, seq_len, vocab_size] shift_logits logits[:, :-1, :].contiguous() # 预测部分 shift_labels batch_input[:, 1:].contiguous() # 目标部分 loss criterion(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) loss loss / accumulation_steps # 损失按累积步数缩放 # 反向传播 loss.backward() if (step 1) % accumulation_steps 0: # 梯度裁剪防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() optimizer.zero_grad() total_loss loss.item() * accumulation_steps工程技巧梯度累积当你的GPU内存无法容纳理想的大批次时梯度累积是救星。它在多个小批次上累积梯度只在累积步数达到后才更新一次参数从而模拟了大批次训练的效果。注意损失需要除以累积步数以保证梯度大小的正确性。梯度裁剪是训练RNN和Transformer类模型的标配它能防止梯度变得过大导致训练不稳定。5. 调试、评估与迭代优化模型开始训练并不意味着结束恰恰是开始。我们需要一套方法来监控其学习过程评估其性能并据此进行调试和优化。5.1 训练过程监控与指标解读最基础的监控指标是训练损失和验证损失。训练损失应稳步下降验证损失在初期随训练损失下降后期可能开始上升这标志着过拟合的开始。绘制两者的曲线图至关重要。困惑度是评估语言模型最直观的指标。它衡量模型对数据的不确定性计算公式为PPL exp(loss)。困惑度越低越好。一个在WikiText-103上训练良好的模型困惑度可以降到20-30左右。对于我们的百万参数模型在较小数据集上能达到100以下的困惑度就是一个不错的起点。除了损失还应监控梯度范数如果梯度范数突然变得极大或接近于零可能意味着学习率不当、网络结构有问题或出现了梯度爆炸/消失。参数更新比率参数更新的幅度与参数本身大小的比率。这个值通常应该在一个合理的范围内例如1e-3左右过大或过小都可能是问题的信号。激活值分布使用TensorBoard或类似工具可视化各层激活值的直方图。如果大量激活值饱和如ReLU输出很多0或sigmoid输出接近0或1可能说明网络初始化不当或学习率太高。5.2 常见训练问题与排查清单从零训练一个LLM你几乎一定会遇到下面这些问题问题现象可能原因排查与解决思路损失不下降Nan/Inf1. 学习率过高。2. 网络初始化不当导致梯度爆炸。3. 数据中存在异常值或未处理的特殊字符。1. 大幅降低学习率如从1e-3降到1e-5试试。2. 检查权重初始化。对于Linear层可尝试nn.init.xavier_uniform_对于Embedding层可缩小初始化范围。3. 添加梯度裁剪clip_grad_norm_。4. 仔细检查数据预处理流程确保输入是合理的整数索引。损失下降很慢1. 学习率过低。2. 模型容量不足层数/维度太小。3. 优化器状态不佳如Adam的动量被污染。1. 尝试增大学习率或使用学习率finder工具寻找合适范围。2. 适当增加模型深度层数或宽度d_model。3. 在训练中期重启优化器重置动量有时有奇效。验证损失先降后升过拟合1. 模型复杂度过高训练数据不足。2. 训练时间过长。3. 正则化不足。1. 增加Dropout比率或在FFN、注意力输出后增加Dropout层。2. 使用权重衰减AdamW已包含。3. 尝试标签平滑Label Smoothing将硬标签0或1替换为软标签如0.1或0.9。4. 尽早停止训练Early Stopping。生成文本重复或 nonsense1. 训练不充分。2. 采样策略单一如总是贪婪解码。3. 模型学到了数据中的噪声或重复模式。1. 继续训练观察验证损失是否还有下降空间。2. 在推理时使用核采样或温度采样增加多样性。3. 检查训练数据质量去除大量重复、低质文本。5.3 模型评估与生成文本分析训练结束后我们需要定量和定性评估模型。定量评估可以用在保留的测试集上计算困惑度。定性评估则更有趣让模型生成文本。编写一个简单的生成函数使用自回归的方式每次生成一个词并将其作为下一次输入的上下文。def generate_text(model, prompt, tokenizer, max_len50, temperature0.8, top_k50): model.eval() tokens tokenizer.encode(prompt).ids generated tokens.copy() with torch.no_grad(): for _ in range(max_len): inputs torch.tensor([generated]).to(device) # 创建因果掩码 mask torch.tril(torch.ones(len(generated), len(generated))).bool().unsqueeze(0).unsqueeze(0).to(device) logits model(inputs, mask) # 取最后一个位置的logits next_token_logits logits[0, -1, :] / temperature # Top-k 采样 indices_to_remove next_token_logits torch.topk(next_token_logits, top_k)[0][..., -1, None] next_token_logits[indices_to_remove] -float(Inf) probs torch.softmax(next_token_logits, dim-1) next_token torch.multinomial(probs, num_samples1).item() generated.append(next_token) if next_token tokenizer.token_to_id([EOS]): # 结束符 break return tokenizer.decode(generated)分析生成文本时关注以下几点连贯性句子在语法和基本逻辑上是否通顺相关性生成内容是否与提示相关多样性多次生成的结果是否丰富还是总在重复几个模式事实性对于小模型要求不能太高是否会产生明显的事实错误如果生成效果不佳回到训练监控指标检查模型是否欠拟合训练损失还很高或过拟合训练损失低但生成效果差。可能需要调整模型大小、数据量、正则化强度或训练时长。6. 项目总结与进阶思考完成一个百万参数LLM的从零构建其意义远超得到一个可用的模型。你获得的是对Transformer架构每一个运算细节的掌控是对语言模型训练中数据流、梯度流、损失变化的直观感受是对“为什么这样设计”的深刻理解。这份经验是阅读十篇论文也无法替代的。在实践过程中有几个点我体会尤深。第一是初始化的重要性。一个合适的初始化能让你在训练初期就走在正确的道路上否则可能连损失都降不下去。第二是学习率调度与批量大小的协同。它们共同决定了优化过程的“步伐”需要耐心调试。第三是监控与日志。一定要把训练过程中的关键指标损失、学习率、梯度范数可视化出来它们是诊断模型健康状况的唯一依据。这个项目本身也是一个绝佳的起点。在此基础上你可以进行无数有趣的扩展扩大规模尝试将参数增加到千万级观察性能变化体会“规模定律”。修改架构尝试替换激活函数如Swish使用不同的归一化方式如RMSNorm或引入最新的注意力变体如FlashAttention。改进训练尝试不同的优化器如Lion引入更复杂的调度策略或者实现混合精度训练以节省显存、加快速度。转向指令微调收集指令-回答对数据在预训练好的模型基础上进行监督微调让你的模型学会遵循指令这是通向ChatGPT类模型的关键一步。最后记得将你的代码、配置和实验记录妥善保存。重现实验结果是研究的基本要求而一个混乱的代码库会让你在尝试新想法时举步维艰。这个项目就像你亲手打造的第一把工具它可能粗糙但每一个细节都凝聚着你的理解而这正是你构建更复杂、更强大AI系统的基石。