CS336：从零手搓LLM-预训练部分

基于 TinyStories 数据集，从 BPE Tokenizer 训练到 Transformer 文本生成，完整实现一个现代语言模型的所有基础组件。

目录

1. 项目概览

2. 第一部分：BPE Tokenizer

3. 第二部分：神经网络基础组件

4. 第三部分：Transformer 模型

5. 第四部分：训练与推理

6. 完整训练及成果展示

7. 踩坑记录与经验教训

项目概览

整个项目的文件结构如下：

assignment1-basics/
├── cs336_basics/              # 核心实现库
│   ├── module.py              # 模型组件（Linear、Embedding、Attention、Transformer 等）
│   ├── functional_imp.py      # 基础函数（softmax、SiLU、交叉熵损失、梯度裁剪）
│   ├── optimizer_imp.py       # AdamW 优化器
│   ├── run_get_batch_imp.py   # 批次数据加载
│   ├── run_train_bpe_imp.py   # BPE 训练算法
│   ├── get_tokenizer_imp.py   # Tokenizer 编码/解码
│   ├── model_imp.py           # 手动实现 forward pass（基于权重字典）
│   ├── serialization_imp.py   # 模型保存与加载
│   └── pretokenization_example.py  # 预分词示例
├── combine/                   # 流水线脚本
│   ├── config.py              # 统一配置
│   ├── train_bpe.py           # 训练 BPE tokenizer
│   ├── run_tokenize.py        # 文本 → token ID 批量转换
│   ├── train_model.py         # 模型训练主入口
│   ├── generate_text.py       # 自回归文本生成（推理）
│   └── utils.py               # tokenizer 加载工具
└── tests/                     # 单元测试（共 8 个测试文件）

第一部分：BPE Tokenizer

1.1 BPE 训练算法

Byte Pair Encoding 的核心思想是：从字节级开始，反复合并最高频的相邻 token 对，直到达到目标词表大小。

实现位于 run_train_bpe_imp.py，核心流程：

flowchart TD
    A["原始文本字节序列"] --> B["预分词 pre-tokenization<br/>用正则把文本拆成 chunks"]
    B --> C["每个 chunk 拆成单字节<br/>初始化 pair 列表"]
    C --> D["统计所有相邻 pair 的频率<br/>找出最高频 pair"]
    D --> E{"词表大小够了？"}
    E -->|否| F["合并该 pair：生成新 token<br/>加入 vocab 和 merges"]
    F --> G["更新所有 chunk 的 pair 列表"]
    G --> D
    E -->|是| H["输出 vocab + merges"]

这里有一个关键设计——预分词。用正则表达式把文本拆成独立的 chunk（单词、标点等），每个 chunk 内部独立做 BPE 合并，跨 chunk 不合并。这避免了合并跨越语义边界。

本项目用的正则与 GPT-2 一致：'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+，按 | 顺序取最长/最靠前匹配，把文本切成「字母串 / 数字串 / 标点串 / 空白串」。训练阶段（BpeTrainer）先按 special token 切分，再对每段用 re.finditer 扫描预分词；推理阶段（Tokenizer.encode）优先最长匹配留出 special token 整体，普通文本才走正则，故特殊 token 不会被拆散。例："Hello world! 123" → Hello / world / ! / 123 四个 chunk，各自 UTF-8 字节化后在内部独立合并。

1.2 训练细节

BPE 训练使用 TinyStories 的一个约 5MB 采样语料 tinystories_sample_5M.txt（而不是全量训练集，以控制训练时间）。

通过 train_bpe.py 调用训练：

vocab, merge = run_train_bpe_imp.run_train_bpe_imp(path, 1000, special_tokens)

训练产物：

• vocab.json：{token_id: bytes} 的映射，词表大小 1000

• merge.json：合并规则列表，按优先级排序 [(token_a, token_b), ...]

1.3 Tokenizer 编码与解码

实现位于 get_tokenizer_imp.py，编码过程：

1. 文本 → UTF-8 字节 → 单字节 list → 贪婪合并（按 merges 优先级）→ token IDs

2. 解码过程：token IDs → 查 vocab 得到 bytes → 拼接 → UTF-8 解码回文本

1.4 批量化处理

run_tokenize.py 将整个 TinyStories 数据集编码为 np.int16 二进制文件（int16 足够覆盖 1000 词表，比 int64 节省 4 倍空间）。编码后训练集约 1.4GB、验证集约 14MB。

第二部分：神经网络基础组件

所有组件都在 module.py 中手写实现，不调用 nn.Linear 等高级 API。

2.1 Linear（线性层）

class Linear(torch.nn.Module):
    def forward(self, x):
        return einsum(x, self.weight, '... d, o d -> ... o')

• 无偏置（bias=None），简化设计

• 权重形状 [out_features, in_features]，与 PyTorch 原生一致

• 使用 einsum 实现，... 自动适配任意 batch 维度

• 权重初始化：N(0, 0.02²)

2.2 Embedding（词嵌入）

class Embedding(torch.nn.Module):
    def forward(self, token_ids):
        return self.weight[token_ids]
# [batch, seq_len] 整数 ID → [batch, seq_len, d_model] 浮点向量

本质上是一个可学习的查找表，用 PyTorch 的高级整数索引（weight[token_ids]）实现。每个 token ID 被替换成对应行的高维向量。

2.3 RMSNorm（均方根归一化）

与 LayerNorm 的区别在于不做去均值（centering），计算量更小。LLaMA 系列模型就使用 RMSNorm。

class Rmsnorm(torch.nn.Module):
    def forward(self, x):
        x_fp32 = x.to(torch.float32)          # 转高精度防溢出
        rms = torch.sqrt(x.pow(2).sum(-1)/self.d_model + self.eps)
        return (x_fp32 / rms * self.weight).to(orig_dtype)

2.4 交叉熵损失（CrossEntropyLoss）

位于 functional_imp.py，使用了数值稳定的 log-sum-exp trick：

target_logits = torch.gather(inputs, -1, targets.unsqueeze(-1))
max_inputs = torch.max(inputs, dim=-1, keepdim=True).values
n_log = -target_logits + max_inputs + torch.log(torch.sum(exp(inputs - max_inputs)))
return n_log.mean()

核心公式：

2.5 其他辅助函数

第三部分：Transformer 模型

3.1 模型整体结构

Transformer
├── token_embeddings: Embedding(1000, 512)       # Token → 向量
├── layers: ModuleList[4 × TransformerBlock]
│   └── TransformerBlock
│       ├── ln1: RMSNorm                          # Pre-Norm ①
│       ├── attn: MultiheadAttentionWithRope      # RoPE 多头注意力
│       ├── ln2: RMSNorm                          # Pre-Norm ②
│       └── ffn: SwiGLU(d_model=512, d_ff=1344)   # SwiGLU 前馈网络
├── ln_final: RMSNorm                            # 最终归一化
└── lm_head: Linear(512, 1000)                   # 输出投影 → 词表分布

架构选择：Pre-Norm + RoPE + SwiGLU，这是现代 GPT 系列（GPT-2/3、LLaMA）的标准配置。

3.2 Multi-Head Attention with RoPE

flowchart LR
    subgraph 单个 TransformerBlock
        A["输入 [B, T, 512]"] --> B["ln1 (RMSNorm)"]
        B --> C["Q/K/V 投影 (Linear)"]
        C --> D1["Q: [B, 16, T, 32]"]
        C --> D2["K: [B, 16, T, 32]"]
        C --> D3["V: [B, 16, T, 32]"]
        D1 --> E1["RoPE 旋转编码"]
        D2 --> E2["RoPE 旋转编码"]
        E1 --> F["Q @ K^T / sqrt(32)"]
        E2 --> F
        F --> G["Causal Mask"]
        G --> H["Softmax → Attention Weights"]
        H --> I["Weighted Sum with V"]
        D3 --> I
        I --> J["Concat Heads (16×32→512)"]
        J --> K["Output Projection"]
        K --> L["+ Residual"]
    end

关键实现细节：

拆分多头：用 einops 将 [B, T, 512] 重排为 [B, 16, T, 32]

q_split = rearrange(q, "... s (h k) -> ... h s k", h=16)

Casual Mask：动态根据实际序列长度创建下三角 mask

causal = torch.tril(torch.ones(seq_len, seq_len)).bool()
scores = scores.masked_fill(~causal, float("-inf"))

合并多头：反转拆分，将 [B, 16, T, 32] 恢复为 [B, T, 512]

rearranged_head = rearrange(head, "... h s d -> ... s (h d)")

3.3 RoPE（旋转位置编码）

RoPE 的核心是对特征维度做两两配对的 2D 旋转，不同 pair 使用不同频率：

key_pair = rearrange(x, "... (k p) -> ... k p", p=2)  # 512维 → 256对×2
theta_p_i = pos × θ^(-2i/d)                            # 每对独立的旋转角度
x1, x2 = key_pair[..., 0], key_pair[..., 1]
y1 = x1 * cos - x2 * sin                               # 2D 旋转
y2 = x1 * sin + x2 * cos

• 高频对 (i 小)：旋转快，对近邻位置变化敏感

• 低频对 (i 大)：旋转慢，适合编码长距离位置关系

这样 RoPE 用旋转矩阵的性质自然地编码了绝对和相对位置信息。

3.4 SwiGLU FFN

class Swiglu:
    def forward(self, x):
        return self.w2( self.w3(x) * silu(self.w1(x)) )

公式：$\text{FFN}(x) = W_2(\text{SiLU}(W_1x) \odot W_3x)$

相比传统 ReLU FFN（两个矩阵），SwiGLU 多了一个门控矩阵 W3，用 SiLU 激活做门控，表达能力更强。

3.5 前向传播流程

[token IDs] → Embedding → [Pre-Norm → Attention → +Residual
                           → Pre-Norm → SwiGLU   → +Residual] × 4
             → Final RMSNorm → Linear → [logits]

我们还在模型层增加了输入长度自动截断保护：

def forward(self, in_indices):
    if in_indices.shape[-1] > self.context_length:
        in_indices = in_indices[..., -self.context_length:]  # 保留最后 256 token
    ...

这样调用方无需关心上下文长度限制，模型内部自动处理。

第四部分：训练与推理

4.1 优化器：AdamW

从头实现 AdamW（位于 optimizer_imp.py）：

m_t = β₁·m_{t-1} + (1-β₁)·grad           # 一阶动量
v_t = β₂·v_{t-1} + (1-β₂)·grad²          # 二阶动量
lr_t = lr · √(1-β₂^t) / (1-β₁^t)         # 偏差修正
θ -= lr_t · m_t / (√v_t + ε)             # Adam 更新
θ -= lr · weight_decay · θ               # Weight Decay（解耦，AdamW 关键改进）

超参数：lr=3e-4, β=(0.9, 0.95), weight_decay=0.01, ε=1e-8

4.2 数据加载

使用 np.memmap 以内存映射方式读取二进制的 token ID 文件——不占用内存，按需从磁盘读取：

train_token_ids = np.memmap("train.bin", dtype=np.int16, mode="r")

每个 batch 的训练步骤：

# 随机采样 32 个不重复的起始位置
begin_indexes = random.sample(range(0, n - 256), 32)
x = [dataset[i : i+256] for i in begin_indexes]      # 输入序列
y = [dataset[i+1 : i+257] for i in begin_indexes]     # 标签（右移一位）

这是标准的自回归语言建模方式：输入 tokens[0:256]，预测 tokens[1:257]。

4.3 训练循环

flowchart TD
    subgraph 每个 batch
        A["随机采样 (x, y)"] --> B["zero_grad"]
        B --> C["model.forward → logits [32,256,1000]"]
        C --> D["CrossEntropyLoss → loss"]
        D --> E["loss.backward"]
        E --> F["梯度裁剪 (max_l2_norm=0.1)"]
        F --> G["optimizer.step"]
    end
    G --> H{iter % 2000 == 0?}
    H -->|是| I["验证集评估"]
    H -->|是| J["保存 checkpoint"]

每 2000 步做一次验证集评估并保存 checkpoint，每 100 步打印训练 loss。

4.4 文本生成（推理）

generate_text.py 实现了自回归生成：

flowchart TD
    A["给定 prompt 文本"] --> B["tokenizer.encode → token IDs"]
    B --> C["model.forward → logits [T, 1000]"]
    C --> D["取最后一个位置: logits[-1, :]"]
    D --> E["temperature 缩放"]
    E --> F{使用 top_k?}
    F -->|top_k=50| G["取 top-50 概率最高 token"]
    F -->|否| H["全词表"]
    G --> I["softmax → 概率分布"]
    H --> I
    I --> J["torch.multinomial 随机采样"]
    J --> K["tokenizer.decode → 文本"]
    K --> L{"遇到结束符?"}
    L -->|否| B
    L -->|是| M["结束"]

采样策略：top-k = 50 + temperature。从概率最高的 50 个 token 中按概率随机选一个，避免低概率垃圾 token，同时保持生成多样性。

完整训练及成果展示

训练环境与配置

训练吞吐

在 mps 后端上，单个训练 step（前向 + 反向 + 梯度裁剪 + 参数更新）实测约 3.16 step/s（约 316ms/step），比纯 CPU 训练快一个数量级。

训练规模

• 训练集 token 化后约 1.4GB（int16，约 7.5 亿 token），一个 epoch 约 9.1 万步

• 验证集 token 化后约 14MB（约 734 万 token），约 866 个验证 batch

• 每 2000 步做一次验证集评估并保存 checkpoint

训练结果

完整跑完 1 个 epoch：

成果使用

终于到了激动人心的效果测试环节了，虽然不指望这个训练了几个小时的模型能吐出什么惊艳的结果，但是好歹也算一个transformer语言模型hhh，训练好的模型可以用 generate_text.py 做自回归文本生成。使用前需要把 config.py 中的 output_model_path 指向最终模型：

output_model_path = model_data_path + "/final_model.pt"

然后运行 python generate_text.py。

生成样例展示

我实际用训练完成后的 final_model.pt 跑了三组不同开头，结果如下。

样例 1：典型 TinyStories 风格开头

Prompt:

Once upon a time, there was a little rabbit who lived in a cozy burrow under a big oak tree.

模型生成：

Once upon a time, there was a little rabbit who lived in a cozy burrow under a big oak tree. Rabbit loved to hop and play all day long. One day, Rabbit found a shiny scale on the ground. Rabbit picked it up with his mouth and wanted to keep it.
But Rabbit forgot about the sharp scale. He hopped around the tree and scooped up some leaves and put them on his face. He was very careful not to drop any of the leaves. Rabbit kept hopping and playing with his rabbit, and slowly, he stopped.
After some time, Rabbit felt his tummy go funny. He hopped away, feeling very proud. Rabbit hopped back to his tree and went back inside. The forest was now full of sunshine for Rabbit to enjoy. He knew that even though the leaves and sunshine were so fun, he would always remember the fun day playing with them in the forest.

这一组效果最好，已经能生成结构完整、语气稳定、明显带有儿童故事风格的短篇故事。

样例 2：更偏想象力的月夜开头

Prompt:

On a quiet night, a little girl looked up at the moon and made a secret wish.

模型生成：

On a quiet night, a little girl looked up at the moon and made a secret wish. Under the light apoint, the caterpillar was gone forever.
The little girl was sleepy, but she knew it was going to be one morning. She went home and kept going. Eventually suddenly she fell from the sky. She was not so happy.
The bad ending of the night was that the little girl had not gone anywhere. And the moon was rather sad, leaving them all.

这一组也能继续往下写，但开始出现不太自然的短语和跳跃的叙事，说明模型对这类更抽象的开头掌握得还不够稳定。

样例 3：家庭场景开头

Prompt:

Tom and his little boat floated on the pond behind Grandma's house.

模型生成：

Tom and his little boat floated on the pond behind Grandma's house. Grandma told them that she loved them and wanted them to grow up. Grandma cooked dinner with them the time and it was time for Tom to stay in. He was happy and loved grandma very much. He said, "Thank you, Grandma! This is the best day ever and I love you more. The fish and the sea are always big and beautiful. I love you and I love your family."
Grandma smiled and said, "You are very sweet of you, Tom. I love you very much. I am glad you love me and you will always be happy." Tom said, "I love you too, Grandma. You are the best Grandma ever."
They waved goodbye to Grandma and went back to the bedroom. Then they played with their baby fish and had a good day. They were happy to have a new Grandma and to be their creatures. They were the best cooks ever.

这一组的流畅度中等，能延续出明确的人物互动，但后半段开始出现重复表达和语义漂移。

效果小结

emm这个效果确实比较抽象。生成的语句还算比较流畅，但逻辑性有点差。有种意大利面拌42号混凝土的既视感。测试中还发现发现模型有时会吐出根本不存在的单词，因为bpe分词器会把一个单词拆分得更细的片段，所以有时会组合出意想不到的组合。总之只能说图个好玩🤣