斯坦福大学 | CS336 | 从零开始构建语言模型 | Spring 2025 | 笔记 | Assignment 1: BPE Tokenizer

目录

• • 前言
  • [1. 作业总览](#1. 作业总览)
  • [2. Byte-Pair Encoding (BPE) Tokenizer 作业要求](#2. Byte-Pair Encoding (BPE) Tokenizer 作业要求)
  • • [2.1 The Unicode Standard](#2.1 The Unicode Standard)
    • • [Problem (unicode1): Understanding Unicode (1 point)](#Problem (unicode1): Understanding Unicode (1 point))
    • [2.2 Unicode Encodings](#2.2 Unicode Encodings)
    • • [Problem (unicode2)：Unicode Encodings (3 points)](#Problem (unicode2)：Unicode Encodings (3 points))
    • [2.3 Subword Tokenization](#2.3 Subword Tokenization)
    • [2.4 BPE Tokenizer Training](#2.4 BPE Tokenizer Training)
    • [2.5 Experimenting with BPE Tokenizer Training](#2.5 Experimenting with BPE Tokenizer Training)
    • • [Problem (train_bpe): BPE Tokenizer Training (15 points)](#Problem (train_bpe): BPE Tokenizer Training (15 points))
      • [Problem (train_bpe_tinystories): BPE Training on TinyStories (2 points)](#Problem (train_bpe_tinystories): BPE Training on TinyStories (2 points))
      • [Problem (train_bpe_expts_owt): BPE Training on OpenWebText (2 points)](#Problem (train_bpe_expts_owt): BPE Training on OpenWebText (2 points))
    • [2.6 BPE Tokenizer: Encoding and Decoding](#2.6 BPE Tokenizer: Encoding and Decoding)
    • • [2.6.1 Encoding text](#2.6.1 Encoding text)
      • [2.6.2 Decoding text](#2.6.2 Decoding text)
      • [Problem (tokenizer): Implementing the tokenizer (15 points)](#Problem (tokenizer): Implementing the tokenizer (15 points))
    • [2.7 Experiments](#2.7 Experiments)
    • • [Problem (tokenizer_experiments): Experiments with tokenizers (4 points)](#Problem (tokenizer_experiments): Experiments with tokenizers (4 points))
  • 结语
  • 参考

前言

本篇文章记录 CS336 作业 Assignment 1: Basics 中的 BPE Tokenizer 作业要求，仅供自己参考😄

Assignment 1 ：https://github.com/stanford-cs336/assignment1-basics

reference ：https://chatgpt.com/

1. 作业总览

以下内容均翻译自 cs336_spring2025_assignment1_basics.pdf，请大家查看原文档获取更详细的内容

在本次作业中，你将 从零开始构建训练一个标准的 Transformer Language Model（LM）所需的全部组件，并完成若干模型训练与评估任务

这次作业的目标并不是简单"调用现成模块"，而是帮助你真正理解一个现代语言模型在工程与算法层面是如何被一步步搭建出来的

What you will implement

你将亲手实现以下核心组件：

1. Byte-Pair Encoding（BPE）分词器（§2）
2. Transformer 语言模型（§3）
3. 交叉熵损失函数（Cross-Entropy Loss）与 AdamW 优化器（§4）
4. 训练循环（training loop） ，并支持
   • 模型参数的序列化（save）
   • 模型与优化器状态的加载（load）（§5）

What you will run

完成实现后，你将进行以下实验流程：

1. 在 TinyStories 数据集 上训练一个 BPE tokenizer
2. 使用训练好的 tokenizer 将数据集转换为 整数 ID 序列
3. 在 TinyStories 数据集上训练一个 Transformer LM
4. 使用训练好的模型：
   • 生成文本样本（sampling）
   • 计算并评估 perplexity
5. 在 OpenWebText 数据集上训练模型，并将你获得的 perplexity 提交到课程 leaderboard

What you can use

本作业 明确要求你从零实现这些组件，因此对可使用的 PyTorch API 有严格限制。

❌ 不允许使用的内容

• torch.nn
• torch.nn.functional
• torch.optim

✅ 允许使用的内容

• torch.nn.Parameter
• torch.nn 中的 容器类（container classes） ，例如：
  • Module
  • ModuleList
  • Sequential
• torch.optim.Optimizer 基类

Note ：完整的 container 类列表可参考：https://pytorch.org/docs/stable/nn.html#containers

除此之外，你可以使用 任何其他 PyTorch 提供的功能 。如果你不确定某个函数或类是否允许使用，可以在 Slack 上提问，一个简单的判断原则是： 使用它是否违背了 "from-scratch" 的设计初衷？

Statement on AI tools

课程 允许 使用 LLM（如 ChatGPT）来：

• 询问 低层次编程问题
• 讨论 语言模型相关的高层概念

但 明确禁止：

• 直接让 AI 替你完成作业代码或解答

此外，我们 强烈建议 在完成作业时：

• 关闭 IDE 中的 AI 自动补全功能，例如：Cursor Tab、GitHub Copilot
• 普通的非 AI 自动补全（如函数名补全）是可以的

课程团队的经验是：AI 自动补全会显著降低你对作业内容的深入理解程度

What the code looks like

作业的所有代码和本说明文档均位于 GitHub：https://github.com/stanford-cs336/assignment1-basics，请先 git clone 该仓库

仓库结构说明：

1. cs336_basics/*
   • 这是你编写代码的地方
   • 目录中 没有任何初始代码
   • 你可以完全自由地从零实现
2. adapters.py
   • 定义了一组 你必须提供的接口功能
   • 对于每个功能（例如 scaled dot-product attention）
     • 你只需要在对应函数中 调用你自己实现的代码
   • ⚠️ 注意：
     • adapters.py 不应包含任何实质性算法逻辑
     • 它只是 glue code
3. test_*.py
   • 包含你必须通过的所有测试
   • 测试会调用 adapters.py 中定义的接口
   • 请勿修改测试文件

How to submit

你需要向 Gradescope 提交以下内容：

• writeup.pdf
  • 回答所有书面问题
  • 请使用排版良好的 PDF（LaTeX / Markdown 转 PDF 均可）
• code.zip
  • 包含你编写的全部代码

若你希望参与课程的 perplexity 排行榜，请向以下仓库提交 PR：https://github.com/stanford-cs336/assignment1-basics-leaderboard，具体提交流程请参考该仓库中的 README.md。

Where to get datasets

本作业使用两个 已预处理好的纯文本数据集：

• TinyStories [Eldan and Li, 2023]
• OpenWebText [Gokaslan et al., 2019]

二者均为单个大型 plaintext 文件。

• 如果你在课程服务器上：
  • 数据位于任意非 head node 的 /data 目录
• 如果你在本地完成作业：
  • 可使用 README.md 中提供的命令下载

---

Low-Resource / Downscaling Tips: Init

在整个课程作业说明中，我们都会提供一些建议，帮助你在 CPU 资源较少甚至没有 GPU 的情况下 完成作业的部分内容。举例来说，我们有时会建议你对数据集或模型规模进行 下采样（downscaling） ，或者说明如何在 macOS 集成 GPU 或 CPU 上运行训练代码

你会发现，这类"低资源提示"通常以 蓝色方框 的形式出现（就像这里这样）。即使你是斯坦福在读学生，能够使用课程提供的计算资源，这些提示也依然可以帮助你更快迭代、节省时间，因此我们强烈建议你认真阅读。

---

---

Low-Resource / Downscaling Tips: Assignment 1 on Apple Silicon or CPU

使用课程工作人员提供的参考实现代码，我们可以在 Apple M3 Max（36GB 内存） 的芯片上，在 不到 5 分钟 内通过 Metal GPU（MPS） 训练出一个能够生成相对流畅文本的小型语言模型；而如果仅使用 CPU ，训练时间大约为 30 分钟。

如果你对这些硬件名词并不熟悉，也无需担心。只要你拥有一台 配置较新的笔记本电脑 ，并且你的实现是 正确且高效的 ，你就能够训练出一个可以较为流畅地生成 简单儿童故事 的小型语言模型。

在作业的后续部分中，我们还会进一步说明：如果你使用的是 CPU 或 MPS，需要对实现做出哪些相应调整。

---

2. Byte-Pair Encoding (BPE) Tokenizer 作业要求

以下内容均翻译自 cs336_spring2025_assignment1_basics.pdf，请大家查看原文档获取更详细的内容

在作业的第一部分，我们将训练并实现一个 byte-level 的 Byte-Pair Encoding（BPE）tokenizer [Sennrich+ 2016] [Wang+ 2019]，核心思路是：我们首先将任意（Unicode）字符串表示为 字节序列（bytes），然后在这些字节上训练 BPE tokenizer，之后，该 tokenizer 会被用于将文本（字符串）编码为 token ID 序列，用于语言模型训练

2.1 The Unicode Standard

Unicode 是一种将字符映射为整数 码点（code points） 的文本编码标准，截至 Unicode 16.0 （发布于 2024 年 9 月 ），该标准共定义了 154,998 个字符，覆盖 168 种书写系统（scripts）

例如，字符 "s" 的码点是 115 ，通常记作 U+0073 ，其中 U+ 是约定俗成的前缀，而 0073 是其十六进制表示，字符 "牛" 的码点是 29275

在 Python 中，你可以使用 ord() 函数将单个 Unicode 字符转换为其对应的整数码点，而 chr() 函数则可以将一个整数码点转换为对应的 Unicode 字符串

>>> ord('牛')
29275
>>> chr(29275)
'牛'

Problem (unicode1): Understanding Unicode (1 point)

(a) chr(0) 返回的是哪个 Unicode 字符？

Deliverable：一句话回答。

(b) 该字符的字符串表示（__repr__()）与其打印结果（print）有什么不同？

Deliverable：一句话回答。

(c) 当该字符出现在文本中时会发生什么？

你可以在 Python 解释器中尝试以下代码，看看是否符合你的直觉：

>>> chr(0)
>>> print(chr(0))
>>> "this is a test" + chr(0) + "string"
>>> print("this is a test" + chr(0) + "string")

Deliverable：一句话回答。

2.2 Unicode Encodings

虽然 Unicode 标准定义了从字符到码点（整数）的映射关系，但直接在 Unicode 码点上训练分词器在实际中并不可行，因为这样得到的词表会异常庞大（大约 15 万个词表项），而且非常稀疏（许多字符出现频率极低）

因此，我们通常会使用 Unicode 编码方式 ，将一个 Unicode 字符转换为一串字节序列。Unicode 标准本身定义了三种编码方案：UTF-8、UTF-16 和 UTF-32 ，其中 UTF-8 是互联网上的主流编码方式（超过 98% 的网页 使用 UTF-8）。

在 Python 中，可以使用 encode() 函数将一个 Unicode 字符串编码为 UTF-8；如果想查看 bytes 对象底层的字节值，可以对其进行迭代（例如调用 list()）。最后，可以使用 decode() 函数将 UTF-8 字节串解码回 Unicode 字符串。

>>> test_string = "hello! こんにちは!"
>>> utf8_encoded = test_string.encode("utf-8")
>>> print(utf8_encoded)
b'hello! \xe3\x81\x93\xe3\x82\x93\xe3\x81\xab\xe3\x81\xa1\xe3\x81\xaf!'
>>> list(utf8_encoded)
[104, 101, 108, 108, 111, 33, 32, 227, 129, 147, ...]
>>> len(test_string)
13
>>> len(utf8_encoded)
23

通过将 Unicode 码点序列转换为字节序列（例如使用 UTF-8 编码 ），我们实际上是把一个由 0 到 154,997 范围内整数（码点）组成的序列，映射为一个由 0 到 255 范围内整数（字节值）组成的序列，长度为 256 的字节级词表在实际处理中要更易于管理。

在采用 字节级分词 时，我们无需担心词表外（out-of-vocabulary, OOV）token 的问题，因为任何输入文本都可以表示为一个由 0 到 255 之间整数构成的字节序列

Problem (unicode2)：Unicode Encodings (3 points)

(a) 为什么在 tokenizer 训练中更倾向使用 UTF-8，而不是 UTF-16 / UTF-32？

Deliverable：1--2 句话说明。

(b) 下列 UTF-8 解码函数是错误的，为什么？请给出一个会产生错误结果的输入示例，并解释原因。

def decode_utf8_bytes_to_str_wrong(bytestring: bytes):
    return "".join([bytes([b]).decode("utf-8") for b in bytestring])

>>> decode_utf8_bytes_to_str_wrong("hello".encode("utf-8"))
'hello'

Deliverable：示例 + 1 句话说明。

(c) 给出一个 无法解码为任何 Unicode 字符的 2-byte 序列。

Deliverable：示例 + 1 句话说明。

2.3 Subword Tokenization

虽然字节级（byte-level）分词可以缓解词级（word-level）分词器所面临的词表外（out-of-vocabulary, OOV）问题，但文本直接分解为字节会 导致输入序列极其冗长，这会拖慢模型训练速度，因为在模型的每一步计算中都需要处理更长的序列

例如，在词级语言模型中，一个包含 10 个单词的句子可能只有 10 个 token；但在字符级（character-level）模型中，这个句子可能会变成 50 个甚至更多的 token（具体取决于单词长度）。处理这些更长的序列会显著增加每一步的计算量，此外，在字节级序列上进行语言建模也更加困难，因为更长的输入序列会在数据中引入更强的长期依赖关系。

子词分词（subword tokenization）位于词级（word-level）分词和字节级（byte-level）分词之间，是二者的折中方案。需要注意的是，字节级分词器的词表大小是 256（对应字节值 0-255），子词分词器通过使用更大的词表来换取对输入字节序列更好的压缩结果。举例来说，如果字节序列 b'the' 在原始数据中频繁出现，那么在词表中为它分配一个条目，就可以将原本长度为 3 的 token 序列压缩成一个单独的 token。

那么，我们该如何选择要加入词表的这些子词单元呢？[Sennrich+ 2016] 提出 字节对编码 （Byte Pair Encoding, BPE [Gage 1994]），这是一种压缩算法，它通过迭代地将出现频率最高地字节对 "合并"（merge）成一个新的、未使用过的索引。该算法不断向词表中添加子词 token，以最大化输入序列的压缩率 -- 如果某个词在输入文本中出现得足够频繁，它最终就会被表示为一个单独的子词单元。

通过 BPE 构建词表的子词分词通常被称为 BPE 分词器 ，在本次作业中，我们将实现一个 字节级 BPE 分词器，其中词表项可以是单个字节，也可以是由多个字节合并而成的序列，这种方式在处理词表外（OOV）问题和保持输入序列长度可控之间取得了良好的平衡，构建 BPE 分词器词表的过程通常被称为对 BPE 分词器进行 "训练"

2.4 BPE Tokenizer Training

BPE 分词器的训练过程主要由于三个步骤组成：

Step 1: Vocabulary Initialization.

分词器的词表是从 字节字符串 token 到整数 ID 的一一映射，由于我们训练的是一个字节级 BPE 分词器，初始词表就是所有可能字节的集合，因为字节共有 256 种取值，所以初始词表的大小为 256

Step 2: Pre-tokenization.

在拥有词表之后，原则上你可以统计文本中哪些字节经常相邻出现，并从出现频率最高的字节对开始进行合并，然而，这样做的计算开销非常大，因为每一次合并都需要对整个语料进行一次完整遍历。此外，直接在整个语料上合并字节，可能会生成一些仅在标点符号上有所差异的 token（例如 dog! 和 dog.），尽管它们在语义上高度相似（仅标点符号不同），但却会被分配为完全不同的 token ID

为了解决这个问题，我们会先对语料进行 pre-tokenize （预分词），你可以将预分词理解为一种粗粒度的分词方式，它有助于统计字符对的共现频率。举例来说，单词 "text" 在语料中可能出现了 10 次，在这种情况下，当我们统计字符 't' 和 'e' 的相邻出现频率时，就可以直接将计数增加 10，而不需要逐字节地扫描整个语料。由于我们训练的是字节级 BPE 模型，每个预分词都会被表示为一段 UTF-8 字节序列

[Sennrich+ 2016] 提出的原始 BPE 实现中，预分词方式只是简单地按空白字符进行分割即 s.split(" ")，相比之下，我们将使用一种 基于正则表达式的预分词器 （GPT-2 所采用的方法，[Radford+ 2019]）

本作业使用 GPT-2 风格的 regex pre-tokenizer：

>>> PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+|?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

代码来自：https://github.com/openai/tiktoken/pull/234/files

为了更直观地理解这个预分词的行为，可以通过交互方式对一些文本进行分割，例如：

>>> # requires `regex` package
>>> import regex as re
>>> re.findall(PAT, "some text that i'll pre-tokenize")
['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

在代码实现中，你应当使用 re.finditer，以避免在构建预分词到频次统计映射时，将所有预分词结果一次性存储在内存中

Step 3: Compute BPE merges.

在将输入文本转换为预分词，并将每个预分词表示为 UTF-8 字节序列之后，我们就可以计算 BPE 合并操作，也就是训练 BPE 分词器

从整体流程上看，BPE 算法会迭代地统计所有字节对的出现频率，并找出频率最高的一对，例如 ("A", "B")，该最高频字节对的每一次出现都会在合并步骤中被替换为一个新的 token "AB"，这个新生成的 token 会被加入到词表中，因此，BPE 训练完成后的最终词表大小等于初始词表大小（本例中为 256）加上训练过程中执行的 BPE 合并次数

为了提高 BPE 训练过程的效率，我们不会考虑跨越预分词边界的字节对

Note ：原始的 BPE 定义 [Sennrich+ 2016] 中包含了词尾标记（end-of-word token），在字节级 BPE 模型中我们不再显式加入词尾标记，因为所有字节（包括空格和标点符号）本身就已经包含在模型词表中，由于我们显式表示了空格和标点，学习到的 BPE 规则自然会反映出词边界信息

当存在多个字节对具有相同的最高频率时，我们会通过词典序（lexicographically）比较来打破并列情况，选择 词典序更大的字节对 进行合并。例如，如果字节对 ("A", "B")、("A", "C")、("B", "ZZ") 以及 ("BA", "A") 都具有最高频率，那么我们会选择合并 ("BA", "A")

>>> max([("A", "B"), ("A", "C"), ("B", "ZZ"), ("BA", "A")])
('BA', 'A')

Special tokens（特殊 token）

在实际应用中，常常会使用某些特殊字符串（例如 <|endoftext|>）来编码元信息（如文档之间的边界），在对文本进行编码时，通常希望将这些字符串视为特殊 token，它们不应被拆分成多个 token，而是始终作为一个整体保留下来。例如，序列结束标记 <|endoftext|> 必须始终作为一个单独的 token（即一个整数 ID），这样我们才能明确知道模型何时应当停止生成文本。因此，这些 token 必须显式地加入词表，并且它们拥有 固定且唯一的 token ID

[Sennrich+ 2016] 在其论文中的 Algorithm 1 给出了一个 BPE 分词器训练的实现，本质上遵循了我们前面介绍的推理步骤，但该实现并不高效，作为一个练习，实现并测试该函数有助于加深你对 BPE 分词器训练流程的理解

---

Example (bpe_example): BPE training example

下面是一个改写自 [Sennrich+ 2016] 的示例，假设我们有如下语料

low low low low low
lower lower widest widest widest
newest newest newest newest newest newest

并且词表中包含一个特殊 token：<|endoftext|>

Vocabulary Initialization

我们将词表初始化为包含特殊 token <|endoftext|> 以及所有 256 个可能的字节值

Pre-tokenization

为了简化示例并将注意力集中在合并过程上，这里假设预分词仅按照空白字符进行切分，对语料进行预分词并统计频次后，可以得到如下频率表：

{ low: 5, lower: 2, widest: 3, newest: 6 }

为了方便实现，我们可以将上述统计结果表示为一个字典 dict[tuple[bytes], int]，例如：

{ (l, o, w): 5, ... }

需要注意的是，在 Python 中，即便是单个字节，其类型也是 bytes 对象，Python 并不存在专门表示单字节的类型，正如它也没有单独的 char 类型来表示单个字符

Merges

接下来，我们遍历所有相邻的字节对，并统计它们在各个词中出现的频率，统计结果如下：

(lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 9, st: 9, ne: 6, ew: 6)

其中，字节对 ('es') 和 ('st') 具有相同的最高频率，根据规则，我们选择 词典序更大的字节对 ，即 ('st') 作为本轮要合并的对象

合并之后，预分词序列将变为：

{ (l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,e,st): 3, (n,e,w,e,st): 6 }

在第二轮合并中，(e, st) 成为最常见的字节对（出现次数为 9），于是将其合并为：

{ (l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,est): 3, (n,e,w,est): 6 }

不断重复这一过程，最终得到的合并序列为：

['s t', 'e st', 'o w', 'l ow', 'w est', 'n e', 'ne west', 'w i', 'wi d', 'wid est', 'low e', 'lowe r']

如果我们执行 6 次合并操作，那么最终得到的合并规则为：

['s t', 'e st', 'o w', 'l ow', 'w est', 'n e']

此时，词表中的元素包括：

[<|endoftext|>, [...256 个字节字符], st, est, ow, low, west, ne]

在使用该词表进行分词时，单词 newest 将被分解为：

[ne, west]

---

2.5 Experimenting with BPE Tokenizer Training

我们将在 TinyStories 数据集上训练一个 字节级 BPE 分词器，有关如何查找和下载该数据集的说明可以在第 1 小节中找到，在开始之前，我们建议你先浏览一下 TinyStories 数据集，以对其中的内容有一个整体认识

Parallelizing pre-tokenization（并行化预分词）

你会发现，预分词步骤是整个流程中的一个主要性能瓶颈，你可以通过 Python 内置的 multiprocessing 库对代码进行并行化，从而加速预分词过程。

具体来说，我们建议在并行实现预分词时，对语料进行 分块（chunking），并确保每个块的边界都处于一个特殊 token 的起始位置，你可以使用下面链接中提供的起始代码（保持原样使用）来获取这些边界，然后将工作分配到不同的线程中：

https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/pretokenization_example.py

这种分块方式始终是安全的，因为我们 绝不会再文档边界之间进行合并 ，在本次作业中，你可以始终采用这种切分策略，不必担心语料中不存在 <|endoftext|> 这一极端情况

Removing special tokens before pre-tokenization（在预分词前移除特殊 token）

在使用正则模式进行预分词（例如使用 re.finditer）之前，你应当从语料（或在并行实现中，从每个语料块）中 移除所有特殊 token ，同时，要确保你是 按特殊 token 本身进行切分的，这样就不会在文本分隔符的两侧发生合并

例如，如果你的语料（或语料块）是：

[Doc 1]<|endoftext|>[Doc 2]

你应当先按特殊 token <|endoftext|> 进行切分，并分别对 [Doc 1] 和 [Doc 2] 进行预分词，从而保证不会在文档边界之间发生合并

这一过程可以通过 re.split 来完成，并使用 "|".join(special_tokens) 作为分隔符，需要注意对正则中的 | 字符进行正确转义，因为该字符本身也可能出现在特殊 token 中，测试用例 test_train_bpe_special_tokens 会对此进行验证

Optimizing the merging step（优化合并步骤）

在前面的示例中，BPE 训练的朴素实现速度较慢，这是因为每一次合并操作都需要遍历所有字节对，以找出出现频率最高的那一对

然而，在一次合并之后，真正发生变化的计数只包括那些 与被合并字节对发生重叠的字节对 ，因此，可以通过为所有字节对建立索引，并在合并后对相关计数进行 增量更新，而不是每次都显式地重新统计所有字节对，从而显著提升 BPE 训练速度

通过这种缓存与增量更新的策略，可以获得明显的性能提升，不过需要注意的是，在 Python 中，BPE 训练中的合并步骤本身是无法并行化的

---

Low-Resource/Downscaling Tip: Profiling

你应当使用诸如 cProfile 或 scalene 等性能分析工具来识别你在实现过程中存在的性能瓶颈，并将优化工作重点放在这些瓶颈上

---

---

Low-Resource/Downscaling Tip: "Downscaling"

与其一开始就在完整的 TinyStories 数据集上训练分词器，我们更推荐你先在数据的一个小子集上进行训练，即所谓的 "调试数据集（debug dataset）"

例如，你可以先在 TinyStories 的验证集上训练分词器，该数据集的规模约为 22K 文档 ，而不是完整数据集的 212 万文档 。这体现了一种通用的开发策略：在条件允许的情况下，通过下采样来加速开发流程，例如使用更小的数据集、更小的模型规模等

在选择调试数据集或超参数配置的规模时，需要仔细权衡：一方面，调试数据集应当足够大，以便暴露与完整配置相同的性能瓶颈（这样你所做的优化才能在完整训练中同样生效）；另一方面，它又不能大到导致运行时间过长，从而影响开发效率

---

Problem (train_bpe): BPE Tokenizer Training (15 points)

Deliverable ：请编写一个函数：给定一个输入文本文件的路径，用于训练一个 字节级 BPE 分词器，你的 BPE 训练函数至少需要支持以下输入参数"

• input_path: str：指向包含 BPE 分词器训练数据的文本文件路径
• vocab_size: int：一个正整数，用于指定最终词表的最大大小，包括初始字节词表、合并过程中生成的词表项以及所有特殊 token
• special_tokens: list[str]：需要加入词表的字符串列表，这些特殊 token 不会以其他方式影响 BPE 的训练过程

你的 BPE 训练函数应当返回训练得到的 词表 和 合并规则：

• vocab: dict[int, bytes]：分词器的词表，一个从 int（词表中的 token ID）到 bytes（token 对应的字节序列）的映射
• merges: list[tuple[bytes, bytes]]：训练过程中产生的 BPE 合并列表，列表中的每一项是一个由 bytes 组成的二元组 (<token1>, <token2>)，表示 <token1> 与 <token2> 被合并，合并规则应当按照 创建顺序 排列

为了使用我们提供的测试用例来验证你的 BPE 训练函数，你需要先实现测试适配器 [adapters.run_train_bpe]，然后运行：

uv run pytest tests/test_train_bpe.py

你的实现应当能够通过所有测试

作为可选项（这可能需要投入较多时间），你可以使用系统编程语言（例如 C++，可考虑使用 cppyy 或 Rust，使用 PyO3）来实现训练流程中的关键部分，如果你选择这样做，需要注意哪些操作涉及从 Python 内存中复制数据，哪些操作是直接读取；同时请确保提供正确的构建说明，或者保证项目可以通过 pyproject.toml 直接构建

此外需要注意的是，GPT-2 所使用的正则表达式 在大多数正则引擎中支持并不好，并且在这些引擎中运行速度通常也较慢，我们已经验证过 Oniguruma 在这方面表现尚可，并且支持负向前瞻（negative lookahead），而 Python 中的 regex 包在性能上甚至更快

Problem (train_bpe_tinystories): BPE Training on TinyStories (2 points)

(a) 在 TinyStories 数据集上训练一个 字节级 BPE 分词器 ，最大词表大小设为 10,000 ，请务必将 TinyStories 的特殊 token <|endoftext|> 加入词表中，将训练得到的词表和合并规则序列化保存到磁盘中，以便后续检查

请回答以下问题：

• 训练过程大约耗时了多少小时？占用了多少内存？
• 词表中最长的 token 是什么？这个结果是否合理？

Resource requirements：≤ 30 分钟（不使用 GPU），≤ 30GB 内存

Hint ：如果在预分词阶段使用 multiprocessing，你应当能够在 2 分钟以内 完成 BPE 训练，可以利用以下两个事实：

• (a) <|endoftext|> token 在数据文件中用于分隔不同文档；
• (b) <|endoftext|> token 在 BPE 合并开始之前会作为特殊情况单独处理。

Deliverable：用 1-2 句话作答。

(b) 对你的代码进行性能分析（profiling），在整个分词器训练过程中，哪一部分耗时最多？

Deliverable：用 1-2 句话作答。

Problem (train_bpe_expts_owt): BPE Training on OpenWebText (2 points)

接下来，我们将尝试在 OpenWebText 数据集上训练一个字节级 BPE 分词器，和之前一样，建议你先浏览数据集内容，以更好地理解其中的数据分布

(a) 在 OpenWebText 数据集上训练一个 字节级 BPE 分词器 ，最大词表大小设为 32,000，将训练得到的词表和合并规则序列化保存到磁盘中，以便后续检查

请回答以下问题：词表中最长的 token 是什么？这个结果是否合理？

Resource requirements：≤ 12 小时（不使用 GPU），≤ 100GB 内存

Deliverable：用 1-2 句话作答。

(b) 比较并分析你在 TinyStories 与 OpenWebText 上训练得到的分词器之间的异同

Deliverable：用 1-2 句话作答。

2.6 BPE Tokenizer: Encoding and Decoding

在作业的前一部分中，我们已经实现了一个函数，用于在输入文本上训练 BPE 分词器，从而得到分词器的词表以及一系列 BPE 合并规则，现在，我们将实现一个 BPE 分词器，它可以 加载已有的词表和合并规则，并利用这些信息在文本与 token ID 之间进行编码与解码。

2.6.1 Encoding text

使用 BPE 对文本进行编码的过程与训练 BPE 词表的流程高度相似，主要包括以下几个步骤：

Step 1: Pre-tokenize.

首先对输入序列进行预分词，并将每个预分词表示为一段 UTF-8 字节序列，就像在 BPE 训练阶段所做的那样。随后，我们会在每个预分词内部，将这些字节合并成词表中的 token，每个预分词都是 独立处理的，不会跨越预分词边界进行合并。

Step 2: Apply the merges.

接着，我们按照 BPE 训练时合并规则的生成顺序，将得到的词表元素合并序列依次应用到预分词上

---

Example (bpe_encoding): BPE encoding example

例如，假设输入字符串为 "the cat ate"，词表为：

{0: b' ', 1: b'a', 2: b'c', 3: b'e', 4: b'h', 5: b't', 6: b'th', 7: b' c', 8: b' a', 9: b'the', 10: b'at'}

并且已经学习到的合并规则为：

[(b't', b'h'), (b' ', b'c'), (b' ', b'a'), (b'th', b'e'), (b'a', b't')]

首先，预分词器会将输入字符串拆分为：

['the', 'cat', 'ate']

接下来，我们依次处理每一个预分词，并按照 BPE 合并规则对其进行编码

第一个预分词 "the"，最初被表示为字节序列 [b't', b'h', b'e']，查看合并规则列表后，我们发现第一个可应用的合并规则是 (b't', b'h')，于是将其合并，得到 [b'th', b'e']，然后再次从合并规则列表中查找下一个可应用的规则，发现 (b'th', b'e') 可用，于是进一步合并为 [b'the']。此时，再回顾合并规则列表，已经没有可以继续应用的规则（因为整个预分词已经合并成了单个 token），因此合并过程结束，对应的 token ID 为 [9]

第二个预分词 "cat"，在应用 BPE 合并规则后，被表示为 [b'c', b'a', b't']，对应的 token ID 序列为 [7, 1, 5]。第三个预分词 "ate"，在应用 BPE 合并规则后，被表示为 [b'at', b'e']，对应的 token ID 序列为 [10, 3]

因此，输入字符串 "the cat ate" 的最终编码结果为 [9, 7, 1, 5, 10, 3]

---

Special tokens（特殊 token）

在对文本进行编码时，你的分词器应当能够正确处理 用户自定义的特殊 token，这些 token 会在构建分词器时提供

Memory considerations（内存方面的考虑）

假设我们需要对一个无法整体载入内存的大型文件进行分词，为了高效地处理这样的文件（或其他数据流），我们需要将其拆分为 可管理的小块（chunks），并逐块进行处理，从而使内存使用量保持与文本规模线性相关的常数级

在这种分块处理的过程中，必须确保 token 不会跨越块边界，否则，得到的分词结果将与一次性将完整序列载入内存并进行分词的朴素做法不同，从而产生不一致的结果

2.6.2 Decoding text

要将一串整数形式的 token ID 解码回原始文本，我们可以直接在词表中查找每个 ID 对应的条目（即一段字节序列），将这些字节序列拼接起来，然后再将拼接后的字节解码为一个 Unicode 字符串

需要注意的是，输入的 token ID 并不一定保证能够解码为合法的 Unicode 字符串 ，因为用户可能传入任意的整数 ID 序列。如果在解码过程中得到的字节序列无法形成合法的 Unicode 字符串，你应当将这些格式错误的字节替换为官方的 Unicode 替换字符 U+FFFD

bytes.decode 的 errors 参数控制 Unicode 解码错误的方式，使用 errors="replaces" 时，所有格式错误的数据都会被自动替换为该替换字符

Note ：关于 Unicode 替换字符的更多信息，可参考 en.wikipedia.org/wiki/Specials_(Unicode_block)#Replacement_character

• token ID → bytes → concatenate → UTF-8 decode
• 若非法 UTF-8：
  • 使用 Unicode replacement character U+FFFD
  • bytes.decode(errors="replace")

Problem (tokenizer): Implementing the tokenizer (15 points)

Deliverable ：请实现一个 Tokenizer 类 ，该类在给定词表和合并规则列表的情况下，能够将文本编码为整数形式的 token ID 序列，并将 token ID 序列解码回文本。你的分词器还应当支持 用户自定义的特殊 token，如果这些 token 尚未存在于词表中，则需要将其追加到词表中

我们推荐使用如下接口设计：

def __init__(self, vocab, merges, special_tokens=None)

使用给定的词表、合并规则列表以及（可选的）特殊 token 列表来构造一个分词器

参数说明：

• vocab: dict[int, bytes]：词表，从 token ID（整数）映射到对应的字节序列
• merges: list[tuple[bytes, bytes]]：BPE 合并规则列表
• special_tokens: list[str] | None = None：可选的特殊 token 列表

def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None)

这是一个类方法，用于从 序列化保存的词表和合并规则文件 中构建并返回一个 Tokenizer 实例，文件格式应与你的 BPE 训练代码输出格式一致，并可选地接收一组特殊 token

额外参数说明：

• vocab_filepath: str：词表文件路径
• merges_filepath: str：合并规则路径
• special_tokens: list[str] | None = None：可选的特殊 token 列表

def encode(self, text: str) -> list[int]

将输入文本编码为 token ID 序列

def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]

给定一个字符串的可迭代对象（例如 Python 的文件句柄），返回一个 惰性生成器 ，逐个生成 token ID，该接口用于 内存高效地对无法整体载入内存的大型文件进行分词

def decode(self, ids: list[int]) -> str

将一组 token ID 解码为文本字符串

为了使用我们提供的测试用例来验证你的 Tokenizer 实现，你需要先实现测试适配器 [adapters.get_tokenizer]，然后运行：

uv run pytest tests/test_tokenizer.py

你的实现应当能够通过所有测试

2.7 Experiments

Problem (tokenizer_experiments): Experiments with tokenizers (4 points)

(a) 分别从 TinyStories 和 OpenWebText 中各随机抽取 10 篇文档，使用你之前训练好的 TinyStories 分词器和 OpenWebText 分词器（词表大小分别为 10K 和 32K），将这些抽样文档编码为整数形式的 token ID

请计算：每个分词器的压缩率是多少（以 bytes/token 表示）？

Deliverable：用 1-2 句话作答。

(b) 如果你使用 TinyStories 分词器 对 OpenWebText 的样本进行分词，会发生什么情况？请比较其压缩率，并对观察到的现象进行定性描述

Deliverable：用 1-2 句话作答。

(c) 估计你的分词器的吞吐率（例如，以 bytes/second 表示），对 82GB 的文本数据进行分词大概需要多长时间？

Deliverable：用 1-2 句话作答。

(d) 使用你训练好的 TinyStories 和 OpenWebText 分词器，将各种的训练集与开发集编码为整数形式的 token ID 序列，我们将在后续使用这些数据来训练语言模型。我们建议将 token ID 序列序列化保存为 uint16 类型的 NumPy 数组，为什么 uint16 是一个合适的选择？

Deliverable：用 1-2 句话作答。

结语

这篇文章我们完整的梳理了 CS336 Assignment 1 中 BPE Tokenizer 的全部作业要求与设计背景，包括 Unicode 和 UTF-8 编码、字节级分词的动机、BPE 合并规则的训练流程以及在工程实现中需要重点关注的性能与内存问题

更详细的内容大家可以查看官方提供的相关文档

下篇文章我们就来一起看看 BPE Tokenizer 具体该如何实现，敬请期待🤗

参考

• https://github.com/stanford-cs336/assignment1-basics