CS336作业1：朴素版优化BPE分词器初版训练与后续优化方向分析

BPE分词器训练优化：从朴素版到增量更新

在阅读 cs336_assignment1_basics.pdf 中关于BPE的原理后，若仍有困惑，可先参考 [[assigment1_overview&bpe_basics]] 中的翻译与细节解释。该文中曾指出朴素版BPE的时间复杂度极高；在 [[train_bpe_naive#测试]] 部分也验证了其训练速度远低于题目要求。因此，首要任务是在朴素版本基础上进行优化。

朴素版详细分析

朴素版的核心瓶颈在于每一轮合并操作的三步：

1. count_pairs：遍历所有word编码中的所有相邻字节对，时间复杂度 O(N)，N 为token总数。
2. max()：从统计结果中找出频率最高的pair。
3. merge_encoding：再次遍历所有word的编码，执行替换。

问题出在第一步：每轮合并后，绝大多数word的编码并未改变，但仍需全部重新扫描。例如文档中的 bpe_example：

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

第一次合并 (s, t)，仅影响3个word；第二轮仍需扫描全部word。这种“无差别攻击”导致效率低下。

优化：告别全量扫描，拥抱增量更新

优化思路清晰：维护全局 pair_counts 字典，每次合并只做增量更新，而非全部重新统计。

核心思想：合并 (A, B) → AB 时，pair_counts 的变化仅与两边的邻居有关。序列 ... X A B Y ... 合并后变为 ... X AB Y ...，则：

• 旧相邻对 (X, A) 和 (B, Y) 消失，需减去其计数。
• 新相邻对 (X, AB) 和 (AB, Y) 出现，需加上其计数。
• 被合并的 (A, B) 自然清零。

优化的精髓在于：

每轮 merge: 直接从 pair_counts 字典里找到最大值，复杂度 O(unique pairs)。只需扫描【含有 (p0,p1) 的 word】执行替换，复杂度 O(受影响的 word)。然后针对这些 word：旧邻对计数 -1，新邻对计数 +1，完成增量更新。

注意：合并 (p0, p1) → new_id 后，pair_counts 的变化由新旧 encoding 的差异决定。只需对旧 encoding 的相邻对逐个减计数，对新 encoding 的相邻对逐个加计数，无需分类讨论左右邻居，逻辑清晰且不易出错。

优化后的实现

# 4. BPE training loop
num_merges = vocab_size - size
pair_counts = defaultdict(int, self.count_pairs(word_counts, word_encodings))

for merge_idx in range(num_merges):
    if not pair_counts:
        print("No more pairs to be merged, quit.")
        break

    # a. Find the max frequency pair to be merged
    merge_pair = max(pair_counts, key=lambda x: (pair_counts[x], self.vocab[x[0]], self.vocab[x[1]]))
    
    # b. Merge and update the word encodings
    token_id = size
    for word in word_encodings:
        old_encoding = word_encodings[word]
        
        # skip if the encoding not changed
        if not any((old_encoding[i], old_encoding[i+1]) == merge_pair for i in range(len(old_encoding) - 1)):
            continue
        
        # update the encoding
        new_encoding = self.merge_encoding(old_encoding, merge_pair, token_id)
        
        # update count for pair
        cnt = word_counts[word]
        for i in range(len(old_encoding) - 1):
            pair_counts[(old_encoding[i], old_encoding[i+1])] -= cnt
        for i in range(len(new_encoding) - 1):
            pair_counts[(new_encoding[i], new_encoding[i+1])] += cnt
        
        # update word encoding
        word_encodings[word] = new_encoding
    
    # clear the pairs whose count is no more than 0
    del_keys = [k for k, v in pair_counts.items() if v <= 0]
    for k in del_keys:
        del pair_counts[k]

实现中遇到的几个坑

1. any() 的用法失误

最初写成：

for i in range(len(old_encoding) - 1):
    if not any((old_encoding[i], old_encoding[i+1]) == merge_pair):
        continue

运行时报 TypeError: 'bool' object is not iterable。any() 接收的参数需为可迭代对象，直接传入布尔值导致错误。修正为使用生成器表达式。

2. 清理失效pair的时机

在扣除旧pair计数时立刻删除归零的pair会导致 KeyError。因为同一word的编码中同一pair可能出现多次，第一次扣除归零后删除key，第二次扣除时该key已不存在。正确做法是等整个word的所有pair更新完毕（扣除和加回都完成）后，再统一判断哪些pair需要清理。

3. pair_counts 未使用 defaultdict

self.count_pairs 返回普通 dict，访问不存在的key会抛 KeyError。改为 defaultdict(int) 可自动初始化为0，避免错误。

pair_counts = defaultdict(int, self.count_pairs(word_counts, word_encodings))

测试

三个测试点均顺利通过，验证了优化后的效率提升。

在 TinyStories 数据集上训练

题目要求：

• 词汇表大小最大10,000，必须包含特殊token "<|endoftext|>"。训练时长 ≤ 30分钟（不使用GPU），占用内存 ≤ 30GB RAM。提示可使用多线程处理预分词加速至2分钟内。
• 分析tokenizer训练过程中最耗时的部分。

按作业要求分三步：编写训练脚本、运行时性能分析、检查结果（找出最长Token）。

训练脚本

数据

首先检查数据格式：

head -n 5 data/TinyStoriesV2-GPT4-train.txt

加载数据时，为方便读写文件，获取项目根路径：

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

获取运行时内存：

def get_memory_usage_mb():
    """Get current process memory usage in MB"""
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024

保存模型：将 vocabulary 和 merges 保存到磁盘。vocab 存为 JSON（字节转可打印字符串），merges 存为文本文件。参考测试文件 tests/test_train_bpe.py 中的 gpt2_bytes_to_unicode 函数。也可直接使用 pickle 保存。

from tests.common import gpt2_bytes_to_unicode

def save_vocab_and_merges(vocab, merges, output_dir='results'):
    Path(output_dir).mkdir(exist_ok=True, parents=True)
    byte_encoder = gpt2_bytes_to_unicode()
    
    vocab_str = {}
    for idx, token_bytes in vocab.items():
        vocab_str[idx] = ''.join(byte_encoder[b] for b in token_bytes)
    with open(f'{output_dir}/vocab.json', 'w', encoding='utf-8') as f:
        json.dump(vocab_str, f, ensure_ascii=False, indent=2)
    
    with open(f'{output_dir}/merges.txt', 'w', encoding='utf-8') as f:
        for p1, p2 in merges:
            t1 = ''.join(byte_encoder[b] for b in p1)
            t2 = ''.join(byte_encoder[b] for b in p2)
            f.write(f'{t1} {t2}\n')

训练主函数：

def run_training(input_path, vocab_size, special_tokens, output_dir='results'):
    print(f'Initial Memory: {get_memory_usage_mb():.2f} MB')
    trainer = train_bpe.BPETrainer()
    start_time = time.time()
    print(f'Starting training on {input_path}...')
    vocab, merges = trainer.train(input_path, vocab_size, special_tokens)
    end_time = time.time()
    duration = end_time - start_time
    peak_memory = get_memory_usage_mb()
    
    print('-' * 100)
    print('Training Complete.')
    print(f'Time Taken: {duration:.2f} seconds ({duration/60:.2f} minutes)')
    print(f'Final Memory: {peak_memory:.2f} MB')
    print('-' * 100)
    
    save_vocab_and_merges(vocab, merges, output_dir)
    
    print("\n=== Statistics (Problem b) ===")
    longest_token_bytes = max(vocab.values(), key=len)
    try:
        longest_token_str = longest_token_bytes.decode('utf-8')
    except:
        longest_token_str = str(longest_token_bytes)
    print(f"Longest Token: {longest_token_str!r}")
    print(f"Length in bytes: {len(longest_token_bytes)}")
    print(f"Total Merges: {len(merges)}")
    print('-' * 100)

执行训练：在 main 中定义参数并调用，终端执行：

uv run python ./train_bpe_tinystories.py --input_path data/TinyStoriesV2-GPT4-train.txt --vocab_size 10000 --profile

训练与分析

训练结果

训练过程中CPU仅使用单核，内存捉襟见肘。最终耗时28分钟，占用内存2302.11 MB，与训练集大小（2.1G）相符。vocabulary 和 merges 已保存，性能分析数据保存至 training.prof。

由于数据集较大（OpenWebText 约9.8G），需分块加载内存，并考虑多块并行预分词（如使用多进程）以优化。

性能分析

使用 snakeviz 分析 training.prof：

1. 安装：uv pip install snakeviz
2. 启动：snakeviz --server training.prof

结果显示总耗时1687秒，主要瓶颈分布如下：

函数	耗时	占比
pretokenize	619秒	37%
train 自身（合并循环）	416秒	25%
any + genexpr（快速跳过）	441秒	26%
max（选最优pair）	161秒	10%

瓶颈1：pretokenize 重复编译正则表达式

_compile 被调用271万次，耗时30秒。每次调用 re.finditer 都会重新编译正则。语料被分为271万个chunk（文档），每个chunk都触发一次编译。优化方向：提前编译正则表达式。

compiled_pattern = re.compile(self.pattern)

同样，切分用的正则也建议提前编译。读文件耗时68秒（read 35秒 + utf_8_decode 32秒），对于2GB文件属正常，无优化空间。

瓶颈2：“快速跳过”逻辑反成最大瓶颈

any 被调用5.84亿次，内部生成器表达式被调用18.3亿次，合计441秒。对应代码：

if not any((old_encoding[i], old_encoding[i+1]) == merge_pair for i in range(len(old_encoding) - 1)):
    continue

该“快速跳过”未提升效率，反而因Python生成器调用开销高而拖慢。更优方案是维护 pair_to_words 索引，直接查找受影响的word。

pair_to_words = defaultdict(set)
for word in word_encodings:
    enc = word_encodings[word]
    for i in range(len(enc) - 1):
        pair_to_words[(enc[i], enc[i+1])].add(word)

for word in pair_to_words[merge_pair]:
    # 直接处理受影响的word

瓶颈3：max() 每轮遍历 pair_counts

max 被调用9746次，每次遍历整个 pair_counts 字典，随着合并进行字典增大，代价逐渐升高。可用堆（heapq）优化，将选最优pair的复杂度从 O(n) 降至 O(log n)，但需处理“失效条目”（lazy deletion）。建议先实现索引优化，再考虑堆优化。

优化方向总结

按预期收益从高到低排列：

1. 建立 pair_to_words 索引：消灭5.84亿次 any 调用，预计节省400秒以上。
2. 预编译正则表达式：消除271万次重复编译，预计节省约30秒。
3. 用堆优化 max：预计节省100秒以上，实现稍复杂。

具体优化实现将在后续文章中详细讲解。