【LLM训练系列】NanoGPT源码详解和中文GPT训练实践:从零搭建到TaoToken统一Key调用

📅 发布时间:2026/10/10 11:08:56
【LLM训练系列】NanoGPT源码详解和中文GPT训练实践:从零搭建到TaoToken统一Key调用
1. 从 NanoGPT 源码到中文 GPT 训练我踩过的坑与跑通路径NanoGPT 是 karpathy 开源的一个极简 GPT-2 复现项目核心代码不到 500 行却完整实现了 Transformer Decoder 的前向、反向与分布式训练。它适合谁适合想真正搞懂 GPT 内部结构、又不想被 HuggingFace 庞大封装淹没的开发者。我自己是从「跑通 shakespeare_char」开始再逐步换成中文语料、换 tokenizer、最后用 TaoToken 统一 Key 做推理验证整个过程踩了不少坑这篇就把可复制的步骤和配置完整写出来。先说结论NanoGPT 本身对 tokenizer 无感知只要输入是 int 类型的 token id 序列就能训练。这意味着你可以用字符级分词、BBPE 自训练分词也可以直接用 Qwen2 的 tokenizer。区别只在于词表大小、压缩率和最终续写效果。中文场景下字符级分词词表约 4000~5000BBPE 自训练约 8000~15000Qwen2 约 15 万。词表越大同样语料下每个 token 的 embedding 训练越不充分小语料容易训崩。下面按「源码理解 → 中文语料准备 → 训练配置 → 推理验证」的顺序展开每一步都给可复制的命令和配置片段。2. NanoGPT 源码逐层拆解CausalSelfAttention 与训练循环关键代码2.1 模型结构Block LayerNorm Attention MLPGPT 的核心就是CausalSelfAttention LayerNorm MLP组成的 TransformerDecoder对应源码里的Block。先看 LayerNorm 和 MLPclass LayerNorm(nn.Module): def __init__(self, ndim, bias): super().__init__() self.weight nn.Parameter(torch.ones(ndim)) self.bias nn.Parameter(torch.zeros(ndim)) if bias else None def forward(self, input): return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5) class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc nn.Linear(config.n_embd, 4 * config.n_embd, biasconfig.bias) self.gelu nn.GELU() self.c_proj nn.Linear(4 * config.n_embd, config.n_embd, biasconfig.bias) self.dropout nn.Dropout(config.dropout) def forward(self, x): x self.c_fc(x) # (B, T, C) - (B, T, 4C) x self.gelu(x) x self.c_proj(x) x self.dropout(x) return xMLP 的升维倍率是 4这是 GPT-2 的经典设定。c_fc把n_embd升到4*n_embdc_proj再降回来中间用 GELU 激活。这个结构在后续的 LLaMA 里被换成 SwiGLU但 NanoGPT 保持 GPT-2 原样。2.2 CausalSelfAttention因果掩码与 Flash Attention注意力部分是理解 GPT 的关键。核心公式是softmax(QK^T / sqrt(d)) V但必须加因果掩码保证每个位置只能看到自己和之前的位置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, biasconfig.bias) self.c_proj nn.Linear(config.n_embd, config.n_embd, biasconfig.bias) self.attn_dropout nn.Dropout(config.dropout) self.resid_dropout nn.Dropout(config.dropout) self.n_head config.n_head self.n_embd config.n_embd self.flash hasattr(torch.nn.functional, scaled_dot_product_attention) if not self.flash: self.register_buffer(bias, torch.tril(torch.ones(config.block_size, config.block_size)) .view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C x.size() q, k, v self.c_attn(x).split(self.n_embd, dim2) k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) q q.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) if self.flash: y torch.nn.functional.scaled_dot_product_attention( q, k, v, attn_maskNone, dropout_pself.dropout if self.training else 0, is_causalTrue) else: 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.attn_dropout(att) y att v y y.transpose(1, 2).contiguous().view(B, T, C) y self.resid_dropout(self.c_proj(y)) return y这里有个容易忽略的点self.c_attn一次性算出 Q、K、V 三个矩阵然后split切开。这样比三个独立 Linear 快因为一次矩阵乘法搞定。scaled_dot_product_attention是 PyTorch 2.0 引入的 Flash Attentionis_causalTrue自动加因果掩码比手动实现快很多。如果你的 PyTorch 低于 2.0会走 else 分支用torch.tril生成下三角掩码。2.3 GPT 主体与权重共享GPT 类把 embedding、Block 堆叠、LayerNorm 和 lm_head 组装起来dataclass class GPTConfig: block_size: int 1024 vocab_size: int 50304 n_layer: int 12 n_head: int 12 n_embd: int 768 dropout: float 0.0 bias: bool True class GPT(nn.Module): def __init__(self, config): super().__init__() self.config config self.transformer nn.ModuleDict(dict( wtenn.Embedding(config.vocab_size, config.n_embd), wpenn.Embedding(config.block_size, config.n_embd), dropnn.Dropout(config.dropout), hnn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_fLayerNorm(config.n_embd, biasconfig.bias), )) self.lm_head nn.Linear(config.n_embd, config.vocab_size, biasFalse) self.transformer.wte.weight self.lm_head.weight # weight tying self.apply(self._init_weights)注意vocab_size50304这是 GPT-2 的 50257 向上取整到 64 的倍数据说能提升约 30% 训练效率。weight tying让输入 embedding 和输出 lm_head 共享权重减少参数量。_init_weights用std0.02初始化残差投影c_proj.weight用0.02/sqrt(2*n_layer)缩放这是 GPT-2 论文的做法。2.4 训练循环get_batch 与梯度累积训练循环里最关键的是get_batch它从train.bin里用np.memmap随机取 batchdef get_batch(split): data np.memmap(os.path.join(data_dir, train.bin if split train else val.bin), dtypenp.uint16, moder) ix torch.randint(len(data) - block_size, (batch_size,)) x torch.stack([torch.from_numpy((data[i:iblock_size]).astype(np.int64)) for i in ix]) y torch.stack([torch.from_numpy((data[i1:i1block_size]).astype(np.int64)) for i in ix]) if device_type cuda: x, y x.pin_memory().to(device, non_blockingTrue), y.pin_memory().to(device, non_blockingTrue) else: x, y x.to(device), y.to(device) return x, yy是x右移一位对应「预测下一个 token」。np.memmap避免把整个大文件读进内存适合几 GB 的 tokenized 数据。梯度累积通过gradient_accumulation_steps模拟更大 batchscaler.scale(loss).backward()配合 fp16 混合精度。3. 中文 GPT 训练可复制配置prepare.py 与 train.py 完整片段3.1 中文语料准备字符级分词最简单的方式是复用shakespeare_char/prepare.py把input.txt换成中文文本。字符级分词词表约 4000~5000适合《红楼梦》这类单本书# 下载《红楼梦》文本重命名为 input.txt mkdir -p data/hongloumeng_char cp input.txt data/hongloumeng_char/input.txt cp data/shakespeare_char/prepare.py data/hongloumeng_char/prepare.py python data/hongloumeng_char/prepare.py运行后会生成train.bin、val.bin和meta.pkl。meta.pkl里存了stoi和itos映射推理时要用。字符级分词的好处是词表小、训练快坏处是序列长、压缩率低。3.2 训练配置train_hongloumeng_char.py新建config/train_hongloumeng_char.pyout_dir out-hongloumeng-char eval_interval 250 eval_iters 200 log_interval 10 always_save_checkpoint False wandb_log False dataset hongloumeng_char gradient_accumulation_steps 1 batch_size 32 block_size 256 n_layer 12 n_head 8 n_embd 512 dropout 0.2 learning_rate 1e-3 max_iters 5000 lr_decay_iters 5000 min_lr 1e-4 beta2 0.99 warmup_iters 100 device cuda compile False启动训练python train.py config/train_hongloumeng_char.py如果显存不够把batch_size降到 16block_size降到 128。n_layer12, n_head8, n_embd512大约 40M 参数单卡 8GB 显存能跑。3.3 基于 BBPE 的中文 tokenizer字符级分词续写效果一般可以自训练 BBPE tokenizer。用tokenizers库from tokenizers import Tokenizer, models, trainers, pre_tokenizers tokenizer Tokenizer(models.BPE()) tokenizer.pre_tokenizer pre_tokenizers.ByteLevel(add_prefix_spaceFalse) trainer trainers.BpeTrainer(vocab_size8000, special_tokens[|endoftext|]) tokenizer.train([data/hongloumeng_char/input.txt], trainer) tokenizer.save(data/hongloumeng_bbpe/tokenizer.json)然后用这个 tokenizer 把文本 encode 成 token id存成train.bin、val.bin。推理时用同一个 tokenizer decode。BBPE 词表 8000 左右压缩率比字符级高续写更流畅。3.4 基于 Qwen2 tokenizer 的尝试直接用 Qwen2 的 tokenizer 也可以但词表 15 万小语料下大部分 token 的 embedding 训练不充分容易输出乱码。我试过在《红楼梦》上训结果基本是乱码。如果要用 Qwen2 tokenizer语料至少要到亿级 token且模型要更大n_layer24, n_embd1024 以上。4. 验证请求与成功结果sample.py 推理与 TaoToken 统一 Key 接入4.1 本地推理验证训练完用sample.py续写python sample.py --out_dirout-hongloumeng-char --devicecuda --start贾宝玉见了黛玉便说如果训练充分会输出类似「贾宝玉见了黛玉便说『妹妹近日可好』黛玉笑道『好只是夜里睡不安稳。』」的续写。训练 250 step 时还是乱码2500 step 后开始通顺5000 step 后人物对话基本合理。4.2 TaoToken 统一 Key 接入推理验证本地小模型续写能力有限可以用 TaoToken 统一 Key 调用更大的模型做对比验证。TaoToken 提供统一的 API 通道兼容 OpenAI 格式Base URL 是https://taotoken.net/api。先申请 API Key访问https://taotoken.net/api-keys?utm_sourcetaotoken_aicg_blog_endutm_contentapi_keysutm_campaignrewrite创建 Key 后复制。然后用 Python 调用from openai import OpenAI client OpenAI( api_key你的TaoToken Key, base_urlhttps://taotoken.net/api ) response client.chat.completions.create( modelclaude-sonnet-4-20250514, messages[ {role: user, content: 用《红楼梦》风格续写贾宝玉见了黛玉便说} ], max_tokens200 ) print(response.choices[0].message.content)如果返回正常文本说明 Key 和通道都通了。这一步的意义是你可以用大模型生成高质量续写作为小模型训练的目标参考或者直接对比两者的风格差异。4.3 用 TaoToken 做训练数据增强TaoToken 还可以用来生成 SFT 数据。比如让大模型对《红楼梦》片段做问答对生成prompt 请根据以下《红楼梦》片段生成3组问答对格式为 JSON 片段贾宝玉见了黛玉便说『妹妹近日可好』黛玉笑道『好只是夜里睡不安稳。』 response client.chat.completions.create( modelclaude-sonnet-4-20250514, messages[{role: user, content: prompt}], max_tokens500 ) print(response.choices[0].message.content)生成的问答对可以拿去 SFT 小模型提升对话能力。TaoToken 的 Coding Plan 适合长期做这类批量生成任务访问https://taotoken.net/coding-plan?utm_sourcetaotoken_aicg_blog_endutm_contentcoding_planutm_campaignrewrite了解详情。5. 本篇常见错排查401、local proxy failed、reading choices、OAuth5.1 401 Unauthorized调用 TaoToken API 返回 401通常是 Key 没填对或没加Bearer前缀。检查client OpenAI( api_keysk-xxxxxxxx, # 不要加 BearerSDK 会自动加 base_urlhttps://taotoken.net/api )如果还是 401去https://taotoken.net/api-keys?utm_sourcetaotoken_aicg_blog_endutm_contentapi_keysutm_campaignrewrite确认 Key 是否有效、额度是否用完。5.2 local proxy failed报local proxy failed或连接超时先检查网络是否能访问https://taotoken.net/api。用 curl 测试curl -X POST https://taotoken.net/api/v1/chat/completions \ -H Authorization: Bearer sk-xxxx \ -H Content-Type: application/json \ -d {model:claude-sonnet-4-20250514,messages:[{role:user,content:hi}]}如果 curl 通但 Python 不通检查是否有环境变量HTTP_PROXY、HTTPS_PROXY干扰临时 unset 掉再试。5.3 reading choices 报错reading choices通常是响应格式不对比如用了不兼容的 SDK 版本。确保openai库版本 1.0pip install -U openai然后确认response.choices[0].message.content能正常取到。如果返回的是流式要用for chunk in response遍历。5.4 OAuth 相关报错如果用的是 Claude Code 或 Cline 这类工具报 OAuth 错误需要在工具里配置 Base URL 和 Key。以 Claude Code 为例三件套是Base URL:https://taotoken.net/apiAPI Key: 你的 TaoToken KeyModel ID:claude-sonnet-4-20250514配置好后重启工具。如果还报 OAuth检查是否误开了官方登录态清掉~/.claude下的缓存再试。5.5 训练侧常见错CUDA out of memory降batch_size或block_size或者开gradient_accumulation_steps。RuntimeError: shape mismatch检查vocab_size是否和meta.pkl里的词表大小一致。字符级分词词表约 4000~5000BBPE 约 8000Qwen2 约 15 万。loss 不下降检查learning_rate是否太大字符级分词建议 1e-3BBPE 建议 3e-4Qwen2 tokenizer 建议 1e-4。6. 从源码到跑通我的实操建议与 TaoToken 接入入口整个流程走下来最深的体会是NanoGPT 的代码虽然短但每一行都值得细看。CausalSelfAttention里的split和view/transpose操作get_batch里的memmap和右移标签GPT里的权重共享和初始化缩放这些都是 GPT 系列的通用套路。搞懂这些再看 LLaMA、Qwen 的代码会轻松很多。中文训练的关键在 tokenizer 选择。字符级分词适合快速验证BBPE 适合单本书或小语料Qwen2 tokenizer 需要大语料和大模型。我自己的经验是先用字符级跑通全流程确认 loss 能降到 1.5 以下再换 BBPE 提升续写质量。推理验证环节TaoToken 的统一 Key 通道省去了管理多个 API Key 的麻烦。无论是做数据增强、对比验证还是接入 Claude Code 做辅助编码一个 Key 就够了。模型对话入口在https://taotoken.net/chat?utm_sourcetaotoken_aicg_blog_endutm_contentmodel_chatutm_campaignrewrite接入文档在https://taotoken.net/doc?utm_sourcetaotoken_aicg_blog_endutm_contentdocutm_campaignrewrite控制台在https://taotoken.net/console?utm_sourcetaotoken_aicg_blog_endutm_contentconsoleutm_campaignrewrite。最后给一个实用技巧训练时把always_save_checkpoint False只在 val loss 下降时保存避免磁盘被 checkpoint 塞满。推理时用--start指定开头能直观对比不同 step 的续写质量。如果要做多卡训练用torchrun --standalone --nproc_per_node8 train.py config/train_gpt2.pyNanoGPT 自带 DDP 支持小模型上和 DeepSpeed 差别不大。