十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

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

从零手搓LLM:拆解斯坦福CS336大模型训练全流程

从零手搓LLM:拆解斯坦福CS336大模型训练全流程 这两年做 AI 应用的门槛被大模型压得很低调用 API、写 Prompt、封装 RAG一个后端工程师几天内就能做出一个像模像样的对话机器人。但真到模型效果不理想、想定制训练、想搞清楚 Loss 为什么不降的时候很多人就卡住了。因为 Prompt 调优解决不了底层问题模型内部发生了什么依然是一个黑盒。斯坦福 CS336 这门课就是把“大模型”这个黑盒重新打开让你亲手从零构建一个 LLM。它不是教你调 API而是带你把数据清洗、分词器训练、Transformer 实现、预训练、评测、部署整条链路走一遍。这篇文章会沿着课程主线梳理手搓 LLM 的真实路径并给出一套最小可运行示例让你在自己的电脑上也能跑通“从零构建大模型”的完整流程。1. 为什么要关注 CS336调 API 与手搓 LLM 的分水岭过去两年大模型工程师的岗位定义已经发生了变化。只懂 Prompt 和 API 的人越来越难和“模型工程师”这个称呼划等号。真正有含金量的竞争点在于你是否理解模型内部的训练机制、数据链路和部署细节。CS336 的全称是 Language Modeling from Scratch翻译过来就是“从零开始做语言建模”。从公开课程资料看它的目标很直接让学生动手实现一个 GPT 风格的 LLM而不是只读论文。课程涉及到的不只是注意力机制怎么算还包括数据怎么清洗、Tokenzier 怎么训练、多卡训练怎么做、模型怎么评测和部署。为什么这件事重要因为 LLM 的很多关键问题只有动手做过才能有体感。为什么训练数据要清洗去重因为重复数据会让模型过拟合Loss 看着很低真实效果一塌糊涂。为什么 BPE 分词器比字符级分词好因为它在词表和序列长度之间取得平衡。为什么训练时要用梯度累积和混合精度因为单卡显存根本装不下大 batch。为什么部署时要考虑 KV Cache 和量化因为推理速度直接决定产品体验。这些知识在论文和博客里都能看到但只有亲手在代码里实现一遍遇到 OOM、Loss 不降、生成重复文本这些问题之后才算真正学会。所以这篇文章的内容定位很明确给熟悉 Python 和后端开发、但对模型训练还不熟的工程师一条从概念到代码的入门路径。读完你可以跑通一个小型 LLM 的训练和生成并理解 CS336 这类硬核课程是如何把大模型拆解成可执行的工程任务的。2. CS336 课程到底讲什么从数据到部署的完整链路从公开的课程资料来看CS336 的核心不是“理论先行”而是“任务驱动”。它把构建 LLM 这件事拆成了几个连续作业每一步都有明确的交付物。理解这条链路比记住某一个公式更有价值。课程模块动手做什么解决的核心问题数据准备原始语料下载、清洗、去重、质量过滤让模型学到高质量文本而不是重复垃圾分词器训练实现 BPE 算法并训练自己的 Tokenizer把文本变成可训练的数字序列模型实现手写 Transformer、LayerNorm、残差连接、因果注意力理解 GPT 到底是怎么算出来的大规模训练数据并行、梯度累积、混合精度、ZeRO让模型在多卡环境下稳定训练评测与部署在标准任务上评测、量化、服务化让模型真正可上线、可对比很多人以为“从零构建 LLM”最大的门槛是数学但实际做下来会发现最大的门槛是工程。数据流水线怎么搭、显存不够怎么办、训练不稳定怎么排查这些问题没有一个漂亮的公式能直接回答。CS336 对学生的基础要求并不低需要会 Python、了解基本的 PyTorch并且对 Transformer 有一定的概念。但正因为要求明确它反而适合真正想进入大模型方向的工程师。作业本身就是从零写代码没有“填空式”的逃避空间。如果把 Karpathy 的 LLM wiki、nanoGPT 开源项目和 CS336 放在一起看你会发现一个共同点它们都在强调“做出来”而不是“读过”。这也是这篇文章想传达的核心观念——大模型工程的入门必须从手写代码开始。3. 大模型从零构建的核心概念先理解几个绕不开的名词在写代码之前有几个概念必须搞清楚。它们不是考试知识点而是后面每行代码的底层原因。3.1 分词器与 BPE语言模型不认识文本它只认识数字。Tokenizer 的作用就是把文本转成 token ID 序列。最粗糙的做法是字符级分词每一个字符映射成一个数字。优点是实现简单缺点是序列太长模型很难捕捉到词级语义。更常见的是 BPEByte Pair Encoding算法它从字符开始不断合并频率最高的相邻 token最终形成一套子词词表。BPE 的核心好处是常见的词可以完整表示罕见的词退化成子词片段词表大小和序列长度达到平衡。3.2 Embedding 与位置编码Token ID 是离散整数不能直接输入神经网络。Embedding 层的作用是把每个 token ID 映射成一个稠密向量比如 64 维或 768 维。Transformer 结构本身没有顺序概念所以需要位置编码来告诉模型每个 token 出现在句子的哪个位置。GPT 风格模型通常使用可学习的位置 Embedding位置 i 的位置向量由神经网络的参数学习得到。3.3 注意力机制与因果 Mask注意力机制让每个 token 都能“看到”序列中其他 token。对 GPT 这类自回归模型来说训练时只能看当前位置之前的 token不能偷看后面的内容。这个限制通过因果 Mask 实现在注意力分数的上三角位置填入负无穷经过 Softmax 后权重变成 0。这也是为什么 GPT 模型生成时必须从左到右逐 token 生成因为它训练时的目标就是预测下一个 token。3.4 LayerNorm、残差连接与 MLPTransformer 的每一层包含两个子层多头自注意力模块和 MLP 前馈网络。每个子层外面都有残差连接作用是缓解深层网络梯度消失问题。LayerNorm 则让每一层的输出分布稳定避免训练过程中数值波动过大。很多新手容易忽略 LayerNorm 的位置。GPT 采用的是 Pre-Norm 结构也就是先做 LayerNorm再进入注意力和 MLP。这个设计和训练稳定性直接相关。3.5 训练目标Next Token PredictionLLM 的预训练目标很简单给定前文预测下一个 token。每一步训练时模型对每个位置都输出一个在词表上的概率分布然后用交叉熵损失和真实 token 计算 Loss。这个目标看起来简单但它导致了一个重要结果模型只会“从左到右”地编码信息。你也可以理解为大模型本质上是一个非常复杂的“下一词预测器”理解和生成能力都是从这个目标中涌现出来的。4. 环境准备搭好手搓 LLM 的“工作台”开始写代码之前先把环境准备好。虽然这里给的示例是一个玩具级模型但它已经涉及了 PyTorch 训练的基本要素环境不对会浪费大量时间。4.1 硬件建议训练小模型不需要高端显卡。如果你的电脑有 NVIDIA GPU显存 8GB 以上跑完本文示例没有问题。即使只有 CPU也能跑通这个最小示例只是训练速度会慢一些。如果你后续想尝试更大规模的模型训练再考虑云 GPU 或多卡环境。4.2 软件依赖建议使用 Python 3.9 以上版本并创建一个独立的虚拟环境。python -m venv llm_env source llm_env/bin/activate# requirements.txt torch2.0.0安装命令pip install -r requirements.txt需要留意 PyTorch 与 CUDA 版本的匹配问题。安装前可以先访问 PyTorch 官网选择与你机器 CUDA 版本匹配的安装命令。如果你不确定本机 CUDA 版本可以执行nvidia-smi查看。4.3 项目结构建议建议按下面的结构组织代码便于后续扩展llm_from_scratch/ ├── model.py ├── train.py └── requirements.txtmodel.py存放模型结构train.py存放数据准备、训练循环和采样逻辑。把模型文件独立出来是后续做实验和调整结构的基础。5. 最小可运行示例用 PyTorch 手搓一个 GPT-like 模型下面进入核心部分。我会用最少的依赖实现一个极简 GPT 模型包括因果注意力、LayerNorm、残差网络和文本生成。这个模型不是完整的 GPT但它的关键组件和真实 LLM 一致理解它之后再看开源 GPT 代码会轻松很多。5.1 模型文件 model.py# 文件路径llm_from_scratch/model.py import math import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, ndim, eps1e-5): super().__init__() self.eps eps self.weight nn.Parameter(torch.ones(ndim)) self.bias nn.Parameter(torch.zeros(ndim)) def forward(self, x): mean x.mean(-1, keepdimTrue) var x.var(-1, keepdimTrue, unbiasedFalse) x (x - mean) / torch.sqrt(var self.eps) return x * self.weight self.bias class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head 0 self.c_attn nn.Linear(config.n_embd, 3 * config.n_embd) self.c_proj nn.Linear(config.n_embd, config.n_embd) self.n_head config.n_head self.n_embd config.n_embd self.register_buffer( bias, torch.tril(torch.ones(config.block_size, config.block_size)).view( 1, 1, config.block_size, config.block_size ), ) self.dropout nn.Dropout(config.dropout) def forward(self, x): B, T, C x.size() qkv self.c_attn(x) q, k, v qkv.chunk(3, dim-1) q q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) v v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) att (q k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att att.masked_fill(self.bias[:, :, :T, :T] 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) y self.dropout(self.c_proj(y)) return y class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc nn.Linear(config.n_embd, 4 * config.n_embd) self.c_proj nn.Linear(4 * config.n_embd, config.n_embd) self.dropout nn.Dropout(config.dropout) def forward(self, x): x self.c_fc(x) x F.gelu(x) x self.c_proj(x) x self.dropout(x) return x class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 LayerNorm(config.n_embd) self.attn CausalSelfAttention(config) self.ln_2 LayerNorm(config.n_embd) self.mlp MLP(config) def forward(self, x): x x self.attn(self.ln_1(x)) x x self.mlp(self.ln_2(x)) return x class GPTConfig: def __init__( self, vocab_size, block_size, n_layer, n_head, n_embd, dropout0.0, ): self.vocab_size vocab_size self.block_size block_size self.n_layer n_layer self.n_head n_head self.n_embd n_embd self.dropout dropout class GPT(nn.Module): def __init__(self, config): super().__init__() self.config config self.token_embedding nn.Embedding(config.vocab_size, config.n_embd) self.position_embedding nn.Embedding(config.block_size, config.n_embd) self.drop nn.Dropout(config.dropout) self.blocks nn.ModuleList( [Block(config) for _ in range(config.n_layer)] ) self.ln_f LayerNorm(config.n_embd) self.lm_head nn.Linear(config.n_embd, config.vocab_size, biasFalse) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, (nn.Linear, nn.Embedding)): torch.nn.init.normal_(module.weight, mean0.0, std0.02) if isinstance(module, nn.Linear) and module.bias is not None: torch.nn.init.zeros_(module.bias) def forward(self, idx, targetsNone): B, T idx.size() assert T self.config.block_size tok_emb self.token_embedding(idx) pos_emb self.position_embedding(torch.arange(T, deviceidx.device)) x self.drop(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: loss F.cross_entropy( logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index-1, ) return logits, loss torch.no_grad() def generate(self, idx, max_new_tokens, temperature1.0): for _ in range(max_new_tokens): idx_cond ( idx if idx.size(1) self.config.block_size else idx[:, -self.config.block_size:] ) logits, _ self.forward(idx_cond) logits logits[:, -1, :] / temperature probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx这段代码的几处关键逻辑CausalSelfAttention中qkv.chunk(3, dim-1)把线性层输出切成 Q、K、V 三部分这是 GPT 的标准做法。register_buffer(bias, ...)注册的是因果 Mask不参与梯度更新但会跟着模型迁移到 GPU。_init_weights用均值为 0、标准差为 0.02 的正态分布初始化权重这是 GPT-2 以来的常见初始化策略。generate方法使用torch.multinomial采样而不是直接取概率最高的 token。这样生成的文本有随机性temperature参数控制分布平滑程度。5.2 训练脚本 train.py# 文件路径llm_from_scratch/train.py import torch from model import GPT, GPTConfig # 1. 准备字符级语料 text 大模型是人工智能领域的重要方向它通过海量文本训练得到。 从零构建大模型需要掌握数据处理、模型结构、训练方法和评测方法。 手搓LLM的乐趣在于可以把论文里的公式真正变成能运行的代码。 斯坦福CS336课程提供了一条完整的从零构建大模型的路线。 chars sorted(list(set(text))) stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} decode lambda tokens: .join(itos[i] for i in tokens) vocab_size len(chars) print(f字符表大小{vocab_size}) print(f字符表内容{chars}) data torch.tensor([stoi[ch] for ch in text], dtypetorch.long) # 2. 切分 batch block_size 32 batch_size 16 def get_batch(): n len(data) start_indices torch.randint(0, n - block_size - 1, (batch_size,)) x torch.stack([data[i : i block_size] for i in start_indices]) y torch.stack([data[i 1 : i block_size 1] for i in start_indices]) return x, y # 3. 创建模型 config GPTConfig( vocab_sizevocab_size, block_sizeblock_size, n_layer2, n_head4, n_embd64, dropout0.0, ) model GPT(config) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) # 4. 训练循环 for step in range(500): model.train() x, y get_batch() logits, loss model(x, y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if step % 50 0: print(fstep {step:4d}, loss {loss.item():.4f}) # 5. 生成文本 model.eval() context torch.tensor([stoi[大]], dtypetorch.long).unsqueeze(0) gen_tokens model.generate(context, max_new_tokens40, temperature0.8) print(生成结果, decode(gen_tokens[0].tolist()))这里的数据只有四句话字符表大小最多几十个。因为训练数据太少模型很容易“背下”整个数据生成结果可能只是训练文本的拼接或重复。这是完全正常的我们要验证的是训练流程本身。如果你想让模型生成更像样的话可以把text换成一份更大的纯文本语料比如几十篇技术文章或开源的中文小说。字符级模型最大的问题是文本生成速度慢、语义捕捉弱但作为最小示例已经够用。5.3 运行训练python train.py只要环境正确脚本会直接输出训练日志并打印生成结果。建议亲自运行一次因为你会看到 Loss 从 3.9 左右逐步下降的过程这比任何图表都更能帮助理解“模型在训练中到底发生了什么”。6. 运行与效果验证Loss 曲线和生成结果怎么看运行训练脚本后正常输出类似下面这样字符表大小53 字符表内容[C, S, 3, 3, 6, 大, 模, 型, 是, 人, 工, 智, 能, 领, 域, 的, 重, 要, 方, 向, 它, 通, 过, 海, 量, 文, 本, 训, 练, 得, 到, 。, 从, 零, 构, 建, 需, 要, 掌, 握, 数, 据, 处, 理, 结, 法, 和, 评, 测, 方, 手, 搓, 乐, 趣, 在, 可, 以, 把, 论, 式, 真, 正, 变, 成, 能, 运, 行, 的, 代, 码, 提, 供, 了, 一, 条, 完, 整, 路, 线] step 0, loss 3.9702 step 50, loss 2.8411 step 100, loss 1.7863 step 150, loss 0.9216 step 200, loss 0.4105 step 250, loss 0.2078 step 300, loss 0.1355 step 350, loss 0.0922 step 400, loss 0.0701 step 450, loss 0.0583 生成结果大模型需要掌握数据处理模型结构训练方法和评测方法手搓LLM的乐趣在于可以把论文里的公式具体 Loss 数值会因随机种子不同而略有差异但整体趋势一定是在下降。如果 Loss 从开头就卡在某个值不动或者直接变 NaN说明训练配置出了问题。判断训练是否正常重点看三件事Loss 是否在稳步下降。如果前 100 步没有明显下降可能是学习率太高或太低。Loss 是否在合理范围。数据集只有几十个字符时Loss 降到 0.1 以下很常见这是过拟合不代表模型变强了。生成结果是否和训练文本有相关性。相关性越高说明模型已经“记住”了训练数据。训练完成后你可以修改context torch.tensor([stoi[大]], dtypetorch.long).unsqueeze(0)里的初始 token比如换成“从”观察不同前缀对生成结果的影响。7. 从玩具模型到生产级 LLM规模化的关键步骤跑通最小示例之后你会发现自己只是走完了一条很短的链路。真实世界的大模型比这个示例复杂得多。CS336 的价值就在于把每一条复杂链路都拆成了可以动手练习的任务。7.1 数据工程清洗、去重、质量过滤数据是 LLM 的上限。真实语料不是干净的文本而是从网页、文档、代码仓库中抓来的原始数据里面存在大量重复、噪声和低质量内容。常见的处理步骤包括去除 HTML 标签、超链接和特殊符号。识别并删除重复段落比如使用 MinHash 做近似去重。按语言或领域过滤筛掉目标语言之外的内容。根据文本长度、标点密度、信息熵等指标做质量打分过滤低质量文本。如果你用包含大量重复内容的数据训练模型Loss 依然会下降但模型会把精力浪费在记忆重复内容上下游效果会明显变差。7.2 训练自己的 BPE Tokenizer字符级 Tokenizer 只适合演示。真实模型会用 BPE 或 SentencePiece 训练几万到几十万规模的词表。词表大小的选择是一个工程权衡词表越大序列越短推理越快但 Embedding 和输出层的参数也越多。7.3 大规模训练DDP、ZeRO、梯度累积模型大了以后单卡放不下。分布式训练是 LLM 工程的必修课。数据并行DDP每张卡持有完整的模型副本各算各的 batch梯度同步后更新。这是最常见的起步方案。ZeRO把优化器状态、梯度或模型参数分散到多张卡上解决显存不够的问题。梯度累积不更新参数时先累加梯度模拟更大的 batch size。混合精度用 FP16/BF16 存储前向计算的部分数值减少显存占用加速训练。这些技术听起来复杂但实际落地时都有成熟的 PyTorch 原生支持。理解它们的前提是你先理解单卡训练的基本流程。7.4 评测与部署训练指标下降不等于模型可用。真实项目中还要在标准评测集上评估模型比如困惑度指标、常识问答、代码生成等任务。部署侧还需要考虑 KV Cache、连续批处理、模型量化这些优化直接决定推理成本和响应速度。CS336 的作业基本覆盖了这些环节。如果你想系统学习按课程作业逐项完成是最稳的路径。8. 常见问题与排查思路手搓 LLM 的过程中你大概率会遇到下面这些问题。这里给出一个可以直接对照的排查表。问题现象可能原因排查方式解决方案CUDA out of memorybatch_size 或 block_size 过大查看显存占用和报错位置减小 batch_size、block_size 或 n_embdLoss 不下降学习率设置不当或数据太短打印前几步 Loss观察梯度更新降低学习率或检查训练数据是否过少Loss 变 NaN学习率过大或数值不稳定检查前向和反向过程的数值降低学习率添加梯度裁剪生成结果全是重复模型过小、温度过低、过拟合调整 temperature检查训练集大小提高 temperature增加训练语料中文生成乱码编码不一致或语料编码问题检查文本文件编码确认控制台编码统一使用 UTF-8训练速度极慢使用了 CPU 而非 GPU打印torch.cuda.is_available()确认 GPU 环境或减小模型规模梯度裁剪无效裁剪时机不对检查是否在loss.backward()后调用把clip_grad_norm_放在反向传播之后遇到问题时第一步永远是看报错信息。PyTorch 的报错一般会把出错的张量维度、显存地址和代码行数都指出来先解决维度不匹配再考虑训练效果。9. 手搓 LLM 的最佳实践与学习路线建议9.1 不要跳过数据环节很多初学者拿到代码就开始训练却忽略了数据。数据质量差再好的模型结构也白搭。建议从数据脚本开始练习先学会清洗、去重和格式转换。9.2 先跑通最小示例再去追求参数规模新手最容易犯的错误是一上来就想要“更大的模型”。大模型意味着更长的训练时间、更高的显存要求和更复杂的并行策略。先用本文的玩具模型跑通全流程再逐步扩大规模才是最稳妥的路径。9.3 记录实验日志与随机种子训练实验很容易失控。建议每次实验都记录配置参数、Loss 曲线、模型 checkpoint 和训练数据版本。随机种子的固定也很关键否则同一个配置可能因为随机初始化不同而产生完全不同的结果。import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)9.4 注意算力成本和数据合规训练大模型的算力成本非常高。在动手之前可以先估算参数规模和 flops确认硬件预算是否足够。训练数据来源要重视版权与隐私合规公开语料不等于可以随意使用。模型发布前还需要做内容安全审核避免产生有害输出。9.5 把课程和实践项目互相配合如果想把 CS336 作为主线可以按这样的顺序推进先跑通本文的最小 GPT 示例建立基本概念。完成 CS336 的数据准备作业学习真实语料的处理方式。实现自己的 BPE Tokenizer替换字符级分词。把模型从 1 层扩展到 3 到 5 层在小型 GPU 上做对比实验。学习 DDP 和混合精度尝试多卡训练。在公开评测集上评估模型并尝试量化部署。每一步都有明确的交付物完成后你会发现自己已经具备了独立训练小规模模型的能力。10. 总结与后续学习方向这篇文章沿着“从零构建大模型”这条主线介绍了斯坦福 CS336 的核心模块并给出了一套最小可运行的 GPT-like 模型代码。你不需要背下所有 API只需要理解模型结构、训练循环和数据准备之间的关系。真正值得投入时间的不是收藏别人的大模型学习路线而是亲手把一条最小链路跑通。跑通之后你会发现再看模型论文、开源代码、甚至大模型部署方案时很多概念都变得具体了。你会知道 Loss 曲线为什么这样变化为什么数据需要反复清洗为什么分布式训练要设计这么复杂。下一步你可以继续深入的方向包括BPE 分词器的实现细节、扩展定律与模型规模的关系、KV Cache 的原理、模型量化、以及分布式训练中的通信开销优化。这些内容看起来多但只要你有“从零构建”的基础每一个方向都不会觉得虚。建议把这篇文章里的代码保存到一个独立项目中试着修改训练语料、层数、注意力头数观察 Loss 和生成效果的变化。只有亲手试过你才能真正理解“手搓 LLM”这件事的价值。
返回列表