Stanford CS336：从零构建语言模型，6周带你写出自己的 LLM

Stanford CS336从零构建语言模型6周带你写出自己的 LLM2026 年春季斯坦福大学开设了一门全新的课程CS336: Language Modeling from Scratch由 Tatsunori Hashimoto 和 Percy Liang 联合授课。这门课的理念很直接——像操作系统课程让学生写一个完整 OS 一样CS336 要求学生从零构建一个完整的语言模型包括数据收集、模型架构、训练到部署的全流程。课程在 Hacker News 上获得了 340 分的热度社区反响强烈。本文带你深入了解这门课的设计思路和核心内容。课程定位把黑盒拆开大多数深度学习课程教学生如何使用预训练模型调参、微调、做推理。CS336 反过来——它假设你什么都不要从裸数据开始造轮子。这不是一门讲概念的课。课程页面明确标注了前置要求熟练掌握 Python代码量远超其他 AI 课熟悉 PyTorch 和深度学习系统优化内存层次结构、GPU 算子线性代数、概率论、机器学习基础而且这是5 学分的课。言下之意如果你的日常还有其他课要上恐怕时间不够。Assignment 1从零实现 Transformer第一个作业就把门槛拉满了Tokenizer实现 BPEByte Pair Encoding分词器处理原始文本数据模型架构实现完整的 TransformerMulti-Head Attention、Feed-Forward、LayerNorm、Positional Encoding优化器实现 AdamW 优化器训练循环在小规模数据集上完成训练这个阶段的目标是让学生理解 Transformer 的每一个组件如何工作。没有 HuggingFace没有 PyTorch Lightning没有现成的 Trainer——你写的就是框架。# 学生需要自己实现的 Attention简化示例classMultiHeadAttention(nn.Module):def__init__(self,d_model:int,n_heads:int):super().__init__()self.n_headsn_heads self.d_headd_model//n_heads self.w_qnn.Linear(d_model,d_model)self.w_knn.Linear(d_model,d_model)self.w_vnn.Linear(d_model,d_model)self.w_onn.Linear(d_model,d_model)defforward(self,x,maskNone):B,T,Cx.shape qself.w_q(x).view(B,T,self.n_heads,self.d_head).transpose(1,2)kself.w_k(x).view(B,T,self.n_heads,self.d_head).transpose(1,2)vself.w_v(x).view(B,T,self.n_heads,self.d_head).transpose(1,2)attnq k.transpose(-2,-1)/(self.d_head**0.5)ifmaskisnotNone:attnattn.masked_fill(mask0,float(-inf))attnF.softmax(attn,dim-1)out(attn v).transpose(1,2).contiguous().view(B,T,C)returnself.w_o(out)Assignment 2系统优化 — 自己写 FlashAttention第二个作业进入系统优化层面。在前一个作业的基础上Profiling用 PyTorch Profiler 分析和基准测试模型的每一层Triton FlashAttention2用 OpenAI Triton 语言自己实现 FlashAttention2 内核分布式训练实现支持多 GPU 并行训练的内存高效版本这是 CS336 的独特之处——大多数课程只讲怎么用模型但这门课让学生深入到 CUDA kernel 级别去理解 attention 的计算模式。# 学生用 Triton 实现的 FlashAttention2简化示意triton.jitdefflash_attn_fwd_kernel(q_ptr,k_ptr,v_ptr,o_ptr,stride_qh,stride_qt,stride_qd,stride_kh,stride_kt,stride_kd,stride_vh,stride_vt,stride_vd,stride_oh,stride_ot,stride_od,T,D:tl.constexpr,BLOCK_T:tl.constexpr,BLOCK_D:tl.constexpr,):# Triton kernel 实现 FlashAttention 的分块计算# 通过 tiling 避免完整注意力矩阵的内存开销...Assignment 3Scaling — 理解规模效应第三个作业聚焦于理解和验证 LLM 的 Scaling 规律在不同规模的模型Small / Medium / Large上训练验证 Chinchilla 法则在给定计算预算下模型参数和训练 token 的最优比例分析和可视化 Scaling 曲线这个阶段的实验直接验证了 Kaplan et al. (2020) 和 Hoffmann et al. (2022) 的 Scaling Law 论文。为什么这门课值得关注1. 填补了会用和会造之间的空白目前大部分开发者处于两个极端要么只会from transformers import AutoModel要么是工业界做预训练的大厂研究员。CS336 瞄准的是中间地带——让你理解 LLM 的内部机制达到可以自己训练小模型、调优训练流程的水平。2. 系统 ML 的交叉训练这不是一门纯 ML 课。它要求你写 Triton kernel、做分布式训练、分析和优化内存带宽。这种 ML Systems 的交叉能力正是目前 AI 人才市场上最稀缺的。3. 对抗框架黑盒化随着 HuggingFace 等工具链的成熟进入 AI 领域的门槛降低了但开发者也越来越不理解底层原理。CS336 的设计理念是知其所以然——即使最终你还是用 HuggingFace 和 PyTorch但你知道每一行代码背后对应的是什么。对中国开发者的启示CS336 的课程资源在 GitHub 上公开github.com/stanford-cs336这也是它上 Hacker News 榜首的原因之一。对于无法选修这门课的中国开发者来说完全可以按照课程大纲自学自学路线图阶段内容建议时间Phase 1实现 BPE Tokenizer DataLoader1 周Phase 2从零实现 TransformerAttention FFN LayerNorm2 周Phase 3训练循环 小规模验证1 周Phase 4用 Triton 实现 FlashAttention2 周Phase 5分布式训练DDP/FSDP1 周Phase 6Scaling Law 实验分析1 周如果你正在做 AI 应用开发不必完全复刻全部作业但至少理解 Transformer 的内部实现原理——这能在你调试模型行为、选择合适的预训练模型、甚至设计新的模型架构时提供底层直觉。总结CS336 不是一门轻松的课。它要求你写大量代码、理解 GPU 架构、分析系统性能。但正是这种从零开始的硬核方式才能培养出真正理解语言模型的工程师和研究者。课程链接cs336.stanford.edu | GitHub: github.com/stanford-cs336我正在做的在线工具站 zidongai.com.cn 也在持续关注 AI 工程化的最新进展欢迎交流。