1. 从零构建一个230万参数的语言模型一次深度实践如果你对ChatGPT、LLaMA这些大语言模型LLM的内部运作感到好奇但又觉得动辄数十亿参数的庞然大物遥不可及那么这篇文章就是为你准备的。我最近花了不少时间基于Meta开源的LLaMA论文架构亲手从零实现了一个仅有230万参数的“微型”语言模型。整个过程没有使用昂贵的GPU仅凭一台普通笔记本电脑和Python就完成了训练和推理。这听起来可能有些不可思议但它的确证明了理解LLM的核心原理并亲手搭建一个远比想象中要触手可及。市面上很多教程和博客热衷于讨论Transformer的理论却对具体的代码实现和实操细节语焉不详。这常常让学习者陷入“道理都懂就是无从下手”的困境。本文将反其道而行之我会带你一步步走完数据准备、模型架构复现、训练调优和文本生成的完整流程。我们将使用经典的TinyShakespeare数据集它虽然只有约100万个字符但足以让我们验证模型的核心能力。我们的目标是不依赖任何高级框架的封装用最清晰的代码揭示一个可工作的LLM是如何被构建出来的。无论你是希望深入理解LLM机制的研究者还是渴望动手实践的开发者这篇文章都将提供一份可直接运行、可修改的“蓝图”。2. 项目核心思路与架构选型2.1 为什么选择LLaMA架构作为蓝本在众多LLM架构中我选择复现LLaMA主要基于以下几点考量首先LLaMA的设计在效率和效果上取得了很好的平衡。它并非简单堆砌参数而是通过一系列精心设计的改进让模型在更小的参数量下达到甚至超越同类模型的性能。这对于我们构建一个“微型”模型具有重要的指导意义。我们不需要盲目追求参数规模而是应该学习如何让有限的参数发挥最大的效能。其次LLaMA的改进点清晰、独立便于模块化实现。相较于原始TransformerLLaMA主要引入了三项关键改进RMSNorm预归一化、RoPE旋转位置编码和SwiGLU激活函数。每一项改进都有明确的数学定义和独立的实现模块这非常适合教学和拆解。我们可以像搭积木一样将这些模块逐个集成到我们的基础模型中并观察每一步带来的变化从而深刻理解每个组件的作用。最后其开源性和社区支持。LLaMA论文提供了足够详细的描述并且有大量的开源实现和讨论可供参考。这降低了我们复现的难度当遇到理解障碍时可以方便地查阅相关资料进行验证。我们的实现虽然简化但力求在核心思想上与原文保持一致。2.2 整体技术路线图我们的项目将遵循一个自底向上、迭代优化的路径环境与数据准备搭建最基础的Python环境获取并预处理TinyShakespeare数据集构建数据管道Data Pipeline。搭建基线模型实现一个最朴素的、基于嵌入层和前馈网络的神经网络模型。这个模型将作为我们的性能基准和后续改进的“脚手架”。逐步集成LLaMA组件第一步引入RMSNorm。替换掉基线模型中的层归一化LayerNorm观察训练稳定性和收敛速度的变化。第二步引入RoPE。实现旋转位置编码替换传统的绝对或相对位置编码理解其对长序列建模能力的提升。第三步引入SwiGLU。将前馈网络中的ReLU激活函数升级为SwiGLU增强模型的非线性表达能力。训练与评估设计训练循环监控训练集和验证集的损失并使用一个简单的文本生成函数来直观评估模型效果。超参数实验与分析在资源允许的范围内调整批次大小、上下文窗口、模型维度等关键超参数分析它们对模型性能和训练动态的影响。这个路线的优势在于每一步的改变都是增量式的并且效果可观测。你能清晰地看到每加入一个LLaMA的“零件”模型的损失曲线和生成文本的质量是如何一步步改善的。注意关于“从零开始”的定义这里“从零开始”指的是我们不依赖像Hugging Facetransformers这样的高级库来直接调用LlamaModel。我们会使用PyTorch等基础框架亲手实现每一个关键层如RMSNorm, RoPE Attention。这能确保你对模型每一个计算步骤都了如指掌。3. 数据预处理与评估策略详解3.1 数据集的选择与字符级分词我们选择了TinyShakespeare数据集。这是一个极简的文本数据集包含莎士比亚作品中的约4万行文本总计约100万个字符。选择它原因有三一是体积小可以在CPU上快速进行迭代实验二是作为英文文本其词汇分布和语法结构相对规范适合作为语言模型的“入门沙盒”三是它被广泛用于教学如Karpathy的makemore系列社区资源丰富。与LLaMA等现代LLM使用复杂的子词分词器如SentencePiece BPE不同为了极致简化我们采用了字符级分词。这意味着我们的词汇表vocab就是数据集中所有不重复的字符集合。对于TinyShakespeare这个词汇表大小是65包括大小写字母、标点、换行符等。# 读取数据并构建词汇表 with open(tinyshakespeare.txt, r) as f: text f.read() vocab sorted(list(set(text))) vocab_size len(vocab) # 输出: 65 # 创建字符到索引编码和索引到字符解码的映射 char_to_idx {ch: i for i, ch in enumerate(vocab)} idx_to_char {i: ch for i, ch in enumerate(vocab)} def encode(s): 将字符串编码为整数列表 return [char_to_idx[ch] for ch in s] def decode(l): 将整数列表解码回字符串 return .join([idx_to_char[i] for i in l]) # 测试 print(encode(hello)) # 输出例如: [46, 43, 50, 50, 53] print(decode([46, 43, 50, 50, 53])) # 输出: hello字符级分词的优点是实现简单没有未登录词问题。缺点是序列长度会很长一个单词需要多个token且模型需要从字符开始学习拼写规则任务更难。但对于我们的小规模实验这完全可接受。3.2 构建数据批次Batching生成器语言模型训练通常使用“滑动窗口”的方式。给定一个长文本序列我们将其切割成许多固定长度context_window的连续片段作为输入而目标label则是将输入序列向后移动一位。例如对于文本“hello world”如果context_window5那么一个样本可能是输入“hello”目标“ello ”。我们需要一个函数能随机地从训练集或验证集中抽取一批batch_size个这样的样本对。import torch def get_batches(data, split, batch_size, context_window): 从数据中获取批次。 data: 整个编码后的数据集张量 split: train, val, 或 test batch_size: 批次大小 context_window: 上下文窗口长度 # 按8:1:1划分训练、验证、测试集 n len(data) train_data data[:int(0.8*n)] val_data data[int(0.8*n):int(0.9*n)] test_data data[int(0.9*n):] # 根据split选择数据源 split_data {train: train_data, val: val_data, test: test_data}[split] # 随机生成起始索引确保不越界 start_idxs torch.randint(0, len(split_data) - context_window, (batch_size,)) # 构建输入x和目标y x torch.stack([split_data[i:icontext_window] for i in start_idxs]) y torch.stack([split_data[i1:icontext_window1] for i in start_idxs]) return x, y # 将整个文本编码并转为张量 data torch.tensor(encode(text), dtypetorch.long) # 获取一个训练批次 x_batch, y_batch get_batches(data, train, batch_size4, context_window8) print(f输入形状: {x_batch.shape}) # torch.Size([4, 8]) print(f目标形状: {y_batch.shape}) # torch.Size([4, 8])这个函数是数据供给的核心。batch_size决定了每次梯度更新时使用的样本数量context_window决定了模型一次能“看到”多长的历史信息。这两个是影响模型性能和训练效率的关键超参数。3.3 评估策略在训练中持续监控在深度学习项目中尤其是在资源有限的情况下我们不能盲目训练成百上千个epoch然后才看结果。我们需要一个在训练过程中持续评估模型在未见数据上表现的方法即监控验证损失validation loss。import numpy as np torch.no_grad() # 禁用梯度计算节省内存和计算资源 def evaluate_loss(model, data, config): 评估模型在训练集和验证集上的平均损失 model.eval() # 将模型设置为评估模式影响Dropout、BatchNorm等层 out {} for split in [train, val]: losses [] # 多次采样取平均使评估更稳定 for _ in range(20): xb, yb get_batches(data, split, config[batch_size], config[context_window]) _, loss model(xb, yb) # 只取损失不取logits losses.append(loss.item()) out[split] np.mean(losses) model.train() # 恢复训练模式 return out这个evaluate_loss函数会在每个训练日志间隔被调用。它计算模型在训练集和验证集上各20个随机批次的平均损失。验证损失是判断模型是否过拟合overfitting的关键指标。理想情况下训练损失和验证损失应该同步下降。如果训练损失持续下降而验证损失开始上升或停滞通常意味着模型开始记忆训练数据而非学习泛化模式。实操心得验证集的重要性很多初学者会只关注训练损失觉得它越低越好。这是一个误区。在资源有限的小模型上过拟合会来得非常快。坚持在每个epoch或每N个step后计算验证损失并绘制损失曲线图是防止无效训练、及时调整策略如早停、增加正则化的最有效手段。我们的训练函数会将这个评估过程集成进去。4. 搭建基线神经网络模型在引入LLaMA的复杂组件之前我们先构建一个最简单的模型作为基线。这个模型将帮助我们建立整个训练流程并提供一个性能对比的基准。4.1 模型结构设计我们的基线模型非常简单嵌入层Embedding Layer将输入的字符索引整数映射为稠密向量d_model维。线性层Linear Layer一个简单的全连接网络将嵌入向量变换回词汇表大小的空间。输出直接输出未归一化的分数logits用于计算交叉熵损失。import torch.nn as nn import torch.nn.functional as F class SimpleModel(nn.Module): def __init__(self, config): super().__init__() self.config config # 嵌入层词汇表大小65映射到d_model维空间 self.embedding nn.Embedding(config[vocab_size], config[d_model]) # 一个简单的两层前馈网络使用ReLU激活 self.ffn nn.Sequential( nn.Linear(config[d_model], config[d_model]), nn.ReLU(), nn.Linear(config[d_model], config[vocab_size]) ) # 打印参数量 total_params sum(p.numel() for p in self.parameters()) print(f基线模型总参数量: {total_params}) def forward(self, idx, targetsNone): # idx 形状: (batch_size, context_window) x self.embedding(idx) # 形状: (batch_size, context_window, d_model) logits self.ffn(x) # 形状: (batch_size, context_window, vocab_size) if targets is not None: # 计算损失。需要将logits和targets reshape成二维 # F.cross_entropy 期望输入形状为 (batch*seq_len, vocab_size) 和 (batch*seq_len,) loss F.cross_entropy(logits.view(-1, self.config[vocab_size]), targets.view(-1)) return logits, loss return logits这里有一个关键细节在forward函数中我们直接返回logits未经过softmax的分数而不是概率。这是因为PyTorch的F.cross_entropy函数内部已经集成了log_softmax和negative log-likelihood的计算并且数值上更稳定。直接传入logits是标准做法。4.2 训练循环的实现有了模型和数据我们需要一个训练循环来优化模型参数。import time import pandas as pd from matplotlib import pyplot as plt def train_model(model, data, config, optimizer, schedulerNone): losses [] start_time time.time() for epoch in range(config[epochs]): # 1. 清零梯度 optimizer.zero_grad() # 2. 获取一个数据批次 xb, yb get_batches(data, train, config[batch_size], config[context_window]) # 3. 前向传播计算损失 _, loss model(xb, yb) # 4. 反向传播计算梯度 loss.backward() # 5. 优化器更新参数 optimizer.step() if scheduler: scheduler.step() # 6. 定期记录和评估 if epoch % config[log_interval] 0: elapsed time.time() - start_time # 评估当前模型在训练和验证集上的损失 eval_results evaluate_loss(model, data, config) losses.append(eval_results) # 打印日志 print(fEpoch {epoch:4d} | 训练损失 {eval_results[train]:.4f} | f验证损失 {eval_results[val]:.4f} | 耗时 {elapsed:.2f}s) start_time time.time() # 重置计时器 # 训练结束后绘制损失曲线 df_losses pd.DataFrame(losses) df_losses.plot(title训练与验证损失曲线) plt.xlabel(评估间隔) plt.ylabel(交叉熵损失) plt.show() return model4.3 基线模型训练与结果分析现在让我们配置参数并运行基线模型。# 主配置字典 MASTER_CONFIG { vocab_size: 65, # 词汇表大小 d_model: 128, # 模型嵌入维度 batch_size: 32, # 批次大小 context_window: 16, # 上下文窗口长度 epochs: 2000, # 训练轮数 log_interval: 200, # 每200轮评估一次 } # 实例化模型、优化器 model SimpleModel(MASTER_CONFIG) optimizer torch.optim.Adam(model.parameters(), lr1e-3) # 开始训练 trained_model train_model(model, data, MASTER_CONFIG, optimizer)运行后你可能会看到类似以下的输出和损失曲线基线模型总参数量: 33281 Epoch 0 | 训练损失 4.1742 | 验证损失 4.1721 | 耗时 0.12s Epoch 200 | 训练损失 3.0123 | 验证损失 3.0256 | 耗时 12.34s Epoch 400 | 训练损失 2.7456 | 验证损失 2.7689 | 耗时 12.50s ... Epoch 1800 | 训练损失 2.5211 | 验证损失 2.5583 | 耗时 12.41s Epoch 2000 | 训练损失 2.5134 | 验证损失 2.5521 | 耗时 12.38s结果分析参数量基线模型只有约3.3万个参数非常小。损失值初始损失约4.17对于65个类别的随机猜测理论损失约为-ln(1/65) ≈ 4.17说明模型初始化是合理的。经过训练损失下降到2.55左右说明模型确实学到了一些模式。过拟合迹象训练损失和验证损失非常接近且同步下降在早期没有明显过拟合。这是因为模型容量很小数据集也不大。为了直观感受模型学到了什么我们实现一个简单的文本生成函数def generate_text(model, config, start_strThe , max_new_tokens100): model.eval() with torch.no_grad(): # 将起始字符串编码并转为模型输入格式 idx torch.tensor(encode(start_str)).unsqueeze(0) # 形状: (1, seq_len) generated list(idx[0].tolist()) for _ in range(max_new_tokens): # 取最后 context_window 个token作为输入 context idx[:, -config[context_window]:] if idx.shape[1] config[context_window] else idx # 前向传播 logits model(context) # 形状: (1, curr_len, vocab_size) # 取最后一个时间步的logits next_token_logits logits[:, -1, :] # 形状: (1, vocab_size) # 通过softmax得到概率分布并采样下一个token probs F.softmax(next_token_logits, dim-1) idx_next torch.multinomial(probs, num_samples1) # 形状: (1, 1) # 将新token追加到序列中 generated.append(idx_next.item()) idx torch.cat([idx, idx_next], dim1) model.train() return decode(generated) print(generate_text(trained_model, MASTER_CONFIG, start_strKING))基线模型的生成结果可能看起来是乱码或重复的字符序列比如“KINGx;x;x;x;x;x;...”。这完全在意料之中一个仅3万参数、没有注意力机制的简单前馈网络很难捕捉长距离的语法和语义依赖。接下来我们就开始为它注入LLaMA的“灵魂”。5. 集成LLaMA核心组件RMSNorm, RoPE, SwiGLU现在我们将把LLaMA论文中的三个核心改进逐一集成到我们的基线模型中。我们将采取迭代的方式每加入一个组件就重新训练并观察效果变化。5.1 组件一RMSNorm均方根归一化原理与动机传统的LayerNorm会对输入进行减均值、除以标准差的操作。RMSNorm的作者发现减去均值的操作并非必要仅通过重新缩放除以RMS即均方根就能达到相似的性能同时计算量更小。公式如下RMSNorm(x) (x / RMS(x)) * g其中RMS(x) sqrt(mean(x_i^2))g是一个可学习的缩放参数。实现代码class RMSNorm(nn.Module): def __init__(self, dim, eps1e-8): super().__init__() self.eps eps # 可学习的缩放参数初始化为全1 self.scale nn.Parameter(torch.ones(dim)) def forward(self, x): # x 形状: (batch, seq_len, dim) # 计算RMS均方根: sqrt(mean(x^2)) rms torch.sqrt(torch.mean(x.pow(2), dim-1, keepdimTrue) self.eps) # 归一化并缩放 x_norm x / rms return self.scale * x_norm集成到模型中我们将RMSNorm层插入到嵌入层之后前馈网络之前。这种“预归一化”Pre-Norm是Transformer架构的常见做法有助于稳定训练。class SimpleModelWithRMSNorm(nn.Module): def __init__(self, config): super().__init__() self.config config self.embedding nn.Embedding(config[vocab_size], config[d_model]) # 新增RMSNorm层 self.rms_norm RMSNorm(config[d_model]) self.ffn nn.Sequential( nn.Linear(config[d_model], config[d_model]), nn.ReLU(), nn.Linear(config[d_model], config[vocab_size]) ) print(f模型总参数量 (含RMSNorm): {sum(p.numel() for p in self.parameters())}) def forward(self, idx, targetsNone): x self.embedding(idx) x self.rms_norm(x) # 应用RMSNorm logits self.ffn(x) if targets is not None: loss F.cross_entropy(logits.view(-1, self.config[vocab_size]), targets.view(-1)) return logits, loss return logits训练与观察使用相同的配置和训练循环重新训练这个模型。你可能会观察到参数量略有增加因为多了scale参数。训练初期可能更稳定损失下降曲线可能更平滑。RMSNorm通过稳定激活值的尺度有时能带来更快的收敛或允许使用更大的学习率。5.2 组件二RoPE旋转位置编码原理与动机Transformer本身没有位置信息需要位置编码。RoPE的巧妙之处在于它通过在特征空间的复数域进行旋转来编码位置信息。对于位置m的token其第i个维度的嵌入向量会被乘以一个旋转矩阵R^m_θ。这使得两个token之间的注意力分数仅依赖于它们的相对位置m-n并且会随着相对距离增大而自然衰减这符合自然语言的直觉。实现旋转矩阵def get_rotary_matrix(context_window, d_model): 生成RoPE旋转矩阵R。 R的形状: (context_window, d_model, d_model) 这里实现一个简化版本实际论文中θ_i是预设的基数。 R torch.zeros((context_window, d_model, d_model)) for pos in range(context_window): for i in range(d_model // 2): # 按维度对处理 theta 10000.0 ** (-2 * i / d_model) m_theta pos * theta # 构建2x2旋转矩阵并填充到高维张量中 R[pos, 2*i, 2*i] torch.cos(m_theta) R[pos, 2*i, 2*i1] -torch.sin(m_theta) R[pos, 2*i1, 2*i] torch.sin(m_theta) R[pos, 2*i1, 2*i1] torch.cos(m_theta) return R集成RoPE的注意力头为了演示我们实现一个极简的单头注意力并应用RoPE。class RoPESimpleAttentionHead(nn.Module): def __init__(self, config): super().__init__() self.config config self.d_model config[d_model] # 定义Q, K, V的投影矩阵 self.w_q nn.Linear(self.d_model, self.d_model, biasFalse) self.w_k nn.Linear(self.d_model, self.d_model, biasFalse) self.w_v nn.Linear(self.d_model, self.d_model, biasFalse) # 预计算旋转矩阵推理时可缓存 self.R get_rotary_matrix(config[context_window], self.d_model) def forward(self, x): # x 形状: (batch, seq_len, d_model) batch, seq_len, _ x.shape q self.w_q(x) # (batch, seq_len, d_model) k self.w_k(x) # (batch, seq_len, d_model) v self.w_v(x) # (batch, seq_len, d_model) # 应用RoPE将位置信息注入Q和K # 为了简化我们这里使用一个循环。高效实现应使用向量化操作。 q_rotated torch.zeros_like(q) k_rotated torch.zeros_like(k) for pos in range(seq_len): q_rotated[:, pos, :] torch.matmul(q[:, pos, :], self.R[pos].T) k_rotated[:, pos, :] torch.matmul(k[:, pos, :], self.R[pos].T) # 计算注意力分数 (简化版未做缩放和掩码) # attn softmax(Q * K^T) * V attn_scores torch.matmul(q_rotated, k_rotated.transpose(-2, -1)) # (batch, seq_len, seq_len) attn_weights F.softmax(attn_scores / (self.d_model ** 0.5), dim-1) # 缩放 output torch.matmul(attn_weights, v) # (batch, seq_len, d_model) return output将注意力头加入模型现在我们用这个带RoPE的注意力头替换掉第一个线性层。class SimpleModelWithRoPE(nn.Module): def __init__(self, config): super().__init__() self.config config self.embedding nn.Embedding(config[vocab_size], config[d_model]) self.rms_norm1 RMSNorm(config[d_model]) # 使用RoPE注意力头 self.rope_attn RoPESimpleAttentionHead(config) self.rms_norm2 RMSNorm(config[d_model]) self.ffn nn.Sequential( nn.Linear(config[d_model], config[d_model]), nn.ReLU(), nn.Linear(config[d_model], config[vocab_size]) ) print(f模型总参数量 (含RMSNormRoPE): {sum(p.numel() for p in self.parameters())}) def forward(self, idx, targetsNone): x self.embedding(idx) x self.rms_norm1(x) # 残差连接 (简化版未做投影) x x self.rope_attn(x) x self.rms_norm2(x) logits self.ffn(x) if targets is not None: loss F.cross_entropy(logits.view(-1, self.config[vocab_size]), targets.view(-1)) return logits, loss return logits训练与观察加入RoPE后模型参数量会显著增加因为多了Q、K、V的投影矩阵。重新训练后你可能会发现模型对序列中token的顺序更敏感。在文本生成任务中可能会观察到比基线模型更连贯的短短语或单词组合。因为RoPE赋予了模型理解相对位置的能力。5.3 组件三SwiGLU激活函数原理与动机SwiGLU是GLUGated Linear Unit结构的一种变体结合了Swish激活函数。其公式为SwiGLU(x, W, V, b, c) Swish(xW b) ⊗ (xV c)其中⊗是逐元素乘法。门控机制xVc可以控制信息流动让模型学会“忘记”或“传递”哪些信息这比简单的ReLU提供了更强的表达能力。实现SwiGLUclass SwiGLU(nn.Module): def __init__(self, dim): super().__init__() # 两个线性变换对应公式中的W和V self.w nn.Linear(dim, dim, biasFalse) self.v nn.Linear(dim, dim, biasFalse) # Swish激活函数: x * sigmoid(beta * x)通常beta1 self.swish lambda x: x * torch.sigmoid(x) def forward(self, x): # Swish(xW) ⊙ (xV) return self.swish(self.w(x)) * self.v(x)集成到前馈网络我们用SwiGLU模块替换掉原来FFN中的第一个线性层和ReLU激活。class SimpleModelWithAllComponents(nn.Module): def __init__(self, config): super().__init__() self.config config self.embedding nn.Embedding(config[vocab_size], config[d_model]) self.rms_norm1 RMSNorm(config[d_model]) self.rope_attn RoPESimpleAttentionHead(config) self.rms_norm2 RMSNorm(config[d_model]) # 使用SwiGLU的前馈网络 self.ffn nn.Sequential( SwiGLU(config[d_model]), # 替换原来的 LinearReLU nn.Linear(config[d_model], config[vocab_size]) # 输出层 ) total_params sum(p.numel() for p in self.parameters()) print(f完整模型总参数量 (RMSNormRoPESwiGLU): {total_params}) def forward(self, idx, targetsNone): x self.embedding(idx) x self.rms_norm1(x) x x self.rope_attn(x) # 残差连接 x self.rms_norm2(x) logits self.ffn(x) if targets is not None: loss F.cross_entropy(logits.view(-1, self.config[vocab_size]), targets.view(-1)) return logits, loss return logits最终模型训练现在我们拥有了一个集成了LLaMA三大核心组件的“完整版”模型。重新训练这个模型并观察其最终表现。# 更新配置可能需要调整学习率因为模型更复杂了 MASTER_CONFIG.update({ epochs: 3000, log_interval: 300, }) final_model SimpleModelWithAllComponents(MASTER_CONFIG) optimizer torch.optim.Adam(final_model.parameters(), lr5e-4) # 调小学习率 trained_final_model train_model(final_model, data, MASTER_CONFIG, optimizer)最终结果对比 训练完成后你可以系统地对比四个模型的最终验证损失和生成文本质量基线模型 (SimpleModel)验证损失 ~2.55生成文本为乱码。RMSNorm模型验证损失可能略低如~2.50训练更稳定。RMSNormRoPE模型验证损失进一步下降如~2.40生成文本可能出现一些重复的单词对或短句片段。完整模型 (All Components)验证损失应是最低的可能达到~2.30或更低。虽然受限于模型规模和数据集生成完整的莎士比亚句子仍很困难但你应该能观察到明显的进步比如生成的文本中大写字母、标点符号的使用更合理字符组合开始像真实的英文单词。重要提示关于注意力头的简化为了保持代码清晰和可理解我们实现的是单头、无因果掩码causal mask的简化版注意力。在一个真正的自回归语言模型中必须使用因果掩码来确保当前位置的预测只能看到之前的位置信息。此外多头注意力Multi-Head Attention能并行学习不同子空间的特征。如果你希望模型性能有质的提升下一步就是实现带因果掩码的多头RoPE注意力。6. 超参数实验与模型优化在核心架构确定后超参数调优是提升模型性能的关键。我们的模型很小这反而让我们可以快速进行一些实验。6.1 关键超参数解析d_model(模型维度)这是嵌入向量和模型内部表示的维度。它直接决定了模型的容量。增大d_model会显著增加参数量大约与平方成正比和计算量。对于我们的微型模型可以从128开始尝试逐步增加到256或512观察损失下降和生成质量的变化但要注意训练时间会变长。context_window(上下文窗口)模型一次能处理的最大token数量。增大窗口能让模型看到更长的历史信息理论上对语言建模有利但也会增加计算复杂度和内存消耗注意力计算是序列长度的平方级。莎士比亚文本中有长句子可以尝试从16增加到32或64。batch_size(批次大小)每次参数更新使用的样本数。较大的批次通常能提供更稳定的梯度估计可能允许使用更大的学习率但需要更多内存。在CPU上训练批次大小受限于内存。可以尝试在8到64之间调整。学习率 (Learning Rate)最重要的优化器超参数。对于Adam优化器1e-3或5e-4是常见的起点。模型变复杂后可能需要调小学习率如1e-4来保证稳定训练。可以使用学习率调度器如CosineAnnealingLR在训练中动态调整。6.2 设计一个简单的实验我们可以固定其他参数系统地改变其中一个来观察验证损失的变化。def hyperparameter_experiment(param_name, param_values, base_config, num_epochs1000): 运行超参数实验 param_name: 要测试的超参数名如 d_model param_values: 该超参数的取值列表如 [64, 128, 256] base_config: 基础配置字典 num_epochs: 每个实验运行的epoch数 results {} for value in param_values: print(f\n 实验 {param_name}{value} ) config base_config.copy() config[param_name] value config[epochs] num_epochs # 重新初始化模型和优化器 model SimpleModelWithAllComponents(config) optimizer torch.optim.Adam(model.parameters(), lrconfig.get(lr, 5e-4)) # 训练并记录最终验证损失 final_val_loss train_and_return_final_loss(model, data, config, optimizer) results[value] final_val_loss print(f最终验证损失: {final_val_loss:.4f}) return results # 辅助函数训练并返回最终验证损失 def train_and_return_final_loss(model, data, config, optimizer): # ... 简化版的训练循环只返回最后一个评估点的验证损失 ... for epoch in range(config[epochs]): optimizer.zero_grad() xb, yb get_batches(data, train, config[batch_size], config[context_window]) _, loss model(xb, yb) loss.backward() optimizer.step() if epoch % config[log_interval] 0: eval_loss evaluate_loss(model, data, config)[val] return eval_loss # 示例测试不同的 d_model base_config MASTER_CONFIG.copy() base_config.update({epochs: 500, log_interval: 100}) d_model_results hyperparameter_experiment(d_model, [64, 128, 256], base_config) print(d_model_results)通过这样的实验你可能会发现d_model256时模型性能最好但训练速度最慢。context_window从16增加到32可能带来明显收益但再增加到64可能收益递减因为数据集中长距离依赖有限。6.3 模型保存与加载训练一个模型可能需要数小时保存训练好的模型至关重要。def save_model(model, config, optimizer, pathmy_llama_model.pth): torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), config: config, }, path) print(f模型已保存至 {path}) def load_model(path, model_class): checkpoint torch.load(path, map_locationtorch.device(cpu)) config checkpoint[config] model model_class(config) model.load_state_dict(checkpoint[model_state_dict]) # 如果需要继续训练也可以加载优化器状态 # optimizer.load_state_dict(checkpoint[optimizer_state_dict]) print(f模型已从 {path} 加载) return model, config # 保存最终模型 save_model(trained_final_model, MASTER_CONFIG, optimizer, final_2.3M_model.pth)7. 常见问题、排查技巧与扩展方向在复现和实践过程中你几乎一定会遇到各种问题。以下是我踩过的一些坑以及解决方案。7.1 训练问题排查表问题现象可能原因排查步骤与解决方案损失值为NaN或无限大1. 学习率过高。2. 网络层中出现了除零或数值溢出例如RMSNorm中eps太小。3. 梯度爆炸。1.立即降低学习率尝试1e-4,5e-5。2. 检查RMSNorm实现确保eps被加到分母中。3. 使用梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。损失不下降一直在初始值附近波动1. 学习率太低。2. 模型初始化不当或某些层如嵌入层的梯度为零。3. 数据或标签有问题。1. 尝试增大学习率。2. 检查模型参数是否被正确注册和更新。打印部分参数的梯度看看。3. 检查get_batches函数确保x和y的对应关系正确y应该是x向右移动一位。验证损失先降后升过拟合1. 模型容量 (d_model) 相对于数据太大。2. 训练轮数过多。1. 减小d_model或增加正则化如Dropout但我们的小模型暂未加。2. 使用早停Early Stopping当验证损失连续N个epoch不再下降时停止训练。训练速度极慢1. 在CPU上训练且模型/批次过大。2. 代码中存在低效的Python循环如我们RoPE的循环实现。1. 减小batch_size或context_window。2.向量化RoPE计算这是性能瓶颈。可以使用爱因斯坦求和torch.einsum或预先计算所有位置的旋转矩阵并批量运算。生成文本全是重复字符或标点1. 模型训练不充分只学到了最简单的统计规律如空格、标点的频率。2. 温度参数太低采样缺乏随机性。1. 增加训练轮数检查损失是否还在下降。2. 在生成时引入温度Temperature和Top-k/p采样增加多样性。7.2 生成文本的改进技巧我们之前的generate_text函数使用了贪婪采样multinomial基于原始概率这容易导致重复和乏味的输出。可以尝试以下改进def generate_with_temperature(model, config, start_str, max_new_tokens200, temperature0.8, top_k40): 带温度和top-k采样的文本生成 model.eval() with torch.no_grad(): idx torch.tensor(encode(start_str)).unsqueeze(0) for _ in range(max_new_tokens): context idx[:, -config[context_window]:] if idx.shape[1] config[context_window] else idx logits model(context) next_token_logits logits[:, -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 F.softmax(next_token_logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat([idx, idx_next], dim1) model.train() return decode(idx[0].tolist()) # 尝试不同的温度和top-k print(温度0.5 (更确定):, generate_with_temperature(trained_final_model, MASTER_CONFIG, KING, temperature0.5, top_k10)) print(温度1.2 (更多样):, generate_with_temperature(trained_final_model, MASTER_CONFIG, KING, temperature1.2, top_k50))温度TemperatureT1使用原始logitsT1使概率分布更尖锐模型更自信输出更确定但可能重复T1使分布更平缓输出更多样但可能不合逻辑。Top-k采样只从概率最高的k个token中采样排除那些极不可能的选项能有效减少生成 nonsense 的概率。7.3 项目扩展方向如果你已经成功运行了230万参数的模型并想继续深入这里有几个方向实现真正的Transformer Block将我们的单头注意力扩展为多头注意力Multi-Head Attention并添加残差连接Residual Connection和前馈网络FFN的Dropout。这是构建强大LLM的基石。增加模型深度堆叠多个Transformer Block。你可以创建一个nn.ModuleList来存放多个相同的层。这是增加参数量的主要方式我们的模型目前只有“宽”度d_model缺乏深度。使用更大的数据集将TinyShakespeare换成WikiText-2、OpenWebText等更大的数据集。注意这需要更长的训练时间和可能更复杂的预处理如BPE分词。引入更高效的优化器尝试LLaMA论文中使用的AdamW优化器并配合余弦退火学习率调度。代码优化与GPU迁移将RoPE等计算密集型操作向量化并将模型和数据移至GPU使用.to(‘cuda’)以获得百倍的速度提升。添加评估指标除了损失计算在验证集上的困惑度Perplexity这是衡量语言模型质量的更直观指标。PPL exp(loss)。这个230万参数的项目是一个完美的起点。通过亲手实现每一个组件你获得的对LLM内部机制的理解远比调用十次预训练模型API要深刻得多。当你能清晰地解释RMSNorm为什么有效能在白板上画出RoPE的旋转示意图并理解SwiGLU中门控机制的作用时你就已经跨过了LLM原理学习的门槛。接下来的路无论是研读更复杂的论文还是尝试训练更大的模型都将有坚实的实践基础作为支撑。