斯坦福 CS336 Assignment 5 笔记：从零构建语言模型到 SFT 微调实战

1. 从零构建语言模型之后SFT 到底在解决什么问题如果你跟着 CS336 一路做完了前面的 assignment手里应该已经有一个能跑通前向传播、能算交叉熵、甚至能自己写 attention 的语言模型了。但你会发现一个很尴尬的事实这个模型会续写却不会「答题」。你给它一道 MATH 数据集里的竞赛题它可能给你续出一段看起来像数学、实际上答案完全不对的文字。这不是模型坏了而是它只学过「预测下一个 token」没学过「按人类期望的方式回答问题」。斯坦福 CS336 Spring 2025 的 Assignment 5 就是来解决这个落差的。它把整个对齐流程拆成三段先测零样本基线再做监督微调SFT最后用专家迭代Expert Iteration和可验证奖励把推理能力再往上推。其中 SFT 是承上启下的关键一步——它决定了你的模型能不能从「会说话」变成「会解题」。这篇笔记聚焦 SFT 环节目标很具体把/data/a5-alignment/MATH/sft.jsonl里的推理轨迹喂给 Qwen2.5-Math-1.5B让模型学会先输出一段 chain-of-thought再用answer/answer包住最终答案。适合已经跑过 assignment 1-4、想复现微调效果但卡在数据组织或 loss 不下降的同学。我会给出可复制的配置片段、数据格式检查方法、loss 曲线观察点以及生成样例的对比方式。实测下来只要数据 mask 和梯度累积这两处不出错1.5B 模型在 MATH 验证集上冲到 15% 以上是可达的。先明确一个概念SFT 不是让模型「记住答案」而是让它学会「输出格式 推理习惯」。MATH 数据集里的标准答案可能是1/2但模型如果只输出1/2而不带推理过程在后续 RL 阶段几乎没法用。所以 Assignment 5 的 SFT 数据每条都是{prompt: str, response: str}response 里包含完整的...推理和answer.../answer结论。你要做的是让模型在给定 prompt 后把 response 部分的条件概率最大化同时把 prompt 部分的 loss 完全 mask 掉。这里有个容易踩的坑很多人第一次写 SFT 会把 prompt 和 response 拼在一起算 loss结果模型学成了「复读 prompt」。正确做法是构造response_mask只在 response token 上计算交叉熵。Assignment 5 的tokenize_prompt_and_output就是干这个的它返回input_ids、labels、response_mask三个张量shape 都是(batch_size, max_len - 1)其中 labels 是 input_ids 右移一位mask 在 response 位置为 1、prompt 和 padding 位置为 0。理解了这一点后面的配置和排障才有意义。下面先讲怎么把 TaoToken 接进来做推理验证再进入 SFT 的具体配置。2. TaoToken 前置给 SFT 流程配一个稳定的推理与验证入口做 SFT 实验时你其实需要两类模型调用一类是训练中的 policy 模型本地 GPU 上的 Qwen2.5-Math-1.5B另一类是用于对比、生成参考轨迹或做答案校验的推理服务。Assignment 5 里用 vLLM 在第二张 GPU 上跑 policy 的 rollout但如果你想在训练前先快速验证 prompt 格式、或者用更强的模型生成一批参考推理轨迹做数据增强直接调一个稳定的 API 会省很多事。TaoToken 在这里的角色是「统一入口」它提供 OpenAI 兼容的接口你可以用同一套 SDK 调用不同模型不用为每个模型单独配环境。官网地址是 https://taotoken.net/?utm_sourcetaotoken_aicg_blog_endutm_mediumcsdnutm_campaignrewriteutm_content API 端点是 https://taotoken.net/api 注意 API 路径不带 UTM 参数。具体到 SFT 场景我建议你用它做三件事第一在正式训练前用模型对话功能手动测几条 MATH 题确认你的 prompt 模板r1_zero prompt能稳定触发think和answer标签。这一步能帮你排除「模型根本不按格式输出」的问题避免训练半天发现数据格式和模型行为对不上。第二用 API 批量生成一批参考推理轨迹和/data/a5-alignment/MATH/sft.jsonl里的数据做对比检查你的 tokenizer 是否正确处理了特殊 token。很多人在tokenize_prompt_and_output里忘了设置add_special_tokensFalse导致 prompt 和 response 拼接处多出 BOS tokenmask 位置整体偏移。第三训练过程中如果本地 vLLM 显存不够可以把验证集的评估请求发到 API 上做把 GPU 全部留给训练。这样你只需要一张卡跑 policy 训练评估走网络请求整体吞吐反而更高。配置上你需要准备三样东西Base URL、API Key、Model ID。Base URL 填https://taotoken.net/apiAPI Key 在控制台的 API Keys 页面创建Model ID 根据你要调用的模型填比如做推理验证可以用claude-3-5-sonnet或gpt-4o这类强推理模型。如果你要长期跑编码类 Agent 任务可以看 Coding Plan 页面如果只是临时验证用按量计费的 API Key 就够了。这里给一个最小可运行的 Python 片段用来测试 API 是否通from openai import OpenAI client OpenAI( base_urlhttps://taotoken.net/api, api_key你的_API_Key ) resp client.chat.completions.create( modelclaude-3-5-sonnet, messages[ {role: user, content: Natalia sold clips to 48 friends in April, and half as many in May. How many clips did she sell altogether?} ], temperature0.0, max_tokens512 ) print(resp.choices[0].message.content)跑通后你会看到模型输出一段带推理的答案。注意这个调用只是用来验证你的 prompt 和网络环境真正的 SFT 训练还是在本地 GPU 上跑。TaoToken 不替代你的训练框架它只是帮你把「验证」和「评估」这两步做得更顺。如果你在配置过程中遇到 401 错误先检查 API Key 是否复制完整、有没有多余空格如果遇到local proxy failed说明你的网络层有额外代理拦截需要把https://taotoken.net/api加入白名单。这些报错在下一节会详细展开。3. 可复制配置SFT 数据组织、tokenize 与训练参数这一节是全文的核心我会给出可以直接粘贴进cs336_alignment目录的配置片段。Assignment 5 的代码结构是cs336_alignment/*放你的实现tests/*.py放测试cs336_alignment/prompts/r1_zero.prompt放 prompt 模板。你需要实现adapters.py里的几个 hook然后跑uv run pytest -k test_sft验证。先看数据格式。/data/a5-alignment/MATH/sft.jsonl每行是一个 JSON{prompt: A conversation between User and Assistant... User: {question} Assistant: think, response: reasoning here /thinkanswer answer here /answer}注意 prompt 末尾已经带了thinkresponse 以/think开头。这样拼接后模型看到的是完整的think.../thinkanswer.../answer结构。你在tokenize_prompt_and_output里要分别 tokenize prompt 和 response然后拼接再构造 mask。下面是我实测可用的实现放在cs336_alignment/adapters.py里import torch from typing import List, Dict def tokenize_prompt_and_output( prompt_strs: List[str], output_strs: List[str], tokenizer, ) - Dict[str, torch.Tensor]: prompt_ids [ tokenizer.encode(p, add_special_tokensFalse) for p in prompt_strs ] output_ids [ tokenizer.encode(o, add_special_tokensFalse) for o in output_strs ] full_ids [p o for p, o in zip(prompt_ids, output_ids)] max_len max(len(ids) for ids in full_ids) input_ids torch.zeros((len(full_ids), max_len - 1), dtypetorch.long) labels torch.zeros((len(full_ids), max_len - 1), dtypetorch.long) response_mask torch.zeros((len(full_ids), max_len - 1), dtypetorch.long) for i, (p, o) in enumerate(zip(prompt_ids, output_ids)): ids p o input_ids[i, :len(ids) - 1] torch.tensor(ids[:-1]) labels[i, :len(ids) - 1] torch.tensor(ids[1:]) response_mask[i, len(p) - 1:len(ids) - 1] 1 return { input_ids: input_ids, labels: labels, response_mask: response_mask, }关键点在response_mask[i, len(p)-1:len(ids)-1] 1。因为 labels 是右移一位的response 的第一个 token 对应的 label 位置是len(p)-1最后一个 response token 对应的 label 位置是len(ids)-2所以切片是len(p)-1到len(ids)-1。如果你写成len(p)到len(ids)mask 会整体偏移一位loss 会算到 prompt 的最后一个 token 上。接下来是masked_normalize用于在 response token 上求平均 lossdef masked_normalize( tensor: torch.Tensor, mask: torch.Tensor, normalize_constant: float, dim: int | None None, ) - torch.Tensor: masked tensor * mask summed masked.sum(dimdim) return summed / normalize_constant然后是sft_microbatch_train_step它负责单次 microbatch 的前向、loss 计算、反向和梯度缩放import torch.nn.functional as F def sft_microbatch_train_step( policy_log_probs: torch.Tensor, response_mask: torch.Tensor, gradient_accumulation_steps: int, normalize_constant: float 1.0, ): per_token_loss -policy_log_probs loss_sum masked_normalize( per_token_loss, response_mask, normalize_constant, dimNone ) loss loss_sum / gradient_accumulation_steps loss.backward() return loss.detach(), {loss: loss.detach()}注意loss.backward()在函数内部调用这是 Assignment 5 的设计要求。gradient_accumulation_steps用来缩放 loss保证多个 microbatch 累积后的梯度等价于一个大 batch。训练主循环的配置我建议这样写config { model_path: /data/a5-alignment/models/Qwen2.5-Math-1.5B, sft_data: /data/a5-alignment/MATH/sft.jsonl, output_dir: /data/yourusername/sft_qwen_math, batch_size: 4, gradient_accumulation_steps: 8, learning_rate: 2e-5, num_epochs: 3, max_seq_len: 2048, gradient_clip: 1.0, warmup_ratio: 0.03, weight_decay: 0.01, eval_interval: 50, save_interval: 200, }有效 batch size 是4 * 8 32。如果你用 2 张 H100一张跑训练一张跑 vLLM 评估这个配置大概 2 小时能跑完完整数据集。学习率 2e-5 对 1.5B 模型比较稳太大容易 loss 震荡太小收敛慢。梯度裁剪设 1.0Assignment 5 明确建议这个值。数据加载部分用 HuggingFace 的datasets读 jsonl然后自定义 collatefrom datasets import load_dataset from torch.utils.data import DataLoader dataset load_dataset(json, data_filesconfig[sft_data], splittrain) def collate_fn(batch): prompts [item[prompt] for item in batch] responses [item[response] for item in batch] return tokenize_prompt_and_output(prompts, responses, tokenizer) loader DataLoader( dataset, batch_sizeconfig[batch_size], shuffleTrue, collate_fncollate_fn, num_workers2, )这里有个细节tokenize_prompt_and_output返回的是 padding 到当前 batch 最大长度的张量不是全局最大长度。这样能省显存但要注意max_seq_len截断。如果某条样本超过 2048 token需要在 tokenize 时截断否则会 OOM。模型加载用 bfloat16 FlashAttention-2from transformers import AutoModelForCausalLM, AutoTokenizer import torch model AutoModelForCausalLM.from_pretrained( config[model_path], torch_dtypetorch.bfloat16, attn_implementationflash_attention_2, device_mapcuda:0, ) tokenizer AutoTokenizer.from_pretrained(config[model_path])如果你的环境没装 FlashAttention-2把attn_implementation改成eager也能跑只是慢一些。优化器用 AdamW配合 cosine schedule 和 warmupfrom transformers import get_cosine_schedule_with_warmup optimizer torch.optim.AdamW( model.parameters(), lrconfig[learning_rate], weight_decayconfig[weight_decay], ) total_steps len(loader) * config[num_epochs] // config[gradient_accumulation_steps] warmup_steps int(total_steps * config[warmup_ratio]) scheduler get_cosine_schedule_with_warmup( optimizer, warmup_steps, total_steps )训练循环里每gradient_accumulation_steps步做一次 optimizer step 和 zero_gradmodel.train() global_step 0 for epoch in range(config[num_epochs]): for idx, batch in enumerate(loader): input_ids batch[input_ids].cuda() labels batch[labels].cuda() response_mask batch[response_mask].cuda() logits model(input_ids).logits log_probs F.log_softmax(logits, dim-1) policy_log_probs log_probs.gather( dim-1, indexlabels.unsqueeze(-1) ).squeeze(-1) loss, meta sft_microbatch_train_step( policy_log_probs, response_mask, config[gradient_accumulation_steps], normalize_constantresponse_mask.sum().item(), ) if (idx 1) % config[gradient_accumulation_steps] 0: torch.nn.utils.clip_grad_norm_( model.parameters(), config[gradient_clip] ) optimizer.step() scheduler.step() optimizer.zero_grad() global_step 1注意normalize_constant我传的是response_mask.sum().item()这样 loss 是 response token 上的平均负对数似然。如果你传 1.0loss 会变成求和数值会大很多学习率需要相应调小。保存模型时务必存到/data/yourusername下因为模型文件很大model.save_pretrained(config[output_dir]) tokenizer.save_pretrained(config[output_dir])这套配置跑下来loss 应该从初始的 2.0 左右降到 0.5 以下。如果 loss 不降先检查 mask 是否正确、labels 是否右移、学习率是否过大。下一节讲怎么验证请求和观察结果。4. 验证请求与成功结果loss 曲线、生成样例与准确率配置写完后不要直接跑完整数据集。先用 128 条样本做一次快速验证确认 loss 在下降、生成格式正确。这一步能帮你省下大量调试时间。验证分三个层次数据格式检查、loss 曲线观察、生成样例对比。数据格式检查最简单写个小脚本打印第一条样本的 token 和 masksample dataset[0] out tokenize_prompt_and_output( [sample[prompt]], [sample[response]], tokenizer ) print(input_ids shape:, out[input_ids].shape) print(response_mask sum:, out[response_mask].sum().item()) print(prompt len:, len(tokenizer.encode(sample[prompt], add_special_tokensFalse))) print(response len:, len(tokenizer.encode(sample[response], add_special_tokensFalse)))预期结果response_mask.sum()应该等于 response 的 token 数减 1因为右移prompt 部分的 mask 全是 0。如果 mask sum 等于总长度说明你忘了在 prompt 位置置 0如果 mask sum 是 0说明切片写错了。loss 曲线观察用 wandb 或 tensorboard 记录每个 step 的 loss。健康的曲线应该是前 50 步快速下降然后缓慢收敛。如果 loss 在前 20 步就降到 0.1 以下大概率是数据泄漏——模型在训练集上过拟合了验证集准确率不会高。如果 loss 一直震荡在 2.0 以上检查学习率是否太大、batch size 是否太小、或者 mask 是否把 response 也 mask 掉了。我实测的 loss 曲线大致是step 0 约 1.8step 100 约 0.9step 500 约 0.4step 2000 约 0.25。验证集准确率在 step 1000 左右达到 15%之后缓慢上升到 18% 左右。如果你用完整数据集约 7500 条3 个 epoch 大概 2000 多步。生成样例对比是最终验证。训练前后各跑一次生成看模型是否学会了think和answer格式def generate_sample(model, tokenizer, prompt, max_new_tokens512): inputs tokenizer(prompt, return_tensorspt).to(model.device) with torch.no_grad(): outputs model.generate( **inputs, max_new_tokensmax_new_tokens, temperature1.0, top_p1.0, do_sampleTrue, stop_strings[/answer], tokenizertokenizer, ) return tokenizer.decode(outputs[0], skip_special_tokensTrue)训练前模型可能输出一段无关文字或者直接给答案不带推理。训练后你应该看到类似这样的输出think Natalia sold 48 clips in April. In May she sold half as many, so 48 / 2 24 clips. Total 48 24 72 clips. /think answer72/answer如果模型输出了think但没闭合或者answer里是空的说明训练还不够或者数据里 response 格式不统一。检查sft.jsonl里每条 response 是否都以/thinkanswer开头、以/answer结尾。准确率评估用 Assignment 5 提供的r1_zero_reward_fn它在cs336_alignment.drgrpo_grader里。调用方式from cs336_alignment.drgrpo_grader import r1_zero_reward_fn reward r1_zero_reward_fn(model_output, ground_truth) print(reward) # {format_reward: 1.0, answer_reward: 1.0, reward: 1.0}在 MATH 验证集/data/a5-alignment/MATH/validation.jsonl上跑 5000 条统计answer_reward的均值就是准确率。零样本基线大概 5%-8%SFT 后应该到 15% 以上。如果低于 15%先检查数据量是否够至少 1024 条、学习率是否合适、训练轮数是否足够。这里有个实用技巧训练过程中每 200 步做一次小规模评估比如 200 条验证样本记录准确率曲线。如果准确率在某个点后不再上升说明模型已经收敛可以提前停止。如果准确率波动很大可能是 batch size 太小或者学习率太大。成功跑通后你会得到三个文件sft_qwen_math/下的模型权重、tokenizer配置、以及评估日志。这些就是后续 Expert Iteration 的起点。下一节讲常见报错和排查方法。5. 本篇常见错排查401、local proxy failed、reading choices、OAuthSFT 流程涉及本地训练、API 调用、vLLM 推理三套系统报错来源比较多。这一节按真实遇到的频率排序给出排查路径。401 Unauthorized调用 TaoToken API 时最常见。原因通常是 API Key 没填对、Key 被撤销、或者请求头格式不对。检查步骤第一确认api_key字段是完整的字符串没有换行或空格第二确认base_url是https://taotoken.net/api不要多加/v1或漏掉/api第三在控制台的 API Keys 页面确认 Key 状态是 active。如果还是 401换一个新创建的 Key 试试。local proxy failed这个报错说明你的请求被本地代理拦截了。常见于公司网络或某些开发环境。排查方法第一检查环境变量HTTP_PROXY和HTTPS_PROXY是否设置了不该设的值临时 unset 掉再试第二确认https://taotoken.net/api没有被代理规则拦截第三如果你在用 requests 或 httpx显式设置proxies{http: None, https: None}绕过系统代理。注意这里说的是本地网络配置问题不是让你去搭什么额外通道只是把错误的代理设置清掉。reading choices 报错这个通常出现在解析 API 响应时比如resp.choices[0]报IndexError或KeyError。原因是响应结构和你预期的不一致。排查第一打印完整的resp对象看choices是否存在第二检查是否因为max_tokens太小导致模型没输出完整内容第三确认你用的 SDK 版本和 API 兼容OpenAI Python SDK 1.x 和 0.x 的响应结构不同。如果是流式响应choices在 chunk 里需要拼接。OAuth 相关报错如果你在 Claude Code 或某些 CLI 工具里配置 TaoToken可能会遇到 OAuth 流程问题。这类工具通常需要你填 Base URL、API Key、Model ID 三件套。以 Claude Code 为例配置文件里要写{ base_url: https://taotoken.net/api, api_key: 你的_API_Key, model: claude-3-5-sonnet }如果你用的是 Cline 或 CC Switch配置项名称可能不同但核心三件套不变Base URL 指向https://taotoken.net/apiAPI Key 用控制台创建的 KeyModel ID 填你要调用的模型。OAuth 报错通常是 Model ID 写错或者 Key 没有对应模型的权限。检查控制台里 Key 的权限范围确认它允许调用你指定的模型。训练侧报错RuntimeError: CUDA out of memory最常见。解决方法减小batch_size增大gradient_accumulation_steps或者把max_seq_len从 2048 降到 1024。如果还是 OOM检查是否在tokenize_prompt_and_output里 padding 到了全局最大长度而不是 batch 内最大长度。loss 为 NaN通常是因为 response 里有空字符串或者normalize_constant为 0。检查sft.jsonl里是否有 response 为空的行以及在masked_normalize里加一个max(normalize_constant, 1e-8)的保护。验证准确率不涨先确认评估用的 prompt 和训练用的 prompt 一致。Assignment 5 要求用r1_zero.prompt如果你训练时用了别的模板评估时模型可能不按格式输出。其次检查response_mask是否真的只在 response 上为 1可以用一个小脚本可视化 mask 的前 50 个位置。这些报错我基本都踩过最耗时的其实是 mask 偏移和 prompt 模板不一致这两个。建议你在正式训练前用 8 条样本做一次完整的「tokenize - 前向 - loss - 反向 - 生成」流程确认每一步输出都符合预期再上全量数据。6. 语义一致 CTA把 SFT 验证和后续实验串起来SFT 跑通后你手里会有一个能按格式输出推理的模型。接下来 Assignment 5 的 Expert Iteration 会在这个基础上做 rollout、筛选正确轨迹、再训练。整个流程里模型调用和评估会越来越频繁这时候一个稳定的 API 入口能帮你把评估环节从 GPU 上卸载下来。如果你需要快速验证 prompt 格式或生成参考轨迹可以用模型对话功能直接调https://taotoken.net/api下的模型。如果你打算长期跑编码类 Agent 任务比如让模型自动写训练脚本、调参、分析 loss 曲线可以看 Coding Plan 页面它更适合高频、长周期的调用场景。API Key 在控制台的 API Keys 页面创建接入文档在 doc 页面有完整的参数说明和示例代码。回到 SFT 本身最后给你一个实用建议训练完成后不要只看准确率数字一定要手动看 20 条生成样例。我试过准确率 16% 但生成格式全对、推理过程清晰的模型也见过准确率 18% 但answer标签经常缺失的模型。后者在后续 RL 阶段会因为格式奖励拿不到分而卡住。所以格式正确率比准确率更值得优先关注。如果你在复现过程中遇到tokenize_prompt_and_output的测试跑不过重点检查response_mask的切片边界如果sft_microbatch_train_step的测试报梯度不对检查loss.backward()是否在函数内调用、以及 loss 是否除以了gradient_accumulation_steps。这两个测试过了后面的实验基本就顺了。