Tiny LLM Ablation Lab

一个面向大模型算法学习和面试展示的小型 Decoder-only Transformer 预训练实验仓库。项目不调用 AutoModelForCausalLM，核心结构由本仓库独立实现，并提供固定训练 token 预算的结构消融入口。

当前状态：已完成代码实现、单元测试和 CPU smoke test。尚未完成正式 TinyStories 训练，仓库中 不提供也不虚构 PPL、收敛速度或生成质量提升数据。正式结果应在确定硬件、数据子集和随机种子后 真实运行并记录。

算法范围

• RMSNorm 与 LayerNorm
• RoPE 与可学习绝对位置编码
• 参数量近似匹配的 SwiGLU 与 GELU MLP
• Multi-Head Attention 与 Grouped-Query Attention
• causal scaled dot-product attention
• Pre-Norm Transformer block、权重绑定和自回归生成
• AdamW、warmup + cosine decay、梯度累积和梯度裁剪
• TinyStories 流式数据准备、训练/验证、JSONL 指标和 checkpoint
• 固定训练 token 数的单变量消融及组合模型实验

项目结构

tiny_lm/
  config.py       # 类型化配置、校验和 token 预算
  data.py         # byte tokenizer 与 uint16 memmap dataset
  model.py        # RMSNorm、RoPE、MHA/GQA、SwiGLU、Transformer LM
  training.py     # 训练、验证、调度和日志
scripts/
  prepare_tinystories.py
  train.py
  run_ablations.py
configs/
  smoke.json
  tinystories_small.json
tests/

快速验证

在仓库根目录执行：

python scripts/prepare_tinystories.py --mode smoke --output-dir data/smoke
python -m unittest discover -s tests -v
python scripts/train.py --config configs/smoke.json
python scripts/run_ablations.py --config configs/smoke.json

第一条命令使用仓库内置的确定性短故事，仅用于检查数据到 loss、反向传播、验证和 checkpoint 的 完整链路。它不是 TinyStories 的实验结果。

准备 TinyStories

安装依赖后，可流式处理公开数据集，避免先下载全部原始数据：

pip install -r requirements.txt
python scripts/prepare_tinystories.py `
  --mode huggingface `
  --output-dir data/tinystories `
  --max-train-examples 100000 `
  --max-val-examples 5000

去掉 --max-*-examples 可遍历完整 split。输出包含 train.bin、validation.bin 和记录数据来源、 文档数、token 数及 SHA-256 的 metadata.json。当前 tokenizer 是完全可复现的 UTF-8 byte tokenizer； 它减少了外部依赖，代价是序列比 BPE 更长。正式扩展可增加 SentencePiece/BPE，并把 tokenizer 选择 作为独立消融变量。

正式训练

python scripts/train.py --config configs/tinystories_small.json

默认示例约为小型模型配置，不代表任何显卡都能直接使用相同 batch size。显存不足时先降低 batch_size，再提高 gradient_accumulation_steps，保持每个 optimizer step 的有效 token 数不变。 CPU smoke 配置使用 float32；CUDA 可根据硬件使用 bfloat16。

等预算结构消融

仅生成配置和清单：

python scripts/run_ablations.py --config configs/tinystories_small.json

真实执行全部实验：

python scripts/run_ablations.py --config configs/tinystories_small.json --execute

实验包含：

实验	相对 GPT baseline 的变化
baseline_gpt	LayerNorm + learned position + GELU + MHA
rmsnorm_only	仅替换 RMSNorm
rope_only	仅替换 RoPE
swiglu_only	仅替换 SwiGLU
gqa_only	仅将 KV heads 减半
modern_all	RMSNorm + RoPE + SwiGLU + GQA

这里的“等预算”指相同 optimizer steps、batch size、sequence length 和 gradient accumulation，因此每组 看到相同训练 token 数。不同结构的参数量和单步 FLOPs 并非完全相同，清单会单独记录参数量；正式 报告应同时展示 token-matched 和 wall-clock/compute-matched 结果，不能只选择对自己有利的口径。

建议至少运行 3 个随机种子，并报告：验证 loss/PPL、达到目标 loss 的 token 数、tokens/s、峰值显存、 参数量以及长于训练长度时的退化。每个结论都应附训练配置和原始日志。

参考与原创边界

项目在设计上参考了 nanochat 和 nanoGPT 所倡导的“小而可修改的训练闭环”，也参考了公开的 Transformer、RoPE、RMSNorm、SwiGLU 和 GQA 论文。这里没有复制两个上游仓库的完整文件或提交历史， 模型、配置、数据管线、训练循环和消融编排均为本项目的独立实现。

如果后续直接复用任何上游代码，应在对应文件保留其版权和许可证，而不能仅在简历中把 Fork 描述为 个人原创。

面试中的诚实表述

在正式实验完成前，可以表述为：

我独立实现了支持 RMSNorm、RoPE、SwiGLU 和 GQA 的小型 Decoder-only Transformer，并打通了 TinyStories 数据准备、训练验证与固定 token 预算消融框架；目前完成 CPU 链路验证，正式单卡 训练和多随机种子结果仍在进行中。

不能表述为“某结构提升了多少”，直到日志和可复现实验真实产生该结果。