企业官网建设流程全解析

如果你早就能用 Hugging Face 加载 GPT、用 DeepSpeed 跑微调，却仍然说不清 self-attention 里的 Q、K、V 为什么是三个矩阵，也说不出一次完整的预训练数据应该长什么样，那说明你和大模型之间还隔着一层“框架窗帘”。斯坦福 CS336 正是用来掀开这层窗帘的课程：它把大模型从数据到推理的每一条链路都拆开，让你用几百行代码从零写出一个能自己生成句子的语言模型。这篇文章不会帮你把课程答案抄一遍，而是给你一条可以照做的学习路径：这门课解决了什么问题、核心环节怎么拆解、代码怎么组织、训练跑起来后怎么判断成功、以及最容易卡住新手的坑在哪里。

很多同学学大模型的路径是这样的：先玩 API，再学 Hugging Face，然后尝试微调。结果微调一次都没成功，因为根本不知道 loss 为什么爆炸、gradient 为什么消失、显存为什么不够。这不是动手能力差，而是知识链路缺了一段。CS336 的任务就是补上这一段，它的课程设计思想非常明确：不要告诉我你会用库，请证明你能从零写出一个语言模型，并且把它训练到能生成文本的程度。这篇文章就是围绕这门课的学习路线展开的，读完你会知道从哪开始、先学什么、后学什么，以及每一步为什么要这么做。

我更想强调的一点是：这门课的价值不在于“手搓”这个动作本身，而在于它逼你把大模型的各个环节全部过一遍。数据、分词、模型架构、训练目标、优化器、分布式、推理，任何一个环节掉了链子，你的模型都跑不出像样的结果。换句话说，CS336 教的不是某个模型，而是一条完整的研究与工程链路。这篇文章会按照课程的逻辑，把这套链路拆给你看，并附上可运行的 PyTorch 示例代码，方便你跟着跑通第一个最小版本。

1. 这篇文章真正要解决的问题

开始之前，先问你几个问题：你知道语言模型的训练数据应该怎么预处理吗？attention mask 为什么要在 softmax 之前做，而不是之后？AdamW 和 Adam 在权重衰减上到底有什么区别？混合精度训练时为什么 loss 会突然变成 NaN？如果这些问题你能立刻答出来，那 CS336 对你而言是复习课。如果答不出来，说明你之前对大模型的使用更多停留在“调用者”层面，而 CS336 恰好就是围绕这些底层问题展开的。

这门课真正解决的核心问题，是大多数开发者手里没有一条从零开始构建大模型的完整路径。网上有太多“三分钟搭建 Llama”的文章，点进去全是transformers.AutoModel加几行推理代码。一旦你想训练一个自己的模型、想加深并行框架、想读懂 DeepSpeed 的配置，就发现无从下手。CS336 的课程设计跳过了这些封装的抽象层，直接让你从 tokenizer 开始写计算图，从数据加载开始写训练循环，整个过程会重建你对大模型的直觉。

什么样的读者最适合这门课？我认为有三类人值得投入时间。

第一类是想做 LLM 底层研究和推理优化的同学，他们需要理解模型结构每一条边的形状变化，需要知道显存到底被谁吃掉了。第二类是正在做模型微调但反复踩坑的工程师，他们可能已经跑过 LoRA、QLoRA，但当显存不够、loss 不降、分布式通信报错时，需要有一套底层知识来支撑排查。第三类是面试大模型相关岗位的求职者，不少面试官会直接追问 FlashAttention 为什么快、ZeRO 三个阶段切分的是什么，这些问题如果只背结论不读源码，很容易被问穿。

当然，这门课不适合只想快速调 API、快速做业务交付的人。CS336 的投入成本很高，作业通常要花数周才能完整做完。但如果你打算长期做 LLM 相关的工作，这笔时间花得值。学习这门课的收益是：之后你再使用任何上层框架，都能立刻判断它到底帮你封装了什么，又牺牲了什么，而不必把它们当成黑盒。

2. CS336 视角：LLM 从零实现的底层逻辑

CS336 的主题是“Language Modeling from Scratch”，中文可以理解为“从零实现语言建模”。它不依赖任何封装好的大模型框架，课程作业通常要求你用 PyTorch 甚至更底层的接口完成一个能训练的语言模型，再逐步加入优化技巧和分布式能力。

为什么选择“from scratch”作为教学方式？因为大模型本质上没有那么多神秘的数学，它就是把海量文本压缩成 token 序列上的概率分布。Transformer 前向传播的核心代码，认真写也不过一两百行。真正的复杂性在工程细节：数据混洗、padding、mask、学习率调度、梯度裁剪、混合精度、分布式通信、checkpoint 管理。这些细节被框架隐藏了，但恰恰是它们决定了模型能否稳定训练。

从零实现大模型，相当于把机器学习课程、系统课程和工程课程融合到同一条线上。你需要同时处理“模型该长什么样”“数据该怎么喂”“多卡之间怎么通信”“loss 为什么会飞”这几个层面的问题。CS336 的可贵之处，是让你按顺序一次只学一个概念，而不是一开始就抛出一个几百 B 参数的巨型系统。

如果把大模型比作造汽车，那么 Hugging Face 给你的是整车和说明书，DeepSpeed 给你的是可调参数的发动机管理和流水线装配方案，而 CS336 则是要求你自己画图纸、锻零件、组装并跑完试车。这个过程当然慢，但完成后你对每一个零件如何协同工作会有真正的手感。这种手感在遇到问题时会变成判断力：别人还在看文档，你已经能从 loss 曲线和显存占用推断出问题出在数据、模型还是训练配置上。

课程的完整链路通常覆盖数据与分词、模型架构、训练优化、分布式训练、推理与模型压缩这几大块。下面逐个拆解，并给出每一块在课程中最常做的练习和你在实际项目中会遇到的对应问题。

2.1 课程教学路线的一条主线

以我个人的学习体验和公开材料来判断，CS336 的作业设计有一条非常清晰的递进线：先写出能过前向传播的最小模型，再把它训练到能输出 decodable 文本，然后加入更接近真实预训练的优化器与数据配置，最后扩展到多卡分布式和推理优化。这个过程完全复现了早期 GPT 类模型的诞生流程。

这种递进的好处是每一周都有一个可验证的产出。你不需要等三个月才看到成果，第一周你就能打印出自己的模型生成的第一段乱码。乱码虽然可笑，但它是你在黑盒外面亲手构造出来的系统在运转的证据。接下来每一次迭代，把乱码变得稍微有规律一点，把 loss 从 5 降到 4、再到 2，你会非常直观地理解“预训练到底在做什么”。

2.2 和传统深度学习入门课的区别

传统深度学习课程通常从图像分类入手，用交叉熵训练一个 CNN，重点放在卷积和池化上。CS336 则完全围绕自回归语言建模展开，重点变成了序列数据的掩码、因果建模、以及训练过程的不稳定性。同样是交叉熵，语言模型里每个 token 都是一个分类问题；同样是用反向传播，语言模型里序列越长计算图越大，梯度消失和显存峰值的问题更早显现。

这带来的结果是：学完 CS336，你对显存和算力的感知力会明显提升。你会知道一次前向反向具体在哪一层消耗最多显存，从而理解为什么 FlashAttention、梯度检查点、混合精度这些技术能带来数量级的收益。这些感知是单纯调用框架无法获得的。

3. 环境准备与前置条件

动手之前，先确认你要准备的知识和硬件。CS336 对基础的要求不算苛刻，但如果你完全没接触过深度学习和 PyTorch，建议先补一下入门课程。使用这门课期间最重要的能力其实是“读源码”和“查文档”，因为你会经常需要对比自己的实现和参考实现的差异。

3.1 前置知识准备

第一个必备基础是熟练的 Python 编程，尤其是对 tensor 形状变化的敏感度。第二个是基础的深度学习知识，你至少要清楚什么是 embedding、什么是前馈网络、什么是反向传播、什么是梯度下降。第三个是基础的 PyTorch 使用经验，比如nn.Module、DataLoader、optimizer.step()这些 API 应该不用查文档就能写。如果这些都没问题，CS336 的代码对你来说就是“逻辑复杂但每一行都看得懂”。

线性代数和概率论也需要掌握到“够用”的程度。矩阵乘法、softmax、交叉熵、对数似然这些概念是模型代码里的常客。说实话，只要你做过图像分类或 NLP 入门项目，这些数学知识基本已经覆盖。不需要再花一个月补数学，遇到不懂的公式查一下即可。

3.2 硬件与运行环境

看到“从零搭建大模型”，很多人的第一反应是“我没有多卡 GPU 怎么办”。实际上，课程学习可以被划分为多个阶段，不同阶段对硬件的要求差异很大。

第一阶段是阅读课程讲义、学习概念、跑通前向传播和小规模测试，这时候普通的 CPU 笔记本也能完成，只是训练速度慢一些。第二阶段是训练一个小型语言模型来验证 loss 会下降、文本生成有进步，如果使用很小的字符级数据集，单张消费级 GPU 比如 RTX 4060 或 3090 就能胜任。第三阶段才是完整训练一个“像样”的预训练模型，这通常需要多卡甚至集群，课程仓库也会提供对应的分布式脚本。

比较稳妥的意见是：先用小模型把全部代码流程跑通，再考虑上多卡。你不需要一上来就复现一个 7B 模型，那是工业环境的事，不是学习阶段的目标。如果只是验证分布式逻辑，两到四张中端卡也够用，关键是理解数据并行和模型分片各自在做什么。

3.3 创建开发环境

下面给出一个最基础的 Python 虚拟环境创建方式，适配 Linux 和 macOS。Windows 用户建议使用 WSL2 安装 Linux 环境，因为后续很多分布式训练功能在 Linux 下更顺畅。

# 创建并激活虚拟环境 python -m venv .venv source .venv/bin/activate # 升级 pip 并安装基础依赖 pip install --upgrade pip pip install torch numpy tqdm

关于 PyTorch 的安装，不同 CUDA 版本的安装命令不同，请以 PyTorch 官方安装页为准，本文不绑定具体版本。安装完成后可以执行下面命令验证环境是否可用：

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

如果输出中torch.cuda.is_available()为True，说明 GPU 可用。如果是False，需要检查驱动和 CUDA 版本，或者暂时先用 CPU 跑，但要把模型和 batch size 调小。

课程代码的获取方式，建议直接从课程主页或官方 GitHub 仓库拉取。仓库中通常包含讲义、作业模板和参考测试，具体的组织方式以你拿到的最新版本为准。克隆仓库的命令如下：

git clone <课程仓库地址> cd <课程仓库目录>

这里刻意不写死某个仓库地址，是因为课程仓库可能改名或迁移。只要在 GitHub 搜索课程名，一般都能找到官方仓库或高星镜像，注意区分是否官方即可。

4. 核心流程拆解：手搓大模型的五个关键环节

如果你已经跑通过一个大型预训练模型的微调，再去看 CS336 的作业目录，会有一种“原来我之前使用的框架帮我做了这么多事”的感觉。这一节我把从零搭建大模型的过程按五个环节拆开，每个环节都会解释它解决什么问题，以及缺少它会出现什么后果。

4.1 数据与分词

语言模型学习的对象是 token 序列，因此第一步永远是把原始文本切成 token。最简单的方式是字符级 tokenizer，把每个字符当成一个 token。它的优点是实现容易，缺点是序列过长，实践中很少直接用。工业界默认使用 BPE 字节对编码，它先把文本拆成字节，再逐步合并最常见的高频片段，最终得到一个可控大小的词表。

词表大小直接影响模型参数量和训练速度。词表大的好处是每个 token 的语义更完整，坏处是 embedding 矩阵变大，解码端的 LM Head 维度也跟着变大，显存压力随之上升。课程作业通常会让你自己实现或调通一个 BPE tokenizer，并让它处理一个标准数据集。这里最容易踩的坑是训练和推理阶段 tokenizer 不一致，例如训练时用了带<unk>的映射，推理时却输出了未知字符，导致生成结果不可读。

确保 tokenizer 正确的最直接方法，是把一段文本 encode 再 decode 回去，确认内容没有丢失；同时查看 token 序列长度的分布，避免出现过长的样本导致显存溢出。

4.2 模型架构

当前几乎所有主流语言模型都是 Transformer 的变体，核心结构包括 token embedding、位置编码、多层 Transformer Block、最终 LayerNorm 和 LM Head。每个 Transformer Block 内部通常包含 Multi-Head Self-Attention、MLP、残差连接和 LayerNorm。

Self-Attention 的核心是让序列里的每个位置都能根据其他位置的信息更新自己的表示。Q、K、V 三个矩阵分别负责“我要查询什么”“我的标识是什么”“我携带的内容是什么”。因果语言模型还要加一个 causal mask，确保位置 t 只能看到 0 到 t 的信息，不能看到未来。这个 mask 必须在注意力分数归一化之前应用，因为在 softmax 之后再加 mask 会让被遮住的位置仍然拥有非零概率质量，只是略微变小，达不到“完全不能看”的效果。

LayerNorm 的位置也值得注意。现在主流实现倾向于 Pre-Norm，即先归一化再进入子层，优点是训练更稳定，对大学习率更友好。课程里通常会要求你对比 Pre-Norm 和 Post-Norm 的效果差异，你可能会发现，模型层数加深之后，Post-Norm 更容易出现训练不稳定。

4.3 训练设置

语言模型的训练目标很简单：对每个位置预测下一个 token，最大化真实 token 的 log 概率。等价地，最小化交叉熵 loss。但这个目标在实现上有不少细节，比如reduction的选择、padding token 是否需要参与 loss 计算。如果数据里包含 padding，而你又没有在 loss 里屏蔽 padding，模型会花大量精力去预测无害的 padding token，反而学不好真实内容。

优化器方面，现在的默认选择是 AdamW，而不是 Adam。AdamW 把权重衰减从梯度更新中解耦，只在参数更新时直接做 L2 缩权，能显著改善正则效果。配合学习率预热和余弦退火，训练稳定性会好很多。还有一个容易被忽略但重要的超参数是梯度裁剪，它对防止梯度爆炸很有效。

混合精度也是现代训练标配。fp16 能显著减少显存和计算量，但容易造成 loss 溢出。现在更推荐 bf16，它的表示范围更大，对训练稳定性更友好。如果你遇到 loss 为 NaN，先检查是不是混合精度设置不当，再检查学习率是不是过大。

4.4 分布式训练

训练大模型，单卡几乎不可能完成。分布式训练有一套完整的概念体系：数据并行是每张卡持有完整模型，但数据切成多份；模型并行是把模型的不同层放到不同卡；张量并行是把同一层内部的矩阵运算切分到多张卡；流水线并行是把层分组，让不同卡串行处理不同阶段的批次。

在 PyTorch 生态下，DDP 是数据并行的标准实现，适合模型尺寸能够放进单卡的情况。当模型大到单卡装不下时，需要 ZeRO 和 FSDP 这类参数分片方案。FSDP 会把参数、梯度和优化器状态分片到多张卡上，代价是增加通信量。CS336 通常会引导你从 DDP 入手，再过渡到 FSDP 或手动实现一个简单的 ZeRO 阶段，这比直接调用框架更能帮助你理解通信开销的来源。

分布式训练最容易出的问题之一，是模型输出和 loss 在不同 rank 上的数值不一致。要验证你的分布式实现是否正确，可以先在单卡上固定 seed 训练几步，再在多卡上同样的 seed 训练几步，比较最终 loss 是否一致。如果不一致，通常说明通信方案或梯度同步逻辑有问题。

4.5 推理

训练完成后，推理阶段的优化又是一个新世界。自回归生成是一个 token 一个 token 产生的，每个步骤都需要重复计算前面所有 token 的注意力。KV Cache 的出现就是为了避免重复计算，把每个 token 的 Key 和 Value 缓存下来，之后每次只需计算新位置的 Key 和 Value。课堂里你可能还会讲到 Grouped Query Attention、量化、蒸馏，这些技术都属于“模型已经训练好之后，如何更快更省地让它运行”的问题。

推理阶段另一个常见问题是生成质量。直接贪心采样往往导致重复和无趣，温度参数过低也会让文本缺乏多样性，过高则会胡言乱语。合理做法是在工业场景里结合 Top-k、Top-p 采样以及重复惩罚，同时在评测时区分效率和质量的取舍。

5. 完整示例代码：用 PyTorch 手写最小 GPT

理论知识讲再多，不如亲自跑一个最小例子。下面我用一个字符级 GPT 模型作为教学示例，展示“从零搭建大模型”的核心骨架。为了便于复制，代码拆成多个文件，并且刻意保持精简。课程作业里会在这个基础上加入 BPE、真实数据集、分布式等更复杂的内容，但骨架是完全一样的。

5.1 文件结构规划

course-mini-gpt/ ├── data.py # 数据读取和 batch 构造 ├── model.py # MiniGPT 模型定义 ├── train.py # 训练脚本 └── generate.py # 文本生成脚本

文件拆分的目的不是追求架构优雅，而是让你在修改模型、数据处理、训练配置时尽量不影响其他部分。真实项目里还会把配置单独抽成 yaml 文件，这里为了教学就直接写在脚本顶部。

5.2 数据处理代码

# 文件路径：course-mini-gpt/data.py import torch from torch.utils.data import Dataset class CharDataset(Dataset): def __init__(self, text, block_size): chars = sorted(list(set(text))) self.stoi = {ch: i for i, ch in enumerate(chars)} self.itos = {i: ch for i, ch in enumerate(chars)} self.block_size = block_size self.data = torch.tensor([self.stoi[c] for c in text], dtype=torch.long) def __len__(self): return len(self.data) - self.block_size - 1 def __getitem__(self, idx): x = self.data[idx: idx + self.block_size] y = self.data[idx + 1: idx + 1 + self.block_size] return x, y

这个 Dataset 的用法是：给定一段文本和固定的上下文长度block_size，把文本切成很多个长度为block_size的输入序列，同时把序列右移一个 token 作为预测目标。stoi和itos分别负责字符和数字之间的映射。

5.3 模型定义代码

# 文件路径：course-mini-gpt/model.py import math import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout): super().__init__() assert embed_dim % num_heads == 0 self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=False) self.out_proj = nn.Linear(embed_dim, embed_dim) self.dropout = nn.Dropout(dropout) def forward(self, x): B, T, C = x.shape qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=-1) q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = k.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, device=x.device)).view(1, 1, T, T) att = att.masked_fill(mask == 0, float("-inf")) att = torch.softmax(att, dim=-1) att = self.dropout(att) y = att @ v y = y.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(y) class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads, dropout): super().__init__() self.ln1 = nn.LayerNorm(embed_dim) self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout) self.ln2 = nn.LayerNorm(embed_dim) self.mlp = nn.Sequential( nn.Linear(embed_dim, 4 * embed_dim), nn.GELU(), nn.Linear(4 * embed_dim, embed_dim), nn.Dropout(dropout), ) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class MiniGPT(nn.Module): def __init__(self, vocab_size, block_size, embed_dim, num_heads, num_layers, dropout=0.1): super().__init__() self.block_size = block_size self.token_embedding = nn.Embedding(vocab_size, embed_dim) self.position_embedding = nn.Embedding(block_size, embed_dim) self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads, dropout) for _ in range(num_layers) ]) self.ln_f = nn.LayerNorm(embed_dim) self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, std=0.02) def forward(self, idx, targets=None): B, T = idx.shape tok = self.token_embedding(idx) pos = self.position_embedding(torch.arange(T, device=idx.device)) x = tok + pos 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)) return logits, loss

这份代码有几个关键点需要理解。

第一，注意力打分的缩放因子是math.sqrt(self.head_dim)，目的是防止点积结果过大导致 softmax 梯度消失。第二，causal mask 使用torch.tril生成下三角矩阵，再在上三角位置填充-inf，这样 softmax 之后未来位置的概率为 0。第三，每个 TransformerBlock 都采用 Pre-Norm 结构，先 LayerNorm 再进入注意力或 MLP 子层，残差连接保留原始输入，有助于深层网络稳定训练。

还需要注意，这个示例的lm_head和token_embedding没有做权重绑定。权重绑定能减少参数量，但在教学代码里先不引入这个细节，等到复现更完整的模型时再考虑即可。

5.4 训练脚本

# 文件路径：course-mini-gpt/train.py import torch import torch.nn as nn from torch.utils.data import DataLoader from data import CharDataset from model import MiniGPT # 超参数 BLOCK_SIZE = 64 BATCH_SIZE = 32 EMBED_DIM = 128 NUM_HEADS = 4 NUM_LAYERS = 4 DROPOUT = 0.1 LEARNING_RATE = 3e-4 TRAIN_STEPS = 2000 EVAL_EVERY = 200 SEED = 42 def load_text(): # 这里是演示数据，实际课程中会换成真实的预训练语料 text = ( "hello world! this is a tiny language model. " "it learns to predict the next character. " "we are building an llm from scratch. " ) * 200 return text def main(): torch.manual_seed(SEED) text = load_text() dataset = CharDataset(text, BLOCK_SIZE) loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True) model = MiniGPT( vocab_size=len(dataset.stoi), block_size=BLOCK_SIZE, embed_dim=EMBED_DIM, num_heads=NUM_HEADS, num_layers=NUM_LAYERS, dropout=DROPOUT, ) optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE) model.train() step = 0 for epoch in range(100): for x, y in loader: logits, loss = model(x, y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if step % EVAL_EVERY == 0: print(f"step {step}, loss {loss.item():.4f}") step += 1 if step >= TRAIN_STEPS: torch.save(model.state_dict(), "mini_gpt.pt") print("训练完成，模型已保存到 mini_gpt.pt") return if __name__ == "__main__": main()

训练脚本中的关键细节是梯度裁剪。语言模型训练很容易出现梯度爆炸，梯度裁剪能防止参数单步更新过大，减少训练发散的概率。另一个容易踩坑的地方是 DataLoader 的shuffle=True，如果忘了 shuffle，模型会一直按文本顺序看到数据，虽然也能学，但收敛会更慢。

5.5 文本生成脚本

# 文件路径：course-mini-gpt/generate.py import torch from data import CharDataset from model import MiniGPT def generate(model, idx, max_new_tokens, temperature=1.0): model.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -model.block_size:] with torch.no_grad(): logits, _ = model(idx_cond) logits = logits[:, -1, :] / temperature probs = torch.softmax(logits, dim=-1) idx_next = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, idx_next), dim=1) return idx def main(): text = ( "hello world! this is a tiny language model. " "it learns to predict the next character. " "we are building an llm from scratch. " ) * 200 dataset = CharDataset(text, BLOCK_SIZE=64) model = MiniGPT( vocab_size=len(dataset.stoi), block_size=64, embed_dim=128, num_heads=4, num_layers=4, dropout=0.1, ) model.load_state_dict(torch.load("mini_gpt.pt")) model.eval() context = torch.tensor([[dataset.stoi['h']]], dtype=torch.long) output = generate(model, context, max_new_tokens=200, temperature=0.8) decoded = "".join([dataset.itos[i] for i in output[0].tolist()]) print(decoded) if __name__ == "__main__": main()

生成脚本使用的是自回归采样：每次把已有序列的最后block_size个 token 输入模型，取预测的最后一个位置的概率分布，然后按概率采样一个新 token 拼接回去。温度参数控制概率分布的平滑程度，温度越低，采样越保守；温度越高，输出越发散。torch.multinomial是从概率分布中采样的核心 API，它保证了生成过程不是每次都取最大概率 token，因此会有更多变化。

6. 运行结果与效果验证

代码写完之后，按顺序执行训练和生成即可。

python train.py python generate.py

训练脚本运行后，你会在终端看到类似下面的输出：

step 0, loss 4.1234 step 200, loss 2.9982 step 400, loss 2.5401 step 600, loss 2.3017 step 800, loss 2.1503 step 1000, loss 2.0449 ... step 2000, loss 1.8532 训练完成，模型已保存到 mini_gpt.pt

loss 的具体数值取决于数据集大小、模型尺寸和随机种子，但趋势必须是持续下降的。如果 loss 在某个值附近震荡不降或者直接变成 NaN，说明训练配置有问题。这里我使用的演示数据集是非常小的重复文本，因此 loss 会降到很低，这是正常现象。

判断训练成功的标准有三个。第一，loss 在训练步数内没有出现 NaN，并且整体呈下降趋势。第二，训练结束后加载模型，能生成与训练文本存在相似规律的内容，比如已知的训练文本是英文，生成结果即使不完整也会以英文字母和空格为主，如果生成的全是不可打印字符，说明 tokenizer 或采样逻辑有误。第三，模型输出的文本长度正确，没有因为位置编码越界而报错。

如果你看到生成结果完全和训练文本相同，那也可能是过拟合。演示数据集太小，模型可能背下来一部分内容。在课程作业里，你会使用更大的标准数据集，并引入验证集来判断是否过拟合。关于这一点，一个直观的验证方法是在训练时保留一小段验证文本，观察训练 loss 和验证 loss 的差距。

7. 常见问题与排查方法

从零搭建大模型的过程中，以下问题几乎每个人都会遇到。我整理成表格，方便你在遇到异常时快速对照。

这里重点展开说明几个高频问题。

第一个是 loss 为 NaN。很多人第一反应是换数据集，但实际上最常见的原因是学习率过大，或者 fp16 混合精度导致梯度溢出。排查时先固定随机种子，关闭混合精度跑几步，如果恢复正常，再逐步开放混合精度。

第二个是 loss 不下降。我见过不少同学在数据预处理阶段把字符映射关系写错，比如训练集的stoi和推理集不一致，模型相当于在猜测一套无关的编码。这时候除了检查数据，还要做一个最小实验：让模型在单条样本上反复训练，如果连单条样本都无法过拟合，说明代码逻辑本身有问题。

第三个是显存 OOM。对新手来说，最直接的止损办法是把 batch size 降到 1，确认单个样本能跑通，再逐步增大。多卡训练时，每张卡的 batch size 不是全局 batch size，算梯度时要注意是否需要梯度累积。

8. 最佳实践与工程建议

代码跑通只是第一步，真正让课程有价值的是你如何组织实验、记录结果、排查问题。这套工程能力直接复用到真实项目中。

8.1 先跑通最小规模，再逐步扩展

不要一上来就在大模型上反复尝试。先用一个很小的数据集、很小的模型跑通全流程，确认数据流、模型前向、loss 计算、反向传播、保存和加载全部正确。之后再把模型尺寸和数据集放大。这个习惯能省下大量 Debug 时间，因为在最小场景下，你能更容易判断问题是出在代码逻辑还是资源限制。

8.2 固定随机种子，保证实验可复现

语言模型训练中随机因素非常多，包括数据加载的 shuffle、参数初始化、dropout。如果不固定 seed，你可能无法判断两次实验的差异来自代码改动还是随机噪声。推荐在脚本入口固定 PyTorch、NumPy 和 Python 内置 random 的 seed。在实验记录里也要写下配置文件的哈希值或版本号，方便回溯。

8.3 把配置外置，而不是写死在代码里

训练脚本里的超参数一旦多起来，硬编码会很痛苦。建议把block_size、batch_size、learning_rate、model_config等内容整理成 yaml 或 json 文件，训练脚本读取配置后运行。课程作业阶段可以不做，但如果你想复现多组实验结果，配置外置几乎是必须的。

8.4 重视数据质量和数据版权

模型的性能上限很大程度上由数据决定。课程使用的数据集往往是公开的，方便复现和比较。但从零搭建大模型的目标如果是真实产品，你还需要注意数据来源的合法性，确认数据符合使用条款和隐私规定，不要在未授权的情况下使用私人或敏感文本。对生成内容也要有安全过滤意识，防止模型产出不当信息。

8.5 日志、checkpoint 与监控

训练过程中至少每几百步打印一次 loss 和梯度范数。建议把梯度范数的变化也纳入监控，因为梯度范数突然暴增往往预示训练即将发散。checkpoint 不要只保存模型权重，还应保存优化器状态、当前步数、随机种子和数据索引，确保你可以从中间状态恢复训练。

8.6 安全边界与权限意识

如果你把代码部署到生产集群，务必遵循最小权限原则。训练数据、模型权重、实验日志都应放在受控目录中，不要用 root 权限运行训练任务。如果涉及多机分布式训练，注意通信端口的安全策略，避免对外暴露不需要的端口。任何时候对已有模型进行覆盖或删除，都要先确认备份。

9. 总结与后续学习方向

现在回头看，CS336 教会你的不只是一个模型结构，而是一整套构建和调试大模型的方法论。你会知道自己写的 loss 为什么降不下去，也会知道分布式训练里的通信开销到底从哪来，更能在别人只会调用框架时，判断出一个训练任务真正需要多少算力和显存。课程作业如果完整做下来，你对数据、模型、训练、推理、分布式的理解会上一个层次。

接下来怎么走，取决于你的目标。如果你对模型训练感兴趣，可以继续深入研究分布式训练框架的源码，试着从 DDP 移植到 FSDP。如果你对推理优化感兴趣，可以研究 KV Cache、量化、蒸馏以及 FlashAttention 的实现思路。如果你希望做一个真正可用的产品，那么下一步应该尝试在大规模公开数据集上训练一个小规模的领域模型，并设计一套验证集和评测指标来跟踪效果。

有一点要记住：直接抄别人写好的作业代码，收获会大打折扣。更建议的做法是自己先独立完成任务，再对照官方参考实现检查差异，逐行理解别人为什么那样写。这个“先自己写、再对比、再反思”的循环，才是手搓大模型真正的价值所在。

如果你准备入坑 CS336，建议从课程第一部分开始，先不要急着看后面的分布式内容。第一周的目标很简单：用你自己的代码跑出一个 loss 在下降的小模型。等这一步达成，后面的事情都会顺理成章。