斯坦福CS336:从零手搓大模型完整训练流程

发布时间:2026/9/7 2:26:03
斯坦福CS336:从零手搓大模型完整训练流程 这次我们来看一个不是“玩具项目”的项目斯坦福大学公开课 CS336《Building a Large Language Model from Scratch》的完整课程体系和配套代码仓库。它的核心目标很硬核——不借助成熟训练框架从原始文本数据开始自己实现分词器、Transformer 结构、训练循环、评估和推理最后真正“手搓”出一个能对话、能续写的大语言模型。很多同学学大模型调过 API用过 Hugging Face也跑过 Ollama但模型内部发生了什么一直是黑盒。CS336 的价值就在这里它把大模型从“组件拼装”打回“钢筋混凝土”让你亲手把每一层、每一个张量操作都写出来。课程配套了分阶段作业和测试代码基础任务以单 GPU 为最低门槛进阶实验才需要多卡集群所以个人开发者和学生也能跟着跑通。这篇文章我会拆解这套课程的学习路线、环境准备、动手部署流程、训练验证方法、资源占用观察方式以及最容易踩的坑。如果你已经会一点 Python 和深度学习基础但一直没搞清楚“模型到底是怎么训练出来的”这篇文章正好可以帮你建立完整的技术地图。1. CS336 核心能力速览能力项说明项目类型大模型教学课程 开源代码仓库课程来源斯坦福大学 CS336 公开课程核心技术栈Python、PyTorch 或 JAX核心任务数据获取、Tokenizer 实现、Transformer 架构、预训练、评估、推理全流程基础实验目标训练一个小规模 GPT 风格模型课程以单 GPU 为最低门槛设计进阶实验目标扩展到 1.3B 参数级别的多卡分布式训练启动方式克隆仓库后按作业顺序执行脚本配合作业说明逐步实现是否支持 API课程本身不提供 Web API训练产物可导出后接入 vLLM、Ollama 等推理服务是否支持批量任务支持批量数据加载与分布式训练推理阶段可批量处理文本适合人群有 Python 基础的开发者、算法工程师、想理解大模型底层原理的学习者这套课程不是教你“怎么用模型”而是教你怎么“造模型”。它的核心卖点可以归纳为四点全链路动手、单卡可入门、由浅入深、验证闭环。每个阶段都有明确产出不会让你学完只会看论文。2. 适用人群与学习边界先泼一盆冷水CS336 不是零基础入门课。你至少要熟悉 Python 语法、PyTorch 基本张量操作、神经网络前向反向传播的基本概念。如果你只会调 sklearn直接上手会非常痛苦。适合系统性学习的人包括这几类想从“调 API 工程师”变成“模型研发工程师”的开发者。准备大模型方向面试需要手推 Transformer、讲清楚 pre-training 细节的候选人。想给团队做内部技术分享需要一张完整大模型训练技术地图的组内讲师。本科或研究生阶段想通过实战理解 LLM 原理而不是只看论文的在校学生。不适合的人群也很明确只想快速部署一个聊天机器人建议直接去用 Ollama 或 vLLM不要在这门课上浪费时间。想直接给业务提供 API 服务的团队课程是学习项目不是生产级框架。没有 Python 基础和深度学习概念的同学建议先补 PyTorch 基础再回来学这门课。使用边界同样要说清楚。课程训练数据通常使用 RedPajama、TinyStories 等公开数据集但公开不等于可任意商用。国内用户下载和处理数据时要注意数据集的许可协议不要用来源不明的爬虫数据。训练出的模型如果用于商业发布一定要确认训练数据版权和生成内容的合规性。涉及敏感领域的文本生成更要在测试环境内验证并评估潜在风险。3. 环境准备与前置条件从我个人经验看CS336 的环境准备比想象中简单但有些坑需要提前规避。3.1 操作系统与硬件要求首选 Linux 环境。macOS 也可以跑但如果你计划训练超过 100M 参数的模型强烈建议使用 Linux 服务器或本地 Linux 工作站。Windows 用户建议安装 WSL2避免文件路径和 CUDA 兼容性问题。硬件方面最低门槛是单张 NVIDIA GPU。课程基础任务设计目标是单 GPU 可完成但实际显存占用取决于模型参数量、Batch Size、序列长度和优化器配置。更稳妥的做法是第一次实验先用小模型、小 Batch、短序列跑通流程再用nvidia-smi观察显存占用动态调整参数。磁盘空间建议预留 50GB 以上因为训练数据集、词汇表、检查点文件和日志都会占用不少空间。如果你只用少量样本跑通流程20GB 也够但建议不要过于压缩。3.2 Python 与依赖安装Python 版本建议 3.10 及以上。创建虚拟环境并安装依赖以下是一个通用安装流程# 进入课程代码目录后创建并激活虚拟环境 python -m venv .venv source .venv/bin/activate # 安装 PyTorch按你自己的 CUDA 版本调整命令 pip install torch --index-url https://download.pytorch.org/whl/cu121 # 安装课程依赖具体包名以仓库 requirements 为准 pip install -e .需要注意CS336 课程同时提供 PyTorch 和 JAX 两个版本实现不同作业阶段依赖的库可能不同。安装前一定先看仓库的 README 和作业说明避免装错版本。3.3 CUDA 与显卡驱动在训练前先确认 CUDA 可用python -c import torch; print(torch.cuda.is_available())如果输出False优先检查显卡驱动版本和 PyTorch 的 CUDA 版本是否匹配。常见错误是驱动太老或者 PyTorch 安装成了 CPU 版本。4. 课程内容结构与动手路线CS336 的课程设计很有层次大致分为四个阶段。每个阶段都有可提交的代码和可验证的实验结果这正是它比单纯看论文更适合自学的原因。4.1 阶段一数据准备与分词器第一步不是写模型而是处理数据。你需要下载一个文本数据集完成清洗、分片、采样等操作并实现或调用一个分词器把文本切成 token 序列。这个阶段你需要理解Token 与 Vocabulary 的关系。BPE 分词算法的合并规则。分词器词表大小对模型效果和显存的影响。数据分片在大规模训练中的作用。实操建议先使用小规模数据集比如 TinyStories 的几十 MB 子集把分词器训练出来然后用encode/decode往返验证from tokenizer import Tokenizer # 初始化并训练分词器 tokenizer Tokenizer() tokenizer.train(data/sample.txt, vocab_size8192) # 编码与解码往返验证 text CS336 手搓大模型 tokens tokenizer.encode(text) decoded tokenizer.decode(tokens) print(tokens) print(decoded) assert decoded text预期结果是 token ids 输出正确解码后能还原原始文本。如果失败优先检查特殊 token 占位和未知字符处理逻辑。4.2 阶段二实现 GPT 风格 Transformer 模型这一阶段是整个课程的核心。你需要自己实现 GPT-2 风格的模型结构包括Token Embedding 和 Positional Embedding。多头自注意力Multi-Head Self-Attention。LayerNorm 和残差连接。前馈网络MLP Block。最终 LM Head 的 logits 输出。建议按照“前向传播 - 损失计算 - 反向传播 - 梯度更新”的顺序逐模块写。不要一次性把整个模型写完再调试否则报错时定位会很痛苦。可以先用极小参数比如n_layer2, n_head4, n_embd64验证模型结构是否能跑通一次前向。import torch from model import GPT # 用极小配置验证前向传播 model GPT( vocab_size8192, n_layer2, n_head4, n_embd64, block_size128, ) x torch.randint(0, 8192, (2, 128)) logits model(x) print(logits.shape) # 预期 [2, 128, 8192]这一步跑通后你才算真正理解了 Attention 中的 Q、K、V 到底在做什么。4.3 阶段三训练循环与优化器模型有了接下来就是完整的训练流程。你需要实现DataLoader 批处理逻辑。AdamW 优化器。学习率调度。梯度裁剪。检查点保存与恢复。日志记录。如果从零手写 AdamW 会有点复杂课程并不要求你实现优化器内部才可继续但你需要理解优化器状态对显存的影响。基础版训练代码结构如下# 训练循环核心逻辑 import torch.optim as optim optimizer optim.AdamW(model.parameters(), lr3e-4, weight_decay0.1) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max1000) for step, batch in enumerate(train_loader): model.train() optimizer.zero_grad() logits, loss model(batch[input_ids], labelsbatch[labels]) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() if step % 100 0: print(fstep {step}, loss {loss.item():.4f}) torch.save(model.state_dict(), fcheckpoints/ckpt_{step}.pt)训练的关键指标是 loss 是否稳定下降。基础任务的目标通常是把 loss 压到 3.5 以下具体数值取决于数据集和分词器配置。如果 loss 长时间不动就要检查学习率、数据加载顺序、Attention Mask 是否正确。4.4 阶段四评估与推理训练完成后需要评估和推理在验证集上计算 Perplexity。根据输入前缀自动生成后续文本。检查多个 prompt 的输出质量。把检查点导出为 Hugging Face 格式方便后续接入推理服务。推理阶段的生成逻辑需要实现自回归采样逐 token 生成def generate(model, tokenizer, prompt, max_new_tokens50): model.eval() input_ids tokenizer.encode(prompt) input_ids torch.tensor([input_ids], dtypetorch.long) with torch.no_grad(): for _ in range(max_new_tokens): logits model(input_ids) next_token_logits logits[0, -1, :] next_token torch.argmax(next_token_logits, dim-1) input_ids torch.cat([input_ids, next_token.unsqueeze(0).unsqueeze(0)], dim-1) return tokenizer.decode(input_ids[0].tolist())通过对比不同生成策略比如贪心搜索、随机采样、Top-K你能直观感知到“同样权重不同采样策略会产生完全不同的输出风格”。5. 从零手搓 LLM 的实操步骤前面是理论路线我们把它拆成可以直接执行的步骤。以下流程是我按课程学习逻辑整理的标准操作顺序。5.1 第一步获取数据集并采样不要上来就下载整个 RedPajama 全量数据几百 GB 会让大多数个人开发者卡在第一关。建议只下载课程指定样本或使用 TinyStories 一类的轻量数据集。确认数据下载完成用文件行数统计确认样本规模wc -l data/train.txt head -5 data/train.txt5.2 第二步配置实验参数训练前先配置好超参数。下面是一份适合单卡小模型的配置参考# configs/small.yaml model: vocab_size: 8192 n_layer: 8 n_head: 8 n_embd: 512 block_size: 256 dropout: 0.0 training: batch_size: 8 grad_accumulation_steps: 4 learning_rate: 3e-4 max_steps: 5000 warmup_steps: 200 weight_decay: 0.1 gradient_clip: 1.0 log_interval: 50 eval_interval: 500 save_interval: 1000这份配置把训练步数控制在 5000 步以内显存压力可控。Batch Size 建议从 8 开始如果OOM就减半或者打开梯度累积或者缩短序列长度。5.3 第三步执行训练课程代码一般会提供训练入口脚本。执行前先确认当前设备nvidia-smi python train.py --config configs/small.yaml训练开始后观察日志里的 loss 和 learning ratestep 100 | loss 6.42 | lr 0.000120 | time 3.2s step 200 | loss 5.18 | lr 0.000150 | time 3.1s step 300 | loss 4.25 | lr 0.000180 | time 3.0s如果 loss 在几百步内从 10 快速下降说明模型结构和数据流基本正确。如果 loss 直接nan优先检查学习率是否过大以及分词器是否返回了越界的 token id。5.4 第四步测试生成效果训练到一定程度后手动生成几个案例。选择一个可控 prompt比如“今天天气很好我决定”观察模型是否输出了语法基本正确、上下文有逻辑的文本。小模型生成效果有限不要期望它像 GPT-4 一样流畅但至少要比“乱码”强。5.5 第五步保存并导出模型训练完成后的检查点一般只保存了 PyTorch 权重后续想接入推理服务还需要转换为通用格式。以导出 Hugging Face 格式为例from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM # 用课程训练好的权重构建 Hugging Face 模型具体转换方式按项目实现 config AutoConfig.from_pretrained(gpt2) model AutoModelForCausalLM.from_config(config) model.load_state_dict(torch.load(checkpoints/ckpt_5000.pt)) model.save_pretrained(./exported_model) tokenizer.save_pretrained(./exported_model)导出后这个模型就可以用标准的 Hugging Face 接口加载也可以继续转成 GGUF 格式接入 Ollama 做本地私有化推理。6. 接口 API 与批量推理CS336 课程本身不是一个“API 服务”项目它要做的是训练和验证一个模型。但训练完成后模型要落地使用就需要接入推理框架。6.1 使用 vLLM 提供 API 服务如果你的显卡显存足够训练出的模型可以转成 Hugging Face 格式再用 vLLM 提供 OpenAI 风格的 API 服务。vLLM 支持高并发请求适合批量测试和内部工具集成python -m vllm.entrypoints.openai.api_server \ --model ./exported_model \ --tensor-parallel-size 1 \ --host 127.0.0.1 \ --port 8000启动后用 curl 验证curl http://127.0.0.1:8000/v1/completions \ -H Content-Type: application/json \ -d {model: ./exported_model, prompt: 一次有趣的旅行, max_tokens: 100}当然这种方式比较适合显存充足的场景。更常见的是只跑通课程基础任务导出的模型比较小可以直接在 Python 脚本里批量推理。6.2 批量推理脚本以下是一个通用批量推理模板把多个提示词放到prompts列表逐个生成并保存结果import json from model import GPT from tokenizer import Tokenizer model GPT.from_pretrained(exported_model) tokenizer Tokenizer.from_pretrained(exported_model) prompts [ 请用一句话介绍人工智能, 写一个关于电脑的短故事, 11等于几, ] results [] for prompt in prompts: output model.generate(prompt, max_new_tokens128) results.append({prompt: prompt, output: output}) print(fPrompt: {prompt}\nOutput: {output}\n) with open(outputs/generation_results.json, w, encodingutf-8) as f: json.dump(results, f, ensure_asciiFalse, indent2)批量推理要注意两个问题一是设定合理的max_new_tokens避免长文本生成耗尽显存二是增加超时检查防止某个 prompt 触发死循环或异常长输出。7. 资源占用与性能观察模型训练和推理最关心的永远是显存。虽然没有一套固定数据但可以给出清晰的观察方法和调优思路。7.1 显存占用怎么看训练时在另一个终端运行watch -n 1 nvidia-smi重点看进程的Memory-Usage和Volatile GPU-Util。显存占用由模型权重、梯度、优化器状态和激活值共同决定。同样的参数量批量越大显存占用越高序列长度越长相关系数越大。7.2 为什么同样的模型别人能跑你跑不动最常见的原因不是显卡不好而是你直接把参考配置里的batch_size和block_size抄过来。解决方案按优先级排列缩小batch_size。减小block_size把 256 改成 128 或 64。打开梯度累积用小 batch 模拟大 batch。使用混合精度训练PyTorch 环境可以用torch.autocast。清空缓存torch.cuda.empty_cache()。7.3 CPU 和 GPU 的差异课程基础任务用 GPU 训练是标准路径但如果你只想验证代码逻辑在 CPU 上跑通前向和反向也不是不行只是速度会慢很多。建议调试阶段用 CPU正式训练切换到 GPU。7.4 训练日志记录不要只靠print输出 loss。建议把训练步数、显存、token/s、loss、learning rate 都记录到结构化日志里方便后期复盘import json import torch import time log_file open(logs/train_log.jsonl, a) for step, batch in enumerate(train_loader): start_time time.time() loss train_step(batch) elapsed time.time() - start_time gpu_mem torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0 log_entry { step: step, loss: round(loss.item(), 4), gpu_mem_gb: round(gpu_mem, 2), time_per_step: round(elapsed, 3), } log_file.write(json.dumps(log_entry) \n) log_file.flush()观察time_per_step的变化可以判断数据加载是不是瓶颈。设置优化显存的优先级永远是先跑通、再调优。8. 常见问题与排查方法手搓 LLM 过程中会遇到大量问题排查思路比具体报错更重要。这里整理一份高频问题清单。问题现象可能原因排查方式解决方案安装依赖失败Python 版本不匹配或 CUDA 版本不兼容查看报错信息、用python --version确认版本更换 Python 版本按 README 安装对应版本依赖数据加载特别慢数据集过大或读取方式低效观察训练日志中time_per_step使用小规模采样、改进 DataLoader 并行读取Tokenizer 编解码不一致特殊 token 处理异常单条文本往返测试检查 PAD/UNK/BOS/EOS token 的添加和去除逻辑模型前向报维度错误Embedding 或 Attention 维度不匹配打印每层张量 shape按模块逐层调试Loss 不下降学习率太小或数据流错误查看日志 lr 和 loss、检查数据批处理调大学习率检查标签是否错位Loss 直接变 NaN学习率过大、梯度爆炸或数据异常检查数据是否有缺失值降低学习率、梯度裁剪、检查输入 token 是否越界显存不足 OOMBatch Size 或序列长度过大查看 nvidia-smi 显存占用减小 Batch、缩短序列、梯度累积、混合精度GPU 利用率低数据加载瓶颈或批大小太小观察Volatile GPU-Util加大 Batch、使用异步数据加载模型生成全重复文本采样策略问题、训练不充分、模型容量过小对比不同采样策略尝试 Top-K/Top-P 采样、调高温度、增大训练步数检查点无法加载模型结构不匹配或路径错误打印 state_dict 的 key确认配置一致、路径完整9. 最佳实践与使用建议手搓 LLM 是一次很好的工程训练但很多人在中途会放弃。工程化实践我建议做好这几件事第一先小后大。第一次实验不要追求“训练一个 ChatBot”而是先用 65M 左右的小模型在小数据集上跑通整条链路。全流程跑通后再扩大数据、加深模型这才是最快的速度。第二日志比清屏重要。不要只在终端看 loss单独开一个日志文件把训练步数、loss、显存、耗时全部记录方便复盘。如果训练中断可以基于检查点恢复继续训练。第三模块化设计。把数据处理、模型定义、训练循环、评估脚本、推理脚本分开目录管理方便回溯也更接近真实项目结构。第四合规使用。训练数据要选择明确许可的数据集不采集未经授权的个人隐私数据。生成内容要经过人工复核再对外发布。如果模型将用于商业场景务必对训练数据的版权链条做一次完整梳理。第五部署前做最小验证。用导出后的模型先跑一遍推理脚本确认输出效果可以接受再接入 API 服务。避免把训练一半的中间检查点直接丢到生产环境。10. 总结与下一步CS336 最值得尝试的地方在于它给你划定了一条从零实现 LLM 的完整路径数据、分词、模型、训练、评估、推理一个不落。最先验证的功能是模型前向传播和 loss 下降完全跑通基础任务再考虑并行训练和模型导出。最容易踩的坑在环境阶段CUDA 版本不匹配、数据下载不完整、Tokenizer 编解码不一致。这些问题的共性在于前期准备不充分多花一点时间确认依赖和数据格式能省下大量调试时间。后续可以继续扩展的方向包括把训练好的模型导出并接入 Ollama搭建本地私有大模型服务尝试更长的序列和更大的模型对比不同规模下的训练曲线结合 LoRA 做领域数据微调让模型具备特定业务能力或者把课程作业整理成自己的开源项目作为技术能力展示。建议先收藏这篇文章准备好环境以后按步骤跑一遍。手搓大模型并不玄学真正动手之后你会发现它比看一百篇论文都更能建立你对大模型的直觉。