Apology
由于成本 / 技术 / 与主题无过大关联 等原因，我并没有手撕一些组件，列举如下：

torch numpy einops 最基础的函数，如 @ * 这样的矩阵运算、广播机制，还有 arange zeros rearrange 之类的基础数据处理函数
torch 自带的 torch.nn.functional.scaled_dot_product_attention : 我独立实现了SDPA，但是非常不幸由于不会手撕 FlashAttention，又考虑到训练 / 推理时间成本，重新用 Pytorch 优化过的注意力写了一个新的 MultiHeadAttention
CLIP-ViT : 预算实在有限，训练的成本还是太高了，只能在了解其原理、感慨 OpenAI 财大气粗之后直接使用成品了。
Info
项目开源地址：Github

模型权重开源地址：ModelScope

模型体验地址：ModelScope

Part 1. Pre-training
这部分是预训练，目的是让 AI 简单学会基础语法，能够续写语料。

1.0 开始的开始
这个项目的起源在于斯坦福的 AI 神课 CS336 。个人认为这门课的作业部分教育指导意义远大于讲课视频，Transformer、优化器等等等等都是要亲手实现才能理解，看视频、看blog、单纯调用 torch.nn.functional 库中的函数与类永远无法弥补此点。

基础 Transformer 模块中的大部分代码（即 modules.py 中的大多数模块）都是直接从我亲自跟着 Assignment 1 敲下的代码摘除或改编的。

Assignment 1 在带着我们亲手实现手搓完整的 BPE分词器、线性层、Embeddings、RMSNorm、SwiGLU、RoPE、SDPA及MHA、CE损失函数、SGD、AdamW、学习率调节器、梯度裁剪、DataLoader、ckpt设计和基础文本生成后，便让我们自己训练一个模型。然而从各个模块的分别实现到组装为一个完整的训练循环是另一码事。训练期间遇到了很多问题。

1.1 预训练的前戏
首先是训练集获取问题。显然为了模型的能力上限、应用场景考虑，不能使用简单的 TinyStories 训练集。但是由于算力、成本有限，我选择了 CS336 特供 - OpenWebText 训练集。

我先训练了一个 BPE 分词器，设定 vocab_size=32000 。

训练是很慢的，跑了半个小时发现进度很差意识到需要优化了。

BPE 训练无非两步，一步是分词，一步是合并。

分词是很快的，使用作业pdf给出的那个分块即可，没啥优化的价值。
问题就出在合并上。我最初实现的合并关键代码是：

1
2
3
4
for segment in segments:
it = re.finditer(PAT, segment)
for i in it:
list_origin_words[i.group()] += 1
这是对作业pdf原理直接的转写，但是显然对于真实实际应用场景很糟糕。

我在实现作业之初用 flameprof 记录分析了一下对 TinyStories bpe 训练过程，发现整个火焰图几乎都是 regex 相关的内容，而其他函数的占比极小：

显然，训练脚本花了大部分的时间放在 finditer() 正则匹配上面了。那么我们直接空间换时间，只需先进行一次预匹配，创一个词频表，然后就无需在每一个 merge 执行一次匹配了。

但是还有一个很大的问题出在

1
best_pair = max(pair_counts.items(), key=lambda x: (x[1], x[0]))
上面。当 pair 数量极多，会查找极慢。干脆放弃 python 直接转向 C++ ，使用 C++ 的 std::unordered_map 哈希表 ，搭配上 PyBind11 ，直接用 C++ 重写这部分。编译为动态链接库之后就能直接被我们的训练脚本当作 module 引入了。

于是，我们训练出了 vocab_size=32000 的一个 BPE 分词器，并且将词典表和合并序列直接用 pickle 保存下来。

那么下一步的当务之急是将这 12G 的纯文本 Tokenize 一下。显然，我在 CS336 中 Tokenizer 的 encode() 函数对长文本的效果很不好，于是我用 AI 快速跑了一个快速 Tokenizer，依旧使用的了 C++ 来写核心的循环合并部分（毕竟 C++ 的哈希表还是太强了），又加入了并行计算，然后就以 8 核全部 98.7% 以上榨干极致性能实现了快速的 Token 化。经过不精确估算，高性能 Tokenizer 的 encode 使 Token 化的速度提升了大约
46000
%
46000%。将训练语料 Token 化后，我用 np.uint16 格式 + torch.save 储存到了本地的 .npy 文件（def make_npy() 这块）。

然后就可以进入预训练了。

1.2 预训练
我选择的参数是：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
vocab_size = 32000
context_length = 1024
batch_size = 8
d_model = 1024
num_layers = 24
num_heads = 16
d_ff = 2752
rope_theta = 10000
max_iters = 3_000_000
max_learning_rate = 2e-4
min_learning_rate = 2e-5

weight_decay = 0.1
betas = (0.9, 0.95)
eps = 1e-8
max_grad_norm = 1.0
架构则是：

Embedding (
32000
×
1024
32000×1024 )
RoPE
24 x Transformer Block，每个 Block 有：
3.1 RMSNorm (pre‑norm)
3.2 Multi‑Head Self‑Attention (16 头，d_model=1024，4 x 线性层)
3.3 RMSNorm (pre‑norm)
3.4 SwiGLU FFN (3 x 线性层)
RMSNorm
线性层
架构之外，我使用了 AdamW 优化器，还有余弦学习率调节器，并且使用 wandb 记录数据。

这大概是一个 GPT2 - Medium 的参数，参数量大概是 370M 。根据 Chinchilla 最优法则，理想训练需要 20 倍参数量即 7B 的训练 Token 量。不过由于时间、预算有限，我 最终 只跑了
3
4
4
3
​
个 CS336 特供款 OpenWebText Sample 和大约
1
3
3
1
​
个 Wikipedia English Chunks。

然后在某炼丹平台租了一张 RTX PRO 6000 96GB，开始了预训练。

训练时发现训练的速度确实很慢，我想大概是 Scaled Dot-Product Attention 的问题。我之前学过一点点 Triton 算子，这种涉及大量张量融合计算的问题应该使用一个 fused 的版本。那么对于 SDPA 而言便是 FlashAttention V2。但是这个的实现有些偏离主题了，干脆直接拿 ‌torch.nn.functional.scaled_dot_product_attention 代替了。（请原谅我在这块的作弊 Orz）

我最初用的 OpenWebText 跑，跑了一共大概 7 个小时，分了两段。

损失变化如下：

最初这里的学习率其实并没有收敛很好（当时只训练到了橙色 pretrain_owt_p3 那里）。。。我想这大概与那时只训练了 1.6B 左右的 token 有关。

我最初并没有接着训练 Wikipedia 训练集，而是直接开了接下来的 sft 微调训练。结果在疯狂碰壁一直改了两天 bug 耗了 50 多 RMB 后发现还是预训练不充足。

于是继续训练了 wiki，保留 checkpoint 里面的模型权重但是重新初始化了 AdamW 优化器。这里补上了验证集评估的记录（训练 wiki 阶段），可见训练效果还不错。

现在一共训练了大概 3B 的 token 量。

部分测试还可以：