ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

从零实现LLM:拆解斯坦福CS336课程与TinyGPT实战

从零实现LLM:拆解斯坦福CS336课程与TinyGPT实战 刚开始接触大模型时很多人都会经历这样一个阶段用 HuggingFace 的transformers加载一个现成模型跑推理或者套一套微调脚本感觉“我也会大模型了”。但一旦把问题从“调用”变成“实现”——比如让你自己写一个 Transformer、自己训练一个 BPE tokenizer、自己把训练循环跑通——很多人就会卡住。斯坦福 CS336 这门课正好就是解决这个问题的。它不像多数课程那样只讲概念而是要求你从数据准备、tokenizer、模型架构、训练循环到推理采样每个环节都亲手实现一遍。本文就以这门课为主线结合课程作业的思路完整拆解“从零搭建 LLM”的全过程并给出可以运行的实战代码。无论你是想理解大模型底层原理的学生还是正在做 LLM 应用、想深入底层机制的开发者这篇文章都能帮你把“黑盒”拆开。读完你不仅能理解 LLM 的工作原理还能在自己电脑上跑起来一个迷你版 GPT 模型。1. CS336 是什么一门逼你手写 LLM 的课程1.1 课程背景与定位CS336 是斯坦福开设的一门关于大语言模型LLM的课程全称聚焦于“Language Modeling from Scratch”也就是从零开始构建语言模型。课程最大的特色是“no black box”——不允许你把 Transformer 当成一个现成的黑盒来用而是要求你自己把每一个组件写出来。课程作业通常覆盖这样几个方向数据处理与 tokenizer 训练Transformer 模型架构实现训练循环与分布式训练指令微调与对齐每一部分都会给你一个“最小可运行的目标”然后要求你补全中间的核心逻辑。整个过程做下来你对 LLM 的认知会和只调 API 完全不一样。1.2 和 nanoGPT、HuggingFace 的关系很多人会把 CS336 和 Karpathy 的 nanoGPT 项目放在一起比较两者确实理念相似但侧重点不同项目侧重点适合人群nanoGPT用极简代码实现 GPT 训练代码本身是范本想快速看懂 GPT 训练流程的开发者CS336 作业要求你自己写代码要自己补全想通过动手彻底理解 LLM 的学生HuggingFace transformers开箱即用的工业级框架想快速落地应用、不想重复造轮子的人简单说nanoGPT 是“你看懂了别人怎么写”CS336 是“你必须自己写出来”。两者搭配学习效果最好。1.3 学完这门课你能获得什么如果你认真跟完 CS336 的作业你至少能掌握理解 tokenizer 为什么存在BPE 算法的训练和推理流程。理解 Transformer 中 Attention、LayerNorm、残差连接、Feed-Forward 这些模块的细节。理解语言模型的训练目标也就是“预测下一个 token”。理解训练中的采样、学习率调度、梯度裁剪的作用。能独立实现一个可训练的 GPT 模型而不是只会调用现成接口。这些能力恰恰是目前 LLM 应用开发中最稀缺的底层功。2. 课程内容与作业拆解CS336 的作业体系设计得很精巧。下面我们按模块拆解顺便说清楚每个模块背后的“为什么”。2.1 数据与 tokenizer文本怎么变成数字大模型输入的是文本但神经网络只能处理数字。所以第一步是设计一套规则把文本切分成 token再把 token 映射成整数 id。最简单的方案是字符级也就是每个字符当做一个 token。但字符级编码信息密度低一个词的语义被拆得很散。更常用的是子词方案比如 BPEByte Pair Encoding字节对编码。CS336 的作业会要求你自己实现 BPE 的训练过程而不是直接调tiktoken库。这个过程会让你理解为什么需要 BPE 而不是直接按空格分词词表大小如何影响模型效果和显存占用训练语料的清洗对最终效果有多大影响2.2 Transformer 模型实现这一部分是整个课程的核心。你需要自己写Token Embedding 和 Position EmbeddingMulti-Head Self-AttentionFeed-Forward NetworkLayerNorm 和残差连接最后的 LM Head 分类层作业会要求你从论文《Attention Is All You Need》以及 GPT-2 的架构出发逐步实现一个完整模型。这里最容易出错的地方是 Attention 的 mask 处理。语言模型只能看到当前 token 之前的内容不能看到未来内容所以需要引入因果掩码causal mask把未来位置的 attention 分数设为负无穷。2.3 训练与并行策略模型写好后就要开始训练。CS336 这一部分会涉及数据加载与 batch 构造损失函数计算优化器选择与学习率调度梯度累积、混合精度训练分布式数据并行DDP这些内容在实际训练大模型时非常重要。很多人会把模型写对但训练不稳定甚至 loss 不下降问题往往就出在这一环。2.4 指令微调与对齐最后一个阶段是“后训练”即让模型学会聊天、遵循指令而不是只做续写。这里会涉及 SFT监督微调以及 RLHF基于人类反馈的强化学习的基本思想。CS336 把这一部分也放进作业里是为了让你理解ChatGPT 类产品不是“训练出来”的而是“训练 对齐”两个阶段配合完成。虽然实际做 RLHF 的成本很高但理解它的流程对做应用开发非常有帮助。3. 环境准备3.1 软硬件建议先说明一下本文后面给出的实战代码是一个“迷你版 GPT”在你的个人电脑上就能跑通。操作系统Windows / macOS / Linux 都可以。Python建议 3.10 或 3.11。深度学习框架PyTorch 2.x。硬件CPU 就能运行小模型如果你有 NVIDIA GPU 且安装了 CUDA训练速度会更快。如果你的电脑没有 GPU完全不用焦虑。本文示例的模型参数量只有几百万CPU 训练几分钟就能完成一轮。3.2 创建虚拟环境建议用 conda 或 venv 创建一个独立环境避免依赖冲突conda create -n llm-from-scratch python3.11 conda activate llm-from-scratch pip install torch numpy如果安装 PyTorch 速度慢可以根据自己的操作系统到 PyTorch 官网选择对应的安装命令。这里只需要torch和numpy两个核心依赖。后面你会看到整个项目不依赖transformers库所有模型代码都是自己写的。这也是 CS336 课程的核心精神。4. 核心原理拆解4.1 LLM 工作原理预测下一个 token大语言模型做的最核心的一件事就是根据前文预测下一个 token。比如你输入“今天天气很”模型会计算词表中每个 token 的概率选择概率最高的“好”作为下一个 token。然后把这个 token 拼到输入后面继续预测下一个。这个过程叫自回归生成autoregressive generation。训练时的目标是一样的给定一段文本让模型学会预测每一个位置的“下一个 token”。用公式表示就是Loss -Σ log P(token_i | token_1, token_2, ..., token_{i-1})这个交叉熵损失会推动模型不断调整参数让预测越来越准确。4.2 BPE tokenizer 是怎么训练的BPE 是当前主流 LLM 最常用的分词算法。它的核心思想是从字符级开始不断合并出现频率最高的相邻 token 对直到达到目标词表大小。下面是一个简化版的 BPE 训练代码演示核心合并逻辑from collections import Counter def train_bpe(texts, num_merges100): # 1. 将文本拆成单词并统计词频 word_freq Counter() for text in texts: words text.strip().split() word_freq.update(words) # 2. 每个单词初始化为字符 token 列表 word_to_tokens {} for word in word_freq: word_to_tokens[word] list(word) # 3. 初始词汇表所有出现过的字符 vocab set(.join(word_freq.keys())) merges [] # 4. 迭代合并最频繁的相邻 pair for _ in range(num_merges): pair_freq Counter() for word, tokens in word_to_tokens.items(): freq word_freq[word] for i in range(len(tokens) - 1): pair (tokens[i], tokens[i 1]) pair_freq[pair] freq if not pair_freq: break # 选择出现次数最多的 pair best_pair pair_freq.most_common(1)[0][0] merges.append(best_pair) # 合并所有该 pair new_token best_pair[0] best_pair[1] vocab.add(new_token) for word, tokens in word_to_tokens.items(): new_tokens [] i 0 while i len(tokens): if i len(tokens) - 1 and (tokens[i], tokens[i 1]) best_pair: new_tokens.append(new_token) i 2 else: new_tokens.append(tokens[i]) i 1 word_to_tokens[word] new_tokens return vocab, merges训练完成后merges保存了合并顺序。对新文本编码时按同样的顺序依次合并即可def bpe_encode(text, merges): tokens list(text) for pair in merges: merged_token pair[0] pair[1] new_tokens [] i 0 while i len(tokens): if ( i len(tokens) - 1 and tokens[i] pair[0] and tokens[i 1] pair[1] ): new_tokens.append(merged_token) i 2 else: new_tokens.append(tokens[i]) i 1 tokens new_tokens return tokens注意真实生产环境中的 BPE 还要处理空格、UTF-8 字节、特殊 token 等问题远比上面的示例复杂。所以你会在 GPT 系列中看到Ġ和|endoftext|这种特殊 token。不过核心思想是一致的。4.3 Transformer 的前向计算GPT 系列使用的是 Transformer 的 Decoder 部分核心结构如下Token Embedding把 token id 映射成向量。Position Embedding给每个位置加上位置信息。多层 Transformer BlockLayerNorm Masked Multi-Head Self-Attention残差连接LayerNorm Feed-Forward Network残差连接最后的 LayerNorm 线性分类层。Self-Attention 的公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V其中 Q、K、V 分别来自输入向量的三个线性变换。在因果语言模型中QK^T矩阵的上三角位置要被掩码盖住确保位置 i 只能看到位置 0 到 i 的信息。Multi-Head 就是把 hidden_dim 拆成多个子空间每个头独立计算 attention最后拼接并通过一个线性层输出。理解了这些你就可以开始写代码了。5. 实战从零实现一个 TinyGPT这一部分我们动手写一个完整可训练的 GPT 模型。为了便于运行和调试我们使用一个小型英文语料采用字符级 tokenizer。字符级虽然简单但能完整展示“数据 → 模型 → 训练 → 生成”全链路。5.1 数据准备与 tokenizer文件路径data.py先准备一小段语料验证代码能跑通corpus The quick brown fox jumps over the lazy dog. Machine learning is a field of study. Language models predict the next token. Transformers use attention mechanisms. GPT models are based on the decoder architecture. # 字符级 tokenizer chars sorted(list(set(corpus))) vocab_size len(chars) char_to_id {c: i for i, c in enumerate(chars)} id_to_char {i: c for i, c in enumerate(chars)} def encode(text): return [char_to_id[c] for c in text] def decode(ids): return .join(id_to_char[i] for i in ids) # 将整个语料编码成 token id 序列 data encode(corpus)文件路径data_utils.py接下来构造训练样本。语言模型的输入是一段 token 序列监督标签是“往后移一位”的序列import torch def build_samples(token_ids, seq_len64, stride32): samples [] for i in range(0, len(token_ids) - seq_len - 1, stride): x token_ids[i:i seq_len] y token_ids[i 1:i seq_len 1] samples.append( ( torch.tensor(x, dtypetorch.long), torch.tensor(y, dtypetorch.long), ) ) return samples samples build_samples(data, seq_len64, stride32) print(f样本数量: {len(samples)}) print(f输入示例: {samples[0][0]}) print(f标签示例: {samples[0][1]})这里seq_len是输入序列长度stride是窗口滑动步长。使用步长而不用逐个位置滑动是为了让相邻样本之间有一定重叠同时控制总样本数。5.2 Transformer 模型代码文件路径model.py先实现 LayerNorm。GPT-2 使用的 LayerNorm 是在最后一个维度上做归一化这一点和很多图像模型不同import math import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, dim, eps1e-5): super().__init__() self.eps eps self.weight nn.Parameter(torch.ones(dim)) self.bias nn.Parameter(torch.zeros(dim)) def forward(self, x): mean x.mean(dim-1, keepdimTrue) var x.var(dim-1, keepdimTrue, unbiasedFalse) x_norm (x - mean) / torch.sqrt(var self.eps) return self.weight * x_norm self.bias再实现因果自注意力。这里把 Q、K、V 通过一个线性层一次算出来然后拆开这是 GPT-2 的常见做法class CausalSelfAttention(nn.Module): def __init__(self, hidden_dim, num_heads, dropout0.0): super().__init__() assert hidden_dim % num_heads 0 self.hidden_dim hidden_dim self.num_heads num_heads self.head_dim hidden_dim // num_heads self.c_attn nn.Linear(hidden_dim, 3 * hidden_dim, biasTrue) self.c_proj nn.Linear(hidden_dim, hidden_dim, biasTrue) self.dropout nn.Dropout(dropout) def forward(self, x): B, T, C x.shape qkv self.c_attn(x) q, k, v qkv.split(self.hidden_dim, dim2) k k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) q q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) att (q k.transpose(-2, -1)) / math.sqrt(self.head_dim) mask torch.tril(torch.ones(T, T, devicex.device)).view(1, 1, T, T) att att.masked_fill(mask 0, float(-inf)) att F.softmax(att, dim-1) att self.dropout(att) y att v y y.transpose(1, 2).contiguous().view(B, T, C) return self.dropout(self.c_proj(y))然后实现前馈网络和完整的 Transformer Blockclass MLP(nn.Module): def __init__(self, hidden_dim, expansion4, dropout0.0): super().__init__() self.c_fc nn.Linear(hidden_dim, expansion * hidden_dim) self.gelu nn.GELU(approximatetanh) self.c_proj nn.Linear(expansion * hidden_dim, hidden_dim) self.dropout nn.Dropout(dropout) def forward(self, x): x self.c_fc(x) x self.gelu(x) x self.c_proj(x) return self.dropout(x) class Block(nn.Module): def __init__(self, hidden_dim, num_heads, dropout0.0): super().__init__() self.ln_1 LayerNorm(hidden_dim) self.attn CausalSelfAttention(hidden_dim, num_heads, dropout) self.ln_2 LayerNorm(hidden_dim) self.mlp MLP(hidden_dim, dropoutdropout) def forward(self, x): x x self.attn(self.ln_1(x)) x x self.mlp(self.ln_2(x)) return x最后把整个模型组装起来class TinyGPT(nn.Module): def __init__( self, vocab_size, hidden_dim, num_heads, num_layers, max_seq_len, dropout0.0, ): super().__init__() self.token_embedding nn.Embedding(vocab_size, hidden_dim) self.position_embedding nn.Embedding(max_seq_len, hidden_dim) self.blocks nn.ModuleList( [Block(hidden_dim, num_heads, dropout) for _ in range(num_layers)] ) self.ln_f LayerNorm(hidden_dim) self.lm_head nn.Linear(hidden_dim, vocab_size, biasFalse) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean0.0, std0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean0.0, std0.02) def forward(self, idx, targetsNone): B, T idx.shape tok_emb self.token_embedding(idx) pos torch.arange(T, deviceidx.device) pos_emb self.position_embedding(pos).unsqueeze(0) x tok_emb pos_emb for block in self.blocks: x block(x) x self.ln_f(x) logits self.lm_head(x) loss None if targets is not None: B, T, C logits.shape loss F.cross_entropy( logits.view(B * T, C), targets.view(B * T), ) return logits, loss逻辑说明token_embedding负责把 token id 变成稠密向量。position_embedding给每个位置一个可学习的向量。每个Block内部都是“注意力 前馈”的组合并用残差连接包裹。lm_head把最后一个隐藏层向量映射回词表大小得到每个 token 的预测分数。apply(self._init_weights)是对所有子模块执行初始化规则是均值为 0、标准差为 0.02 的正态分布。5.3 训练循环文件路径train.py接下来写训练脚本。这里使用AdamW优化器加了一点权重衰减和梯度裁剪这是训练语言模型很常见的配置。import torch from model import TinyGPT from data_utils import samples def train(): hidden_dim 128 num_heads 4 num_layers 4 max_seq_len 64 batch_size 8 epochs 50 lr 3e-4 model TinyGPT( vocab_sizevocab_size, hidden_dimhidden_dim, num_headsnum_heads, num_layersnum_layers, max_seq_lenmax_seq_len, dropout0.1, ) optimizer torch.optim.AdamW( model.parameters(), lrlr, weight_decay0.1, ) for epoch in range(epochs): perm torch.randperm(len(samples)) total_loss 0 steps 0 for i in range(0, len(samples) - batch_size, batch_size): idx perm[i:i batch_size] x torch.stack([samples[j][0] for j in idx]) y torch.stack([samples[j][1] for j in idx]) logits, loss model(x, y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss loss.item() steps 1 print(fepoch {epoch 1}/{epochs}, loss {total_loss / steps:.4f}) torch.save(model.state_dict(), tinygpt.pt) print(训练完成模型已保存到 tinygpt.pt) if __name__ __main__: train()这里vocab_size需要从data.py中导入完整脚本建议把它们写在同一个文件或做成模块导入。为了演示方便你可以把corpus、chars、vocab_size等定义放在一个config.py中或者直接在train.py顶部定义。5.4 文本生成与效果验证文件路径generate.py训练完成后可以用一个简单的自回归生成函数来采样文本import torch import torch.nn.functional as F from model import TinyGPT torch.no_grad() def generate( model, encode, decode, prompt, max_new_tokens100, max_seq_len64, temperature0.8, top_k40, ): model.eval() ids encode(prompt) input_ids torch.tensor([ids], dtypetorch.long) for _ in range(max_new_tokens): x input_ids[:, -max_seq_len:] logits, _ model(x) # 只看最后一个位置的 logits next_logits logits[:, -1, :] / temperature # top-k 采样只保留概率最高的 k 个候选 if top_k is not None: v, _ torch.topk(next_logits, min(top_k, next_logits.size(-1))) next_logits[next_logits v[:, [-1]]] float(-inf) probs F.softmax(next_logits, dim-1) next_token torch.multinomial(probs, num_samples1) input_ids torch.cat([input_ids, next_token], dim1) gen_ids input_ids[0].tolist() return decode(gen_ids) if __name__ __main__: hidden_dim 128 num_heads 4 num_layers 4 max_seq_len 64 model TinyGPT( vocab_sizevocab_size, hidden_dimhidden_dim, num_headsnum_heads, num_layersnum_layers, max_seq_lenmax_seq_len, ) model.load_state_dict(torch.load(tinygpt.pt, map_locationcpu)) prompt GPT text generate(model, encode, decode, prompt, max_new_tokens80) print(text)这里有两个关键参数temperature控制随机性。temperature 越大输出越随机越小输出越保守。一般取 0.7 到 1.0 比较合适。top_k只从概率最高的 k 个 token 中采样可以避免模型输出一些很奇怪的 token。由于我们用的语料非常小模型输出可能会长得和训练语料很像甚至出现复读现象。这是小模型的正常表现不代表代码有 bug。你可以替换成更大的语料库比如下载一个 txt 格式的英文小说或者使用中文语料重新运行训练。如果换成中文语料字符级 tokenizer 也能工作但训练出的模型效果会比较有限因为中文字符太多、上下文依赖更复杂。做中文实验时建议使用 BPE 或按字切分并适当增大模型。6. 常见问题与排查思路在从零实现和训练 LLM 的过程中你大概率会遇到下面这些问题问题现象常见原因排查与解决思路训练的 loss 一直不下降学习率太大或太小尝试 lr3e-4 到 1e-3 区间观察 loss 波动情况loss 是 NaN 或 inf梯度爆炸 / 混合精度设置错误启用梯度裁剪降低学习率检查输入数据是否包含 NaNattention 能跑但效果差causal mask 写错模型偷看了未来信息检查 mask 是否为下三角矩阵可以手动打印 att 矩阵验证显存不够batch_size 太大 / 序列太长减小 batch_size使用梯度累积降低 max_seq_len生成文本全是重复温度太低 / top_k 太小 / 模型欠拟合适当调高 temperature减少 top_k 限制增加训练步数字符级 tokenizer 效果不好字符级信息密度太低切换为 BPE tokenizer重新训练模型embedding 维度对不上位置编码长度小于输入序列长度确保 max_seq_len 大于等于生成时的输入长度一个很实用的调试技巧先用一个很小的语料做“过拟合测试”。如果你的模型能在几十步内把训练集 loss 压到很低说明模型结构基本正确。如果连训练集都学不进去通常是 bug 在模型结构或数据处理而不是在“语料不够大”。7. 最佳实践与工程建议当你跑通上面的 TinyGPT 之后建议继续往这几个方向深入7.1 先用小模型验证再放大很多人一开始就想着训练一个 7B 甚至 70B 的模型结果卡在算力、显存和调试泥潭里。更合适的路径是用几 MB 的语料训练一个人物模型确认流程没有 bug。把 tokenizer 从字符级升级为 BPE。在更大的公开数据集上训练一个中等规模模型比如 100M 参数。再考虑分布式训练、混合精度等优化手段。7.2 训练稳定性的几个关键点梯度裁剪clip_grad_norm_设个 1.0能避免梯度爆炸。学习率预热前几百步用一个较小的学习率再逐步上升到目标值对训练稳定性很有帮助。固定随机种子在数据打乱和权重初始化时固定 seed方便复现和调试。定期保存 checkpoint建议每个 epoch 都保存一次避免训练中途崩掉后白跑。7.3 工程化建议数据质量 数据数量脏数据对模型的负面影响非常明显宁可少而精不要多而乱。记录实验日志用wandb或tensorboard记录 loss、学习率、显存占用方便对比实验。渐进式替换库从零实现是为了理解原理生产环境该用transformers、tiktoken、vLLM的时候不要犹豫。理解原理的人用框架才能用得明白。注意合法授权训练数据必须来源合法不要使用未经授权爬取的商业平台内容。8. 总结与后续学习建议CS336 课程最值得称道的地方是它把“大模型”从一个模糊的行业热词拆成一个个你可以在电脑上亲手实现的小任务。当你亲自写过 tokenizer、亲手调过 attention mask、亲眼看到 loss 随着训练下降你对 LLM 的理解会有质的变化。如果你刚接触这套内容建议从本文的 TinyGPT 开始把代码跑通然后对照 CS336 的作业要求把字符级 tokenizer 替换成 BPE再把模型维度扩大一点观察训练时间和效果的变化。下一步你可以继续看 Karpathy 的 nanoGPT 项目和 llm wiki 资源把论文《Attention Is All You Need》、GPT-2、GPT-3 的相关章节读一遍结合代码理解每个设计背后的动机。等你觉得这些都清楚了再去尝试微调、LoRA、RLHF会发现很多概念已经有了具象的认知基础。从零搭建 LLM 不是一条轻松的路但它是一条能让你真正“看懂”大模型的路。动手写起来吧哪怕只是先把 TinyGPT 跑出一个 hello world你也已经比别人更接近 LLM 的原理了。
返回列表