从 Bytes 到 Tokens：我在 CS336 中实现 BPE Tokenizer

Table of Contents

Note: This article is available in Chinese only. 本文暂无英文版本。 View original

最近在做 Stanford CS336（Language Modeling from Scratch）的 Assignment 1，其中有一个完整的 BPE Tokenizer 实现——从训练到编解码。做完之后觉得，BPE 这个话题值得好好写一写。不是因为算法本身多复杂，而是它背后的 设计动机和 trade-off 经常被一笔带过。

这篇文章不是作业报告，也不打算逐行讲代码。我想借这次实现的经历，把 BPE Tokenizer 背后的几个关键问题讲清楚：为什么需要它、它在优化什么、以及一个 naive 的 BPE trainer 为什么会很慢。

为什么 LLM 需要 Tokenizer？
#

在讨论 BPE 之前，先退一步问一个更根本的问题：为什么不能直接把原始文本丢给 Transformer？

最朴素的方案：纯 byte 表示
#

最直觉的做法是这样的：

1text → UTF-8 → bytes (0~255) → Transformer

也就是说，vocabulary 只有 256 个 byte token，不做任何合并。

这个方案的优点非常明显：

vocabulary 极小（就 256 个）
任意字符串都可以表示——中文、英文、emoji、代码、URL、奇怪字符串统统没问题
没有传统意义上的 OOV（out-of-vocabulary）

但代价也很明显：sequence 太长。

一个简单的例子：

1"hello" → h e l l o → 5 tokens

如果 tokenizer 已经学到了 hello 是一个高频 pattern，那同样的文本可能只需要 1 个 token。

这里引出一个核心指标——compression ratio：

$$\text{compression ratio} = \frac{\text{UTF-8 bytes 数}}{\text{token 数}}$$

纯 byte tokenizer 的 compression ratio 基本就是 1（每个 byte 一个 token）。而 BPE tokenizer 通常远大于 1——同样的原始文本，用更少的 token 就能表示。

需要强调的是，这里说的「纯 byte 表示」是指 vocab 只有 256 个 byte token 且不做 merge。我们后面实现的 byte-level BPE 其实正是从这 256 个 byte 出发，通过学习高频 merge 来扩展 vocabulary：

1UTF-8 bytes → 256 byte vocabulary → learn frequent merges → larger BPE vocabulary

另一个极端：word-level tokenizer
#

那把每个自然语言单词当成一个 token 呢？

1"machine learning is interesting" → 4 tokens

compression ratio 很好。但问题是：

1run / runs / running / runner
2PyTorch / PyTorch2 / PyTorch2.7
3foobar123

如果 vocabulary 中不存在某个完整 word，就产生 OOV。传统 NLP 里通常用一个 <UNK> token 来兜底，但这意味着所有不认识的 word 全部坍缩成同一个 token——原始信息直接丢失了。

为了减少 OOV，又必须不断往 vocabulary 里加 word。于是 vocabulary 会越来越大，而那些低频 word 在训练时几乎见不到几次，embedding 很难学好。

Tokenizer 真正的 Trade-off
#

把上面两个极端放在一起看，就能看到 tokenizer 真正在做的事情：

1byte-level                                word-level
2    |                                         |
3    |----------- subword / BPE ---------------|
4    |
5small vocab                                huge vocab
6long sequence                            short sequence
7strong coverage                         OOV problem

更细粒度的 token（比如 byte）：vocab 小、coverage 好、几乎不存在 OOV、rare string 也能组合出来。但 token sequence 很长，Transformer 需要处理更多 position，模型还需要自己从大量低层 byte pattern 中学习更高级的结构。

更粗粒度的 token（比如 word）：sequence 短、compression ratio 高、常见 pattern 可以直接作为一个输入单位。但 vocabulary 巨大、rare word 训练不充分、新词和变体导致 OOV。

所以：

BPE 的本质并不是"把单词切成 subword"这么简单，而是在 vocabulary size 和 sequence length 之间寻找一个工程上的平衡点。

BPE：每一次 Merge 都是一笔交易
#

理解了 trade-off 之后，BPE algorithm 的逻辑就很自然了。

初始状态：vocab 就是 256 个 byte。此时 vocabulary 最小、sequence 最长、representation 最通用。

然后 BPE 去 corpus 中发现：某些 adjacent pair 高频出现，比如 t h、h e、i n、e r……

BPE 选择最值得 merge 的 pair，把它合并成一个新 token：

1(a, b) → ab
2vocab size += 1

以后 corpus 中所有 a b 的位置都可以用 ab 来表示，每处少一个 token。

可以这样理解：

1Cost:   +1 vocabulary slot
2Benefit: corpus 中大量 occurrence 少一个 token

高频 pair 显然比极低频 pair 更值得加入 vocabulary。这就很好地解释了 为什么 BPE merge 的核心统计量是 pair frequency——它直接衡量了一次 merge 能带来多少 sequence 压缩。

也可以进一步说：

每增加一个 BPE merge，本质上都是花掉一个 vocabulary slot，换取 corpus 中一部分 token sequence 的缩短。BPE training 可以理解成不断寻找"最值得购买的 vocabulary entry"。

一个小例子
#

用一个 toy corpus 来看具体的 merge 过程。假设 corpus 经过 pre-tokenization 后得到这些 pre-token（括号里是出现次数）：

1"low"    (5)
2"lower"  (2)
3"newest" (6)
4"widest" (3)

初始状态下，每个 pre-token 被拆成 byte 序列。统计所有 adjacent pair 的加权频率（乘以 word 出现次数）：

1('e', 's') → 6 + 3 = 9
2('s', 't') → 6 + 3 = 9
3('l', 'o') → 5 + 2 = 7
4('o', 'w') → 5 + 2 = 7
5...

这里 ('e', 's') 和 ('s', 't') 频率相同。CS336 的要求是 tie-breaking 时取 bytes 字典序更大的 pair。这里有个容易踩的坑——是按照 bytes 的字典序排序，不是按照 token ID。对于初始的 256 byte vocabulary，两者恰好一致，导致这个 bug 很难排查。

选出 best pair 后执行 merge，比如 ('e', 's') → 'es'，这时 newest 的表示从 n e w e s t 变成 n e w es t——少了一个 token。然后重新统计 pair frequency，进入下一轮。

整个流程：

flowchart LR
    A["Corpus"] --> B["Pre-tokenization"]
    B --> C["Byte sequences"]
    C --> D["Count adjacent pairs"]
    D --> E["Choose best pair"]
    E --> F["Merge"]
    F --> D

Pre-tokenization 在解决什么问题
#

在做 BPE 之前有一个预处理步骤：pre-tokenization。一个自然的问题是——为什么不能直接对整个 corpus 的所有 bytes 做一个巨大的 BPE？

如果不做 pre-tokenization，BPE 可能学到跨越 word 边界的 merge。比如 "the cat" 里的 e 和空格后面的 c 可能被合并成 e c——这种 pattern 在语言层面没有意义，却浪费了 vocabulary slot。

Pre-tokenization 的作用就是先把 corpus 划分成局部单元，让 BPE merge 不跨越某些不应该跨越的边界。

CS336 使用的是 GPT-2 风格的 regex pattern：

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

具体 regex 细节不展开。high-level 来看，它把文本切成缩写（'s、't、'll）、单词（可能带前导空格）、数字、标点、空白等几类。

Pre-tokenization 改变的是 BPE 可以学习 pattern 的边界。 它决定了哪些字符可以被 merge 到一起、哪些永远不会。

除了语言层面的合理性，pre-tokenization 还有一个对性能非常重要的副作用：相同的 pre-token 只需要存一次，配上出现次数即可。比如语料中 " the" 出现了十万次，我们不需要存十万份，只需要存 (" the", 100000) 这一条。Pair frequency 统计时用 word frequency 做加权就行。这对 trainer 的性能影响巨大，后面会详细说。

Training 与 Encoding 是两件事
#

初学者容易混淆的一点：训练 tokenizer 和 用 tokenizer 编码文本 是完全不同的两个操作。

flowchart TB
    subgraph training ["Training"]
        A["Training corpus"] --> B["BPE Trainer"]
        B --> C["Vocabulary"]
        B --> D["Merge rules"]
    end

    subgraph inference ["Encoding"]
        E["New text"] --> F["Tokenizer.encode"]
        C --> F
        D --> F
        F --> G["Token IDs"]
    end

Training 是在一个大 corpus 上学习 vocabulary 和 merge rules。学完以后，这些 merge rules 就固定下来了。

Encoding 则是拿着已学好的 merge rules，对新文本做编码。具体来说：

把 text 转成 UTF-8 bytes
Pre-tokenize
对每个 pre-token 应用 learned merges
输出 token IDs

这里有一个微妙但重要的区别：training 时选 merge 的依据是 pair frequency（频率最高的先 merge）；而 encoding 时应用 merge 的依据是 merge rank（训练时先学到的先应用）。

我的 encode 实现中这个逻辑很清楚——每轮在当前 token 序列中找 rank 最小的 mergeable pair，执行 merge，直到没有可 merge 的 pair 为止：

 1while True:
  best_pair: Pair | None = None
  best_rank = float("inf")
  for pos in range(len(tokens) - 1):
      pair: Pair = (tokens[pos], tokens[pos + 1])
      rank = self.merges_dict.get(pair)
      if rank is not None and rank < best_rank:
          best_pair = pair
          best_rank = rank
  if best_pair is None:
      break
  tokens = self._merge(tokens, best_pair)

merges_dict 就是训练时产出的 merge 列表，key 是 pair、value 是它在 merge 列表中的位置（rank）。

另外还有一个实现中踩到的坑值得一提。我一开始以为 special token（比如 <|endoftext|>）既然已经在 vocab 里了，encoder 应该自然就能识别出来。实际上不行——BPE encoding 是贪心的两两合并，不是最长匹配。就算 <|endoftext|> 整个 bytes 在 vocab 里，pre-tokenization 会先把它拆成 <|、endoftext、|> 这几段，merge 永远到不了完整串。所以 encode 必须在 BPE 之前先把 special token 识别出来，作为原子单位处理。“在 vocab 里” 是必要但不充分条件——encoder 必须被显式告知哪些 byte 序列是原子的。

Compression Ratio：Tokenizer 到底压缩了多少？
#

说了这么多 trade-off，来看看实际数据。CS336 要求在两个数据集上训练 tokenizer 并评估压缩比：

样本	Tokenizer	Vocab Size	Bytes	Tokens	Bytes/Token
TinyStories	TinyStories	10K	7,522	1,800	4.18
OpenWebText	OpenWebText	32K	39,245	8,530	4.60
OpenWebText	TinyStories	10K	39,245	11,985	3.27

几个观察：

Vocab 越大，compression ratio 越高——OpenWebText 的 32K vocab 比 TinyStories 的 10K vocab 压缩得更好（4.60 vs 4.18 bytes/token）。更多 vocabulary slot 意味着更多高频 pattern 被学到。

Domain mismatch 会导致压缩效率显著下降——用 TinyStories tokenizer（在简单儿童故事上训练的）去编码 OpenWebText（正式文本），compression ratio 从 4.60 跌到 3.27。Tokenizer 学到的 merge 是针对训练语料分布的，遇到域外文本，很多 token 都用不上，文本被切得更碎。

从 LLM 系统角度看，compression ratio 直接关系到模型的 effective context length。bytes/token 从 3 提升到 4，意味着同样的 context window 能装下多约 33% 的原始文本。这进一步影响 prefill token 数量、attention workload、KV cache 消耗、training token budget 和 inference cost。

但不能简单下结论说 compression ratio 越高越好——因为 vocabulary 不是免费的。

Vocabulary 为什么不能无限大？
#

既然每增加一个 merge 就能缩短 sequence，那为什么不一直 merge 下去、把 vocab 搞到几百万？

因为 vocabulary size 同时影响 Transformer 另一侧的成本。

Sequence-side cost——token 数越多：sequence 越长、attention 成本增加（self-attention 是 $O(n^2)$）、KV cache 更大、context window 能容纳的原始文本更少。

Vocabulary-side cost——vocab 越大：embedding matrix 是 $V \times d$，LM head（output projection）是 $d \times V$。参数量随 $V$ 线性增长。每一步 forward 都要对整个 vocab 做 softmax——logits computation 和 softmax normalization 的开销都跟 $V$ 有关。而且 vocab 里一定会有很多 rare token，它们在训练中只出现极少次，embedding 很难学充分。

所以 tokenizer 实际上控制了一个很基础的系统 trade-off：

$$\boxed{\text{larger vocabulary} \quad \leftrightarrow \quad \text{shorter sequences}}$$

这不只是 tokenizer 自己的事。Tokenizer 看起来只是 preprocessing，但它实际上改变了 Transformer 后面整个 execution workload——从 embedding lookup 到 attention 到 output projection，都受 vocab size 和 sequence length 的影响。

我的 BPE Trainer 实现
#

到这里读者已经理解了 BPE 在做什么、在优化什么。下面来看我在 CS336 中的具体实现。

Trainer 的整体数据流：

1raw corpus
→ stream chunks (text/parquet)
→ parallel pre-tokenize (ProcessPoolExecutor)
→ aggregate word_freq
→ build word_to_int_sequence
→ build global_pair_freq + pair_to_word_set
→ merge loop:
    choose best pair → update affected words → repeat
→ output: vocab + merge rules

核心数据结构有 5 个：

数据结构	类型	作用
`word_freq`	`dict[str, int]`	pre-token 频率聚合，避免存重复
`word_to_int_sequence`	`dict[str, list[int]]`	每个 word 的当前 token 序列，merge 时原地更新
`global_pair_freq`	`dict[tuple[int,int], int]`	全局加权 pair 频率
`pair_to_word_set`	`dict[tuple[int,int], set[str]]`	倒排索引：每个 pair 出现在哪些 word 中
`global_pair_freq_heap`	`list[...]`	max-heap，用于 O(log n) 找最大 pair

Vocabulary 从 256 byte 开始初始化：

1vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}

然后加入 special tokens。接下来就是 merge 循环——这是整个 trainer 计算量最集中的部分。

关于 完整实现代码、naive BPE trainer 为什么慢、以及我是怎么一步步把训练时间从 8 分钟优化到 36 秒的，我把代码和 profiling 过程放在了文末附录里。这里只说结论：

真正昂贵的不是 merge 操作本身，而是每一轮 merge 后重新发现"哪些统计信息发生了变化"。 一次 merge 只影响 corpus 中包含 best pair 的那些 word——其他所有 word 的 pair 统计完全没变。Naive 实现每轮都重新扫描整个 corpus，绝大多数工作都是在重复计算已知不变的信息。

核心优化思路就是 incremental maintenance——不再每轮全量重算，而是只维护发生变化的局部状态。配合 max-heap 加速最大值查找，TinyStories 10K 的训练时间从 ~8 分钟降到了 ~36 秒。

Tokenizer 在 LLM System 中意味什么
#

flowchart LR
    A["Raw Text"] --> B["Tokenizer"]
    B --> C["Token IDs"]
    C --> D["Embedding"]
    D --> E["Transformer"]
    E --> F["Logits"]

Tokenizer 运行在 Transformer 之前，但它决定了同一段原始文本到底产生多少 token——也就是 Transformer 实际需要处理多少工作。

假设 compression ratio 从 3 bytes/token 变成 4 bytes/token。在相同原始文本上，token 数减少了约 25%。这直接影响 training tokens、prefill FLOPs、KV cache、context capacity 和 serving throughput。

Tokenizer 看起来只是 LLM 前面一个不起眼的 preprocessing step，但它其实决定了模型用什么粒度观察语言，也决定了后面的 Transformer 需要处理多少计算。

几个值得记住的 Insight
#

Tokenizer 本质上是在把原始 byte stream 压缩成更适合 Transformer 处理的离散序列。
Byte-level representation 解决 coverage，word-level representation 解决 compression，而 subword tokenizer 在两者之间寻找平衡。
BPE 的每一次 merge，本质上都是用一个 vocabulary slot 换取 corpus 中 sequence length 的下降。
Tokenizer vocabulary 越大并不意味着一定越好，因为模型同时需要为这个 vocabulary 支付 embedding、LM head 和 token sparsity 的成本。
BPE trainer 的性能优化核心不是让 pair counting 本身更快，而是避免每次 merge 后重新计算那些根本没有变化的东西。

参考资料
#

附录：BPE Trainer 实现与性能优化
#

点击展开 BPE Trainer 完整代码（bpe_trainer.py）

下面是我最终版本的 train_bpe 实现，包含并行 pre-tokenization、pre-token 频率聚合、incremental pair update、倒排索引和 max-heap。读优化过程之前，建议先扫一遍整体结构。

  1import regex as re
  2from collections import defaultdict
  3import logging
  4from concurrent.futures import ProcessPoolExecutor, as_completed
  5from itertools import islice
  6import heapq
  7
  8from cs336_basics.input_pipeline import iter_training_chunks
  9from cs336_basics.pre_tokenizer import pre_tokenize
 10
 11
 12def merge_best_pair_and_get_new_seqence(
  token_id_sequence: list[int], best_token_id_pair: tuple[int, int], new_token_id: int
 14) -> list[int]:
  result_token_id_sequence: list[int] = []
  length = len(token_id_sequence)
  pos = 0
  while pos < length:
      if (
          pos < length - 1
          and token_id_sequence[pos] == best_token_id_pair[0]
          and token_id_sequence[pos + 1] == best_token_id_pair[1]
      ):
          result_token_id_sequence.append(new_token_id)
          pos += 2
      else:
          result_token_id_sequence.append(token_id_sequence[pos])
          pos += 1
  return result_token_id_sequence
 30
 31
 32def get_chunk_word_freq(text: str, compiled_pattern: re.Pattern) -> dict[str, int]:
  chunk_word_freq: dict[str, int] = defaultdict(int)
  text_zones_without_special_tokens = compiled_pattern.split(string=text)
  for text_zone in text_zones_without_special_tokens:
      zone_chunks: list[str] = pre_tokenize(text=text_zone)
      for chunk in zone_chunks:
          chunk_word_freq[chunk] += 1
  return chunk_word_freq
 40
 41
 42def train_bpe(
  input_path: str,
  vocab_size: int,
  special_tokens: list[str],
 46) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
 47
  vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
  start_pos = len(vocab)
  vocab.update(
      {start_pos + index: special_token.encode("utf-8") for index, special_token in enumerate(special_tokens)}
  )
  split_pattern = "|".join(re.escape(special_token) for special_token in special_tokens)
  special_token_re = re.compile(split_pattern)
 55
  # 维护 word -> 当前 token 序列；merge 过程中原地更新
  word_to_int_sequence: dict[str, list[int]] = dict()
  word_freq: dict[str, int] = defaultdict(int)
 59
  num_processes = 8
  _, chunk_iter = iter_training_chunks(
      input_path=input_path,
      special_tokens=special_tokens,
      num_processes=num_processes,
  )
 66
  batch_size = 32
  with ProcessPoolExecutor(max_workers=num_processes) as executor:
      while True:
          batch = list(islice(chunk_iter, batch_size))
          if not batch:
              break
          futures = [executor.submit(get_chunk_word_freq, chunk, special_token_re) for chunk in batch]
          for future in as_completed(futures):
              chunk_word_freq = future.result()
              for word, freq in chunk_word_freq.items():
                  word_freq[word] += freq
 78
  for word, freq in word_freq.items():
      utf8_word = word.encode("utf-8")
      word_to_int_sequence[word] = list(utf8_word)
 82
  merged_bytes_list: list[tuple[bytes, bytes]] = []
  best_pair_word_set: set[str] = set()
  global_pair_freq: dict[tuple[int, int], int] = defaultdict(int)
  pair_to_word_set: dict[tuple[int, int], set[str]] = defaultdict(set)
  global_pair_freq_heap: list[tuple[int, bytes, bytes, int, int]] = []
 88
  while len(vocab) < vocab_size:
      if not best_pair_word_set:
          for word, freq in word_freq.items():
              int_sequence: list[int] = word_to_int_sequence[word]
              length = len(int_sequence)
              for pos in range(length - 1):
                  lhs = int_sequence[pos]
                  rhs = int_sequence[pos + 1]
                  global_pair_freq[(lhs, rhs)] += freq
                  pair_to_word_set[(lhs, rhs)].add(word)
 99
          for pair, freq in global_pair_freq.items():
              left, right = pair
              global_pair_freq_heap.append((freq, vocab[left], vocab[right], left, right))
          heapq.heapify_max(global_pair_freq_heap)
104
      else:
          for word in best_pair_word_set:
              freq = word_freq[word]
              int_sequence: list[int] = word_to_int_sequence[word]
              length = len(int_sequence)
              for pos in range(length - 1):
                  lhs = int_sequence[pos]
                  rhs = int_sequence[pos + 1]
                  global_pair_freq[(lhs, rhs)] += freq
                  heapq.heappush_max(
                      global_pair_freq_heap,
                      (global_pair_freq[(lhs, rhs)], vocab[lhs], vocab[rhs], lhs, rhs),
                  )
                  pair_to_word_set[(lhs, rhs)].add(word)
119
      if len(global_pair_freq) == 0:
          logging.error(f"not enough data to train, vocab size = {len(vocab)},target: {vocab_size}")
          break
123
      best_pair: tuple[int, int] | None = None
      while global_pair_freq_heap:
          heap_freq, _, _, lhs, rhs = heapq.heappop_max(global_pair_freq_heap)
          pair = (lhs, rhs)
          if global_pair_freq.get(pair) != heap_freq:
              continue
          best_pair = pair
          break
132
      if best_pair is None:
          logging.error("empty heap")
          break
136
      vocab[len(vocab)] = vocab[best_pair[0]] + vocab[best_pair[1]]
      merged_bytes_list.append((vocab[best_pair[0]], vocab[best_pair[1]]))
139
      best_pair_word_set = pair_to_word_set[best_pair].copy()
      for word in best_pair_word_set:
          id_sequence = word_to_int_sequence[word]
          length = len(id_sequence)
          for pos in range(length - 1):
              lhs = id_sequence[pos]
              rhs = id_sequence[pos + 1]
              global_pair_freq[(lhs, rhs)] -= word_freq[word]
              if global_pair_freq[(lhs, rhs)] > 0:
                  heapq.heappush_max(
                      global_pair_freq_heap,
                      (global_pair_freq[(lhs, rhs)], vocab[lhs], vocab[rhs], lhs, rhs),
                  )
              pair_to_word_set[(lhs, rhs)].discard(word)
154
          word_to_int_sequence[word] = merge_best_pair_and_get_new_seqence(
              token_id_sequence=id_sequence, best_token_id_pair=best_pair, new_token_id=len(vocab) - 1
          )
      del pair_to_word_set[best_pair]
      del global_pair_freq[best_pair]
160
  return (vocab, merged_bytes_list)

点击展开性能优化过程（8 分钟 → 36 秒）

这部分按照实际的开发时间线来讲。每个优化阶段都经历了同一个循环：profiling 定位瓶颈 → 分析哪些工作是冗余的 → 消除冗余 → 验证效果。训练目标是 TinyStories 数据集，vocab_size = 10000。

起点：Naive 实现
#

最朴素的实现长这样：

1for _ in range(num_merges):
  scan_entire_corpus()
  count_all_pairs()
  find_max_pair()        # O(N)
  replace_pair_everywhere()

假设 corpus 中有 $N$ 个 token，需要 $M$ 次 merge。每轮都要扫描整个 corpus 来统计 pair frequency，直觉上复杂度接近 $O(MN)$。

Python 层面的开销——dict update、tuple hashing、list reconstruction、大量临时对象的分配——还会让常数因子变得非常大。

初始版本在 TinyStories 上跑完需要 约 8 分钟，峰值内存 16.14 GB。

优化 1：Pre-token 频率聚合
#

分析： 自然语言 corpus 中存在大量重复。比如 " the" 可能出现十万次。如果每一份都存一个完整的 byte 序列参与 pair counting，pair frequency 统计需要扫描的数据量跟 corpus 原始大小成正比。

方案： pre-tokenization 阶段就把相同的 pre-token 聚合起来，只存 (word, frequency) 对。统计 pair frequency 时对 word 的频率做加权：

1for word, freq in word_freq.items():
  int_sequence = word_to_int_sequence[word]
  for pos in range(len(int_sequence) - 1):
      lhs = int_sequence[pos]
      rhs = int_sequence[pos + 1]
      global_pair_freq[(lhs, rhs)] += freq

这一行 += freq 是关键——pair 的全局频率 = 每个 word 内的 pair 出现次数 × word 的频率。这样操作的数据量就从 corpus 的 token 总数降到了 unique word 数。对 TinyStories 这样重复率极高的语料，差距可以到几个数量级。

优化 2：Incremental update + 倒排索引（27s → 2.4s）
#

分析： 这是整个优化过程中最关键的一步。当 (A, B) → AB 被 merge 时，整个 corpus 中绝大多数 word 的 pair 统计根本没有变化。真正受影响的只有那些包含 pair (A, B) 的 word。

举个例子，word X A B Y 经历了 merge (A, B) → AB：

1before: X-A  A-B  B-Y    (3 pairs)
2after:  X-AB     AB-Y    (2 pairs)

A-B 消失了，X-A 和 B-Y 也消失了，新增了 X-AB 和 AB-Y。但不包含 (A, B) 的 word 完全没有任何变化。一次 merge 本质上只产生 local updates。

方案： 维护倒排索引 pair_to_word_set：对于每个 pair，记录它出现在哪些 word 中。选中 best pair 后，通过倒排索引直接找到所有受影响的 word，只更新这些 word 的统计。

实现是一个 subtract → merge → (下轮 add) 的模式：

 1best_pair_word_set = pair_to_word_set[best_pair].copy()
 2for word in best_pair_word_set:
  id_sequence = word_to_int_sequence[word]
  # 先减去这个 word 对所有 pair 的旧贡献
  for pos in range(len(id_sequence) - 1):
      lhs = id_sequence[pos]
      rhs = id_sequence[pos + 1]
      global_pair_freq[(lhs, rhs)] -= word_freq[word]
      pair_to_word_set[(lhs, rhs)].discard(word)
  # 执行 merge，得到新的 token 序列
  word_to_int_sequence[word] = merge_best_pair_and_get_new_seqence(
      token_id_sequence=id_sequence,
      best_token_id_pair=best_pair,
      new_token_id=len(vocab) - 1,
  )

下一轮循环开始时，再扫描受影响的 word，把新的 pair 贡献加回去。

效果： 单元测试（在 corpus.en 上跑）从 27 秒降到了 2.4 秒。

把"每轮重新计算全局状态"转化成"只维护发生变化的局部状态"——这是一个非常通用的工程优化模式，不只属于 BPE。

下面是优化前的 profiling flamegraph，几乎所有时间都集中在 merge loop 中的全量扫描上：

优化后（只更新受影响的 word），热点转移到了 max() 操作上：

优化 3：并行 Pre-tokenization（8min → 5min）
#

分析： pre-tokenization 阶段是 embarrassingly parallel 的——每个 text chunk 的 regex 匹配和 word 统计完全独立。

方案： 用 ProcessPoolExecutor 并行处理：

 1num_processes = 8
 2with ProcessPoolExecutor(max_workers=num_processes) as executor:
  while True:
      batch = list(islice(chunk_iter, batch_size))
      if not batch:
          break
      futures = [
          executor.submit(get_chunk_word_freq, chunk, special_token_re)
          for chunk in batch
      ]
      for future in as_completed(futures):
          chunk_word_freq = future.result()
          for word, freq in chunk_word_freq.items():
              word_freq[word] += freq

各 worker 独立统计自己 chunk 的 word frequency，主进程 reduce。

效果： TinyStories 训练时间从 ~8 分钟降到了 ~5 分钟（wall clock）。CPU 利用率从 ~100% 涨到 ~192%。峰值内存从 16.14 GB 降到了 7.41 GB（不再需要一次性 f.read() 整个文件）。

parallelizable preprocessing 和 inherently sequential merge loop 是两类完全不同的 workload。pre-tokenization 可以随便并行，但 merge loop 有严格的 sequential dependency——第 $t$ 次 merge 的结果决定了第 $t+1$ 次 merge 的输入。

优化 4：只遍历 best_pair_word_set（5min → 2.4min）
#

分析（profiling）： 并行化之后，瓶颈转移到了 merge loop 内部。虽然已经做了增量更新，但代码里有一个问题：遍历的是整个 word_to_int_sequence 字典，然后用 if word in best_pair_word_set 过滤。

1# 优化前：遍历全量 word，逐个判断
2for word, id_sequence in word_to_int_sequence.items():
3    if word in best_pair_word_set:
4        # ... 更新统计 ...

TinyStories 有几十万个 unique pre-token。遍历全部 word 再做 membership check，光这个循环本身就很耗时。

方案： 直接遍历 best_pair_word_set（通常只有几千甚至几百个 word）：

1# 优化后：直接遍历受影响的 word
2for word in best_pair_word_set:
3    id_sequence = word_to_int_sequence[word]
4    # ... 更新统计 ...

效果： 训练时间从 ~5 分钟降到了 ~2.4 分钟。做完这一步后的 profiling 显示 max() 操作的占比从 56% 涨到了 96%——之前有大量时间浪费在遍历无关 word 上，消除之后真正的瓶颈暴露出来了。

只遍历 affected words 后的 profiling flamegraph

Profiling 指导优化方向——先消除非瓶颈处的浪费，才能暴露真正的瓶颈。

优化 5：Max-heap with lazy deletion（2.4min → 36s）
#

分析（profiling）： 96% 的时间花在这一行上：

1best_pair, best_freq = max(
2    global_pair_freq.items(),
3    key=lambda p: (p[1], vocab[p[0][0]], vocab[p[0][1]])
4)

这是一个 $O(N)$ 的线性扫描——每轮 merge 都要遍历所有 pair 才能找到最大值。

方案： 用 max-heap 代替线性扫描。但 BPE 场景有一个特殊问题：每次 merge 后一些 pair 的频率会变化，堆中对应的条目就过期了。标准 heap 不支持高效的 decrease-key。

解决方案是 lazy deletion：把新的 (freq, pair) 直接 push 进 heap，pop 的时候检查频率是否跟 global_pair_freq 中的实际值一致。不一致就说明是旧条目，跳过。

1while global_pair_freq_heap:
  heap_freq, _, _, lhs, rhs = heapq.heappop_max(global_pair_freq_heap)
  pair = (lhs, rhs)
  # 旧快照已经失效
  if global_pair_freq.get(pair) != heap_freq:
      continue
  best_pair = pair
  break

这里用的是 Python 3.14 新增的 heapq.heapify_max / heappop_max / heappush_max（之前 Python 标准库只有 min-heap）。

Lazy deletion 的代价是 heap 中会积累过期条目，空间占用增加。但在实践中，pop 到第一个有效条目通常只需要几次尝试。相比 $O(N)$ 的全量扫描，$O(\log N)$ amortized 的 heap 操作带来了巨大提升。

效果： 训练时间从 ~2.4 分钟降到了 ~36 秒。

优化效果汇总
#

阶段	TinyStories 训练时间	峰值内存	主要改动
Naive 实现	~8 min	16.14 GB	每轮全量 recount
+ 并行 pre-tokenization	~5 min	7.41 GB	ProcessPoolExecutor, 8 workers
+ 只遍历 affected words	~2.4 min	8.53 GB	直接遍历 `best_pair_word_set`
+ Max-heap	~36 s	8.53 GB	lazy deletion heap

总计约 13 倍加速。

后续还在 OpenWebText 全量数据（80 个 parquet shard、vocab_size = 32000）上跑了一次，9 小时 45 分钟完成，51 GiB 峰值内存，7.7 亿次 major page faults。能跑完，但内存压力巨大——进一步优化空间仍然存在。

性能优化的通用模式
#

回过头看这一系列优化，其实遵循了一个简单的 pattern：

1naive（全量重算）
→ 识别哪些工作是重复的
→ 避免重复工作
→ incremental maintenance

也需要区分两类优化：

Algorithmic optimization：incremental update、inverted index、heap。改变的是渐进复杂度或工作量级。
Constant-factor optimization：pre-token 频率聚合、直接遍历 set 而非 filter 整个 dict。不改变渐进复杂度，但减少常数。

两者都有价值，但 algorithmic optimization 通常是决定性的。

性能优化最有价值的地方通常不是"把一个循环写快"，而是找到哪些工作根本不应该重复做。

为什么 LLM 需要 Tokenizer？#

最朴素的方案：纯 byte 表示#

另一个极端：word-level tokenizer#

Tokenizer 真正的 Trade-off#

BPE：每一次 Merge 都是一笔交易#

一个小例子#

Pre-tokenization 在解决什么问题#

Training 与 Encoding 是两件事#

Compression Ratio：Tokenizer 到底压缩了多少？#

Vocabulary 为什么不能无限大？#

我的 BPE Trainer 实现#

Tokenizer 在 LLM System 中意味什么#

几个值得记住的 Insight#

参考资料#

附录：BPE Trainer 实现与性能优化#

起点：Naive 实现#

优化 1：Pre-token 频率聚合#

优化 2：Incremental update + 倒排索引（27s → 2.4s）#

优化 3：并行 Pre-tokenization（8min → 5min）#

优化 4：只遍历 best_pair_word_set（5min → 2.4min）#

优化 5：Max-heap with lazy deletion（2.4min → 36s）#

优化效果汇总#

性能优化的通用模式#

Related