0. 前言
刚刚完成了 CS336 课程中第一个作业的 BPE 训练部分代码，收获很多，记录下来，以便后续回顾。另一方面也是要锻炼我的写作表达能力，同时整个总结过程又是让我深入理解和学习的过程。
1. 分词器原理
分词器（tokenizer）主要的任务就是将语言文本转换成模型训练或者推理可以直接使用的格式（tokens）。简单来说就是要将语言文本（英文、中文、特殊符号等）转换成数字。
每个字符都要有对应的数字，很自然会想到 Unicode 编码，其规范定义了每个字符对应的码点（code point）数字。但 Unicode 目前收录了 150k+ 个字符，如果直接用它作为我们的词汇表（Vocab），会导致训练时的词嵌入矩阵（embedding table）和输出层投影层变得超级大；而且其中很大一部分字符的使用频率很小。所以直接使用 Unicode 的码点进行转换不行。
退一步，Unicode 编码（Encoding）可以将字符转换成一个字节（bytes）序列，比如用 UTF-8 编码后得到的字节序列中每个值都在 0-255 之间。这样 Vocab 大小只有 256，但代价是编码后的序列长度远大于之前的字符串长度：训练语料变得更庞大、训练要更长时间，同时推理时相同 token 吐出速率下表达能力变弱了。基于表达能力的角度，其实还可以想到直接将词汇（word）映射成 token，但问题在于词汇是无限的，并且将来还会有新的词汇产生。
所以总结下来，我们对分词器的要求是：
- Vocab 不能过大，但也不能太小；
- Vocab 要能覆盖所有的字符/词汇（组合），尽量将使用频率更高的字符/单词放进 Vocab。
综合 UTF-8（byte-level）和 word-level 两者优点，使用 sub-word 的编码方式。Bytes-pair encoding（BPE）是 sub-word 编码方式的一种，在 UTF-8 基础上将出现频率最高的相邻字节（token）合并成一个新的 token 加入 Vocab。我们实现的 BPE 训练过程整体包括"预分词 + 迭代合并"两个阶段：
- 预分词（pre-tokenization）：用 GPT-2 的正则把文本切成"词"，统计每个词的出现次数。
- 迭代合并（Merge）：每一轮统计所有相邻 token 对的出现频次，选出最高频的一对合并成一个新 token，把这次合并追加到
merges，直到词表达到目标大小。
2. Naive版本
基于BPE原理，最直观的实现方式如下：
- 预分词得到
词 → 频次； - 每一轮，对每个词从原始单字节开始重放整个
merges列表，即：对每个词基于当前的merges列表进行合并，得到它当前的 token 序列； - 基于每个词的当前token序列重新统计相邻 pair，选出最高的那对；
- 追加 merge、更新词表。
这个版本的复杂度随"合并轮数 × 词数 × 词长"增长。基线代码如下（find_chunk_boundaries 是作业直接提供的，作用是把文件按特殊 token 切成可以独立计数的块，后面两节不再重复列出）：
import os
from typing import BinaryIO
import re
from collections import defaultdict
PAT = r"""'(?:[sdmt]|ll|ve|re)| ?[^\W\d_]+| ?\d+| ?(?:[^\s\w]|_)+|\s+(?!\S)|\s+"""
def find_chunk_boundaries(
file: BinaryIO,
desired_num_chunks: int,
split_special_token: bytes,
) -> list[int]:
"""
Chunk the file into parts that can be counted independently.
May return fewer chunks if the boundaries end up overlapping.
"""
assert isinstance(split_special_token, bytes), "Must represent special token as a bytestring"
# Get total file size in bytes
file.seek(0, os.SEEK_END)
file_size = file.tell()
file.seek(0)
chunk_size = file_size // desired_num_chunks
# Initial guesses for chunk boundary locations, uniformly spaced
# Chunks start on previous index, don't include last index
chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]
chunk_boundaries[-1] = file_size
mini_chunk_size = 4096 # Read ahead by 4k bytes at a time
for bi in range(1, len(chunk_boundaries) - 1):
initial_position = chunk_boundaries[bi]
file.seek(initial_position) # Start at boundary guess
while True:
mini_chunk = file.read(mini_chunk_size) # Read a mini chunk
# If EOF, this boundary should be at the end of the file
if mini_chunk == b"":
chunk_boundaries[bi] = file_size
break
# Find the special token in the mini chunk
found_at = mini_chunk.find(split_special_token)
if found_at != -1:
chunk_boundaries[bi] = initial_position + found_at
break
initial_position += mini_chunk_size
# Make sure all boundaries are unique, but might be fewer than desired_num_chunks
return sorted(set(chunk_boundaries))
def to_ids(pre_token: str) -> list[int]:
return list(pre_token.encode("utf-8"))
def merge(ids: list[int], merges:list[tuple[bytes, bytes]], vocab_bytes2int: dict[bytes, int]) -> list[int]:
cur_ids = ids
for pair in merges:
new_ids = []
int_pair = (vocab_bytes2int[pair[0]], vocab_bytes2int[pair[1]])
i = 0
while (i < len(cur_ids)):
if i < len(cur_ids)-1 and (cur_ids[i], cur_ids[i+1]) == int_pair:
new_ids.append(vocab_bytes2int[pair[0] + pair[1]])
i += 2
else:
new_ids.append(cur_ids[i])
i += 1
cur_ids = new_ids
return cur_ids
def train_bpe(
input_path: str,
vocab_size: int,
special_tokens: list[str] = None
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
"""
训练BPE分词器
输入：
input_path: 训练数据的文件路径
vocab_size: 最终训练的词汇表大小
sepcial_tokens: 包含特殊token的集合
输出：
vocab: 最终得到的词汇表
merges: merge list
"""
special_tokens = special_tokens or []
# initial vocab
vocab = {i:bytes([i]) for i in range(256)}
vocab_byte2int = {bytes([i]):i for i in range(256)}
cur_vocab_size = 256
for s in special_tokens:
vocab[cur_vocab_size] = s.encode('utf-8')
vocab_byte2int[s.encode('utf-8')] = cur_vocab_size
cur_vocab_size += 1
merges = []
# pre-tokenization and training
with open(input_path, "rb") as f:
num_processes = 4
boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
# The following is a serial implementation, but you can parallelize this
# by sending each start/end pair to a set of processes.
pre_tokens_count = defaultdict(int)
for start, end in zip(boundaries[:-1], boundaries[1:]):
f.seek(start)
chunk = f.read(end - start).decode("utf-8", errors="ignore")
# Run pre-tokenization on your chunk and store the counts for each pre-token
pattern = "|".join(re.escape(t) for t in special_tokens)
split_chunks = re.split(pattern, chunk)
for s in split_chunks:
split_words = re.findall(PAT, s)
for word in split_words:
pre_tokens_count[word] += 1
# loop
while (cur_vocab_size < vocab_size):
pairs = defaultdict(int)
for pre_token, freq in pre_tokens_count.items():
ids = to_ids(pre_token)
ids = merge(ids, merges, vocab_byte2int)
for i in range(len(ids)-1):
pairs[(ids[i],ids[i+1])] += freq
best = max(pairs.items(), key=lambda kv: (kv[1], (vocab[kv[0][0]],vocab[kv[0][1]])))[0]
merges.append((vocab[best[0]], vocab[best[1]]))
vocab[cur_vocab_size] = vocab[best[0]] + vocab[best[1]]
vocab_byte2int[vocab[best[0]] + vocab[best[1]]] = cur_vocab_size
cur_vocab_size += 1
# print("{}th merge: {} and {} into {}".format(cur_vocab_size-257,vocab[best[0]], vocab[best[1]], vocab[best[0]] + vocab[best[1]]))
return vocab, merges
需要注意的是在比较相同出现频率的pair时，是按照其字节对大小比较的，不是按照token数值大小比较，即要用vocab把token转换成bytes。
best = max(pairs.items(), key=lambda kv: (kv[1], (vocab[kv[0][0]],vocab[kv[0][1]])))[0]
另一个需要注意的点在于 re.escape 的使用，pattern 的构造如果不使用 pattern = "|".join(re.escape(t) for t in special_tokens)，会导致比如 <|endoftext|> 特殊字符被当成三个匹配规则，其原理是将特殊字符中的 | 转成对应的转义字符了。
3. 优化后的版本
下面两节只贴相对上一版改动的部分，没有列出的函数与第 2 节完全一致。
3.1 并行预分词
预分词是一次性地扫过整个语料，天然可并行。用 multiprocessing 创建多个进程，可以充分利用有多个 CPU 核心的硬件。可以基于硬件的可用核心数量将文本切分成对应个数的文本块（chunk）。在上面 naive 版本的代码就有这个函数，可以直接用，但是需要注意的是 find_chunk_boundaries 只在特殊 token 的位置切文件。这样每个 pre-token 一定完整落在某个 chunk 内，各 chunk 的局部词频相加就等于全量结果。此外，如果特殊 token 数量较少时，切分出来的 chunk 数可能会少于你设定的 chunk 数量。基于 chunk 的方式进行并行的方案是确定的，但是数据传输的方式值得考量（子进程如何获取对应需要处理的 chunk）。
两种数据传输方案：
- 方案 A：父进程只传"路径 + 起止偏移"，子进程自己
open/seek/read； - 方案 B：父进程读出文本再传给子进程。
方案 B 写起来简单，也是我实现的版本，父进程要串行地做全部 pickle，内存峰值很高；方案 A 把 IO 和解码都并行化，参数永远很小。大语料应该用 A。简单的解释一下：进程之间数据传输都是要先将对象 pickle，然后传输给对应进程再 load 回对应对象。所以如果数据量很大，这里 pickle、数据传输以及 load 回去耗时可能很大。
实现过程中踩过的坑：
- 以为子进程能改父进程的 dict。把
defaultdict作为参数传给Process，指望子进程往里写。结果：进程之间不共享内存，子进程改的是自己的副本，父进程拿到的全是空的，而且不报错。需要用到Queue来完成父进程跟子进程之间的通信。 - 一个 chunk 一个进程。
find_chunk_boundaries会set()去重，实际块数不确定，预分配的列表很容易越界。这个就是前面说的自己定义的num_processes跟实际的进程数量可能不一致，实际进程更少。 q.get()写在start()循环里。父进程起完一个进程就阻塞等它的结果，并行度退化成 1，比串行还慢（多付了进程启动开销）。记住 q.get 会阻塞进程。
修正后的结构：建一个 Queue → 循环里只做读数据和启动进程 → 循环结束后连续收取 N 个结果 → 最后 join。并给 get 加超时，在 worker 里捕获异常。
这部分新增 pre_tokenization 作为子进程入口；train_bpe 里只把原来串行的预分词替换成如下几行，后面的 merge 主循环与第 2 节相同：
import multiprocessing as mp
def pre_tokenization(q, chunk, pattern):
pre_tokens_count = defaultdict(int)
split_chunks = re.split(pattern, chunk)
for s in split_chunks:
split_words = re.findall(PAT, s)
for word in split_words:
pre_tokens_count[word] += 1
q.put(pre_tokens_count)
# ... 前面初始化 vocab / 后面的merges 的部分不变 ...
with open(input_path, "rb") as f:
num_processes = 8
boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
pre_tokens_count = defaultdict(int)
pattern = "|".join(re.escape(t) for t in special_tokens)
processes = []
sub_pre_tokens_count = []
q = mp.Queue()
for start, end in zip(boundaries[:-1], boundaries[1:]):
f.seek(start)
chunk = f.read(end - start).decode("utf-8", errors="ignore")
# Run pre-tokenization on your chunk and store the counts for each pre-token
p = mp.Process(target=pre_tokenization, args=(q, chunk, pattern))
processes.append(p)
p.start()
for _ in range(len(processes)):
sub_pre_tokens_count.append(q.get())
for p in processes:
p.join()
for i in range(len(processes)):
for pre_token, freq in sub_pre_tokens_count[i].items():
pre_tokens_count[pre_token] += freq
3.2 pair 计数的增量更新
第二轮优化针对 merge 主循环。naive版本每次循环要基于每个pre-token(word)重新统计所有字节pair的频率，这里需要获取每个word在当前循环时的字节序列（每次对原始字节序列重新完成所有merge）。很容易看出每次重头merge原始字节序列是一个需要优化的点。可以对每个word维护一个当前的字节序列列表，每次循环后完成merge，新的字节序列供下次统计词频使用。另一个优化点在于每次合并我们统计了所有词中pair的频率，但是一次合并只影响含这个 pair 的那些词，其余词不会影响合并当前pair的计数的变化。所以我们需要知道合并这个pair影响到哪些词，需要对每个pair维护一个所有包含它的词的集合。
总结一下我们需要维护三份状态：
- 词 → 当前 token 序列：避免每轮从原始字节重放 merge；
- pair → 包含它的词的集合（反向索引）：避免每轮扫描全部词。
- pair → 加权计数：选 best 的依据；
有很多需要注意点：
- 一次合并会影响到两侧的pair， 比如 x a a y, 如果合并了a a 为 h 那么原始pair计数中 x a 和 a y都要减少当前这个词的频率次, 同时你还需要加上 x h 和 h y pair的频率 ：
- 重叠的情况，比如 x a a a y， 合并到i=1时，后面接着的a a 要守卫一下，不能把当前合并的pair频率减少。
- 更复杂的情况，比如 x a a a a y, 第一次 合并a a时，你得到了 x h a a y, 这里你会先增加一个h a pair，但是你后面会继续合并第二对 a a, 实际变成了 x h h y, 所以没有 h a pair, 第二次合并的时候你需要减去，同时加上 h h pair。
这一版把 pairs、words2token、pair2words 一次建好，merge 主循环换成下面这样；merge() 这个重放函数和 vocab_byte2int 都不再需要了：
pairs = defaultdict(int)
words2token = {}
pair2words = defaultdict(set)
for pre_token, freq in pre_tokens_count.items():
ids = to_ids(pre_token)
words2token[pre_token] = ids
for i in range(len(ids)-1):
pair = (ids[i],ids[i+1])
pairs[pair] += freq
pair2words[pair].add(pre_token)
# loop
while (cur_vocab_size < vocab_size):
best = max(pairs.items(), key=lambda kv: (kv[1], (vocab[kv[0][0]],vocab[kv[0][1]])))[0]
merges.append((vocab[best[0]], vocab[best[1]]))
vocab[cur_vocab_size] = vocab[best[0]] + vocab[best[1]]
cur_vocab_size += 1
if cur_vocab_size == vocab_size:
break
# merge
words_of_pair = pair2words[best]
pairs[best] = 0
for word in words_of_pair:
ids = words2token[word]
new_ids = []
i = 0
while (i < len(ids)):
if i < len(ids) - 1 and (ids[i], ids[i+1]) == best:
if i < len(ids) - 2:
if (ids[i+1], ids[i+2]) != best:
pairs[(ids[i+1], ids[i+2])] -= pre_tokens_count[word]
pairs[(cur_vocab_size-1, ids[i+2])] += pre_tokens_count[word]
pair2words[(cur_vocab_size-1, ids[i+2])].add(word)
if i > 0:
if new_ids[-1] == ids[i-1]:
pairs[(ids[i-1], cur_vocab_size-1)] += pre_tokens_count[word]
pairs[(ids[i-1], ids[i])] -= pre_tokens_count[word]
pair2words[(ids[i-1], cur_vocab_size-1)].add(word)
else:
pairs[(cur_vocab_size-1, ids[i])] -= pre_tokens_count[word]
pairs[(cur_vocab_size-1,cur_vocab_size-1)] += pre_tokens_count[word]
pair2words[(cur_vocab_size-1, cur_vocab_size-1)].add(word)
new_ids.append(cur_vocab_size-1)
i += 2
else:
new_ids.append(ids[i])
i += 1
words2token[word] = new_ids
pair2words[best] = set()
4. Vocab和Merge数据格式保存
保存训练结果用 GPT-2 的标准格式：
vocab.json：token 字符串 → id的 JSON 对象；merges.txt：一行一条合并，两个 token 用单个空格分隔，行序即优先级。
token 字符串不是原始字节，而是经过 gpt2_bytes_to_unicode() 映射后的字符：把 188 个可见字节原样保留，其余 68 个（控制字符、空格等）偏移到 chr(256+n)。所以空格被写成 Ġ，换行写成 Ċ。这样一个纯字节的 tokenizer 就能安全地存进"空格分隔、换行分条"的文本文件。
@lru_cache
def gpt2_bytes_to_unicode() -> dict[int, str]:
"""
Returns a mapping between every possible byte (an integer from 0 to 255) to a
printable unicode string character representation. This function is taken
from the GPT-2 code.
For example, `chr(0)` is `\x00`, which is an unprintable character:
>>> chr(0)
'\x00'
>>> print(chr(0))
As a result, this function returns a dictionary `d` where `d[0]` returns `Ā`.
The bytes that are visually printable keep their original string representation [1].
For example, `chr(33)` returns `!`, and so accordingly `d[33]` returns `!`.
Note in particular that the space character `chr(32)` becomes `d[32]`, which
returns 'Ġ'.
For unprintable characters, the function shifts takes the integer representing
the Unicode code point of that character (returned by the Python `ord`) function
and shifts it by 256. For example, `ord(" ")` returns `32`, so the the space character
' ' is shifted to `256 + 32`. Since `chr(256 + 32)` returns `Ġ`, we use that as the
string representation of the space.
This function can simplify the BPE implementation and makes it slightly easier to
manually inspect the generated merges after they're serialized to a file.
"""
# These 188 integers can used as-is, since they are not whitespace or control characters.
# See https://www.ssec.wisc.edu/~tomw/java/unicode.html.
bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
cs = bs[:]
# now get the representations of the other 68 integers that do need shifting
# each will get mapped chr(256 + n), where n will grow from 0...67 in the loop
# Get printable representations of the remaining integers 68 integers.
n = 0
for b in range(2**8):
if b not in bs:
# If this integer isn't in our list of visually-representable
# charcters, then map it to the next nice character (offset by 256)
bs.append(b)
cs.append(2**8 + n)
n += 1
characters = [chr(n) for n in cs]
d = dict(zip(bs, characters))
return d
if __name__ == "__main__":
train_data_path = "/home/diego/cs336/data/TinyStoriesV2-GPT4-train.txt"
special_tokens = ["<|endoftext|>"]
t1 = time.perf_counter()
vocab, merges = train_bpe(train_data_path, 10000, special_tokens)
t2 = time.perf_counter()
print("共耗时{}秒".format(t2-t1))
# save
vocab_save_path = "data/my_vocab.json"
merges_save_path = "data/my_merges.txt"
vocab2save = {}
b2s = gpt2_bytes_to_unicode()
for n, b in vocab.items():
s = "".join([b2s[i] for i in b])
vocab2save[s] = n
with open(vocab_save_path, mode='w', encoding="utf-8") as f:
json.dump(vocab2save, f, ensure_ascii=False, indent=4)
with open(merges_save_path, mode='w',encoding="utf-8") as f:
for m1,m2 in merges:
s1 = "".join([b2s[i] for i in m1])
s2 = "".join([b2s[i] for i in m2])
s2save = s1 + " " + s2 + "\n"
f.write(s2save)
最终优化后的版本在 TinyStories 数据集上，词表大小为 10000 时，训练共计耗时 107 秒（num_processes = 8）。