从零手搓的一个 decoder-only 玩具 GPT:Python + PyTorch,不依赖任何现成的 transformer 库,注意力、MLP、LayerNorm、训练循环全是自己写的。模型小到什么程度呢——4 层、128 维、上下文 128,笔记本(M 系 Mac,走 mps)一晚上就能跑完。
为什么写这个:看再多 GPT 的讲解,都不如亲手踩一遍坑印象深刻。这回一路踩了 mask 条件写反、忘了 zero_grad()、targets 没右移、top-k 取反等一串坑(commit 记录里都有),踩完才算真的懂了。语料选的是古腾堡上的公版书《今古奇观》,文言白话小说,看小模型能学成什么样还挺有意思。
pip install -r requirements.txt
python data/prepare.py # 清洗语料 + tiktoken 分词,产出 data/train.bin / data/val.bin
python train.py # 训练 2500 步,loss 打到 logs/train.log,checkpoint 存 ckpt.pt
python sample.py --prompt 話說 --tokens 300 --temperature 0.8 --top_k 40原始语料 data/jinguguan_raw.txt 和 checkpoint ckpt.pt 没进仓库(见 .gitignore),语料自己从古腾堡下载就行,链接在下面。
- 作品:《今古奇观》(明 · 抱瓮老人 辑,公版书)
- 来源:Project Gutenberg eBook #24230,https://www.gutenberg.org/ebooks/24230
- 原始 TXT 约 2.9MB,先剃掉 Gutenberg 的版权头尾再分词
- 分词:tiktoken 的 gpt2 BPE,词表 50257
- 切分:95% / 5%,切分点对齐到段落边界,train 201 万 token,val 7.6 万 token
| 项 | 值 |
|---|---|
| 层数 n_layer | 4 |
| 头数 n_head | 4 |
| 维度 n_embd | 128 |
| 上下文 block_size | 128 |
| dropout | 0.1 |
| 词表 | 50257 |
训练:batch 32,lr 3e-4(200 步 warmup 后 cosine 衰减到 3e-5),weight decay 0.1,grad clip 1.0,2500 步。
最终 loss(完整曲线在 logs/train.log):
- train 2.70
- val 2.84
2500 步的小模型,能学到对话格式、「話說」开头、引号这些表面结构,内容就别指望通顺了。完整样例在 samples/。
temperature 0.8 / top-k 40(节选):
話說,知我個那。貴相被發公,細至相,不處高把了。 那中在馬頭還是,就便色晚,殿般活詩,莫人走景在致,卻不見來。
temperature 0.6 / top-k 10(节选):
話說。」道:「時自身詩,不是起,不得為起至個爺山等。」兒道:「那把還,就是接約了。」那是沒言聲,媒個山。
temperature 1.3 + 不做 top-k 的样例也在目录里,基本是一堆生僻字乱炖,可以用来对比感受下采样参数的影响。另外还留了一条 sample-buggy-topk-inverted.txt,那是 top-k 写反时期的产物,当 bug 纪念了。