Hotdry.
归档
ai-systems

微型 GPT 从零构建：核心训练循环、分词与单 GPU 高效推理

基于 Karpathy microGPT 与 nanoGPT，详解最小 GPT 分词、训练循环及单 GPU 推理参数，提供纯 Python 快速原型参数与监控清单。

在大型语言模型时代，从零构建微型 GPT 不仅是教育原型的最佳方式，更是理解 Transformer 本质的核心路径。Andrej Karpathy 的 microGPT 以 200 行纯 Python、无依赖实现完整 GPT 训练与推理，完美诠释了 “一切效率之外皆算法本质”。本文聚焦单一技术点：核心训练循环、分词机制与单 GPU 高效推理，给出可落地参数、阈值与清单，避免新闻复述，转向工程实践。

分词：从字符级到 BPE 的最小实现

分词是将文本转为整数序列的第一步，决定了模型输入质量与效率。microGPT 采用极简字符级 tokenizer，仅需遍历数据集收集唯一字符（如 names.txt 的 a-z + BOS），vocab_size=27（26 字母 + BOS 开始 / 结束标记）。

uchars = sorted(set(''.join(docs)))  # 唯一字符
BOS = len(uchars)
vocab_size = len(uchars) + 1

每个文档如 “emma” 转为 [BOS, e, m, m, a, BOS]。这种 char-level 简单、直观，适合 tiny 数据集（32k names），训练无需外部工具。但生产中用 BPE（如 nanoGPT 的 tiktoken GPT-2）更高效，vocab_size~50k，减少序列长。

落地参数：

  • Char-level：vocab_size ≤ 100，适合 <1M tokens 数据。
  • BPE：tiktoken.encode，max_token_len=1024（nanoGPT block_size）。
  • 监控：token 分布均匀度 >95%，OOV <1%；否则扩 vocab。

风险：char-level 序列长，内存 / 时间翻倍；回滚至预训练 tokenizer。

模型核心：~4k 参数的 GPT-2 简化版

microGPT 模型 stateless 函数 gpt (token_id, pos_id, keys, values)，输出 logits。架构：token/pos embedding + n_layer=1 块（attn + MLP）+ lm_head。关键简化：RMSNorm（无 bias）、ReLU（替 GeLU）、无 RoPE（学 pos emb）。

  • embd=16，head=4（head_dim=4），总 params=4192。
  • Attn：QKV proj → scaled dot-product → softmax → weighted V → Wo。
  • MLP：fc1(4x embd) → ReLU → fc2。
  • Residual + RMSNorm 稳定梯度。

nanoGPT 扩展至 GPT-2：n_layer=6/12，embd=768/1024，head=12/16，支持预训权重加载。

可调参数清单（单 GPU 友好）：

参数 microGPT nanoGPT tiny (Shakespeare) 作用
n_embd 16 384 特征维，↑内存 / 算力 x^2
n_head 4 6 并行关注，head_dim=embd/head
n_layer 1 6 深度，↑易过拟合
block_size 16 256 上下文长，↑KV cache
dropout 0 0.2 规约

初始化：gauss (0, 0.08)，std 随规模调（小模型 0.02-0.1）。

核心训练循环：前向 - 损失 - 反向 - 优化

训练是 causal LM：预测下个 token。microGPT 单步：选 doc → tokenize → seq 前向（逐 token 建 KV cache）→ CE loss avg → backward → Adam。

关键代码骨架：

  1. 前向：逐 pos，gpt (token, pos, keys, values) 累积 KV。
  2. 损失：-log(probs[target])，avg over seq。
  3. 反向：Value.backward () 拓扑逆序链式法则。
  4. 优化：Adam (β1=0.85, β2=0.99, eps=1e-8)，lr=1e-2 * decay (线性至 0)。

nanoGPT train.py ~300 行：DataLoader (bin 文件) → batch → forward → CrossEntropy → AdamW → LR schedule (cos/decay) → checkpoint。

单 GPU 参数 / 阈值：

  • batch_size=64 (A100) / 12 (CPU)，grad_acc=4 模拟大 batch。
  • lr=6e-4，warmup=10%，decay_iters=max_iters。
  • iters=5000 (tiny Shakespeare，loss~1.47)，eval_interval=100，clip_grad=1.0。
  • 监控：train/val loss Δ<0.01 停训；grad norm>10 降 lr 50%。

示例命令（nanoGPT Shakespeare char）：

python train.py config/train_shakespeare_char.py --device=cuda --batch_size=64 --block_size=256 --n_layer=6

3min A100 达合理 loss。RTX 3090 ~10min。

风险：梯度爆炸（clip），欠拟（↑iters/LR）；单 GPU 慢，回滚预训 finetune (lr=1e-5, iters=1000)。

单 GPU 高效推理：采样 + KV Cache

推理 autoregressive：从 BOS 采样至 EOS/max_new_tokens。microGPT：

  • logits / temp → softmax → choices(prob)。
  • temp=0.5（保守），top_k=40 防乱码。

nanoGPT sample.py：--start="prompt" --num_samples=10 --max_new_tokens=500 --temp=0.8。

效率参数（RTX 3090/24GB）：

  • max_new=512，batch=16 并行样本。
  • KV cache：逐 token append，避免重算 past。
  • 优化：torch.compile (PT2)，bf16 (内存 - 50%)，top_p=0.9 (核采样)。

生成 Shakespeare：

python sample.py --out_dir=out-shakespeare-char --device=cuda --max_new_tokens=200

输出连贯剧本片段。

监控：perplexity<5 质量好；latency<50ms/token (3090)。

快速上手清单：

  1. microGPT：下载 gist，python microgpt.py（1min CPU，loss~2.37，生成 names）。
  2. nanoGPT：git clone，pip install，data/shakespeare_char/prepare.py → train → sample。
  3. 自定义：换数据集（txt→bin），调 config.yaml (layer=4, embd=256)，RTX 3090 ~1h。
  4. 回滚：HuggingFace GPT2-small finetune。

这些参数经 HN 验证（3090 训 tiny 高效，租云～300$ full GPT2）。微型 GPT 非玩具，而是解锁自定义原型的关键，参数微调即见效。

资料来源：

查看归档