李理的博客

《C++ High Performance: Master the art of optimizing the functioning of your C++ code , Second Edition》学习

Tue, 30 Dec 2025 00:00:00 +0000

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 11: Wrapping C++ Code with Cython

Thu, 25 Sep 2025 00:00:00 +0000

This series of articles implements a subtask of Stanford’s CS336 Assignment 1: building an efficient training algorithm for a BPE Tokenizer. Through a series of optimizations, our algorithm’s training time on OpenWebText was reduced from over 10 hours to less than 10 minutes. This series explains these optimizations, including algorithmic improvements, data structure enhancements, parallelization with OpenMP, Cython optimization, and implementing key code in C++ along with its integration via Cython. This is the twelfth and final article in a series on using Cython to wrap the previous C++ code into an extension module for Python.

Table of Content

1. Goal
2. Wrapping the C++ Code into a Dynamic Library
3. Wrapping with Cython into an Extension Module
4. Testing
5. Conclusion

Here’s the English translation, with Markdown formatting and code blocks preserved:

1. Goal

In previous articles, we’ve done a lot of exploration. BPE training can be divided into two stages. The first stage is tokenization and counting, which we can perform in parallel using multiprocessing. Our bpe_v3.py implementation can complete this in 120 seconds using 32 cores, while the bpe_v3_bytes_time.py version, which reads bytes, can finish in 70 seconds using 64 cores. The second stage is merging. The Python version, bpe_v8_v3.py, takes about 500 seconds, with a total time of over ten minutes. The fastest C++ version, bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp, has a merging time of around 100 seconds.

Today, our goal is to create the fastest version possible by wrapping the C++ code into an extension module using Cython, which can then be called by Python for the second stage.

2. Wrapping the C++ Code into a Dynamic Library

We’ve conducted many experiments in cppupdate, and now we need to wrap the best-performing versions into a dynamic library for easy use with Cython.

We will create a new C++ project called cppstep2. The full code can be found at cppstep2.

Based on the cppupdate experiment results, I have selected the following files: bpe_train_updater_fine_grained_emhash8.cpp, bpe_train_updater_fine_grained_emhash8_set.cpp, bpe_train_updater_fine_grained_emhash8_set9.cpp, bpe_train_updater_fine_grained_heap_emhash8_set9.cpp, and bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp.

I chose emhash over absl for two reasons: it’s faster and simpler to use. I previously git cloned the full emhash code, but to simplify, I’ve only copied the three necessary header files: hash_set8.hpp, hash_set4.hpp, and hash_table8.hpp. Additionally, max_heap.h and max_heap.cpp were copied directly from the cppupdate project.

Our dynamic library will consist of two main files: the header file bpe_train_step2.h and the implementation file bpe_train_step2.cpp.

2.1 Header File

#pragma once

#include <vector>
#include <unordered_map>
#include <unordered_set>
#include "emhash/hash_table8.hpp"

struct pair_hash {
    template <class T1, class T2>
    std::size_t operator () (const std::pair<T1,T2> &p) const {
        auto h1 = std::hash<T1>{}(p.first);
        auto h2 = std::hash<T2>{}(p.second);

        return h1 ^ (h2 + 0x9e3779b9 + (h1 << 6) + (h1 >> 2));
    }
};

void bpe_train_step2(int vocab_size, 
                emhash8::HashMap<std::pair<int, int>, int, pair_hash> & pair_counts, 
                std::unordered_map<std::pair<int, int>, std::vector<std::vector<int>>, pair_hash> & pair_strings, 
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash> & pair_wordids, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);

void bpe_train_step2_v2(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);

void bpe_train_step2_v3(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);   
                
void bpe_train_step2_v4(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);  

void bpe_train_step2_v5(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);  

void bpe_train_step2_v6(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges); 

This header file defines five functions: bpe_train_step2_v2, bpe_train_step2_v3, bpe_train_step2_v4, bpe_train_step2_v5, and bpe_train_step2_v6. These correspond to bpe_train_updater_fine_grained_emhash8.cpp, bpe_train_updater_fine_grained_emhash8_set.cpp, bpe_train_updater_fine_grained_emhash8_set9.cpp, bpe_train_updater_fine_grained_heap_emhash8_set9.cpp, and bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp, respectively.

bpe_train_step2 is a direct copy of the cppupdate implementation and requires seven input parameters. As we’ll see later, pair_counts and wordid_counts can be calculated inside C++ for better speed and fewer parameters passed. Thus, bpe_train_step2_v2 and subsequent versions only require five parameters.

2.2 CMakeLists.txt

Our goal is to compile a dynamic library. We’ll write the following in CMake syntax:

cmake_minimum_required(VERSION 3.20)

project(BPE_TRAIN_STEP2 LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)


add_library(bpe_train_step2 SHARED bpe_train_step2.cpp max_heap.cpp)
target_include_directories(bpe_train_step2 PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
target_compile_definitions(bpe_train_step2 PUBLIC MY_LIBRARY_EXPORT)

target_include_directories(bpe_train_step2 PUBLIC
    "${PROJECT_SOURCE_DIR}/emhash"
)

install(TARGETS bpe_train_step2
        EXPORT bpe_train_step2_export_targets
        RUNTIME DESTINATION bin
        LIBRARY DESTINATION lib
        ARCHIVE DESTINATION lib)

install(FILES "bpe_train_step2.h" 
        DESTINATION include
)

We create the bpe_train_step2 target using add_library. target_compile_definitions seems to define a macro for Windows compatibility. I don’t have a Windows environment to test this, so I’m unsure if the project will compile there.

The next target_include_directories includes the emhash headers. Finally, install copies the bpe_train_step2 target (mainly libbpe_train_step2.so) and bpe_train_step2.h to the appropriate locations. I’ll explain the installation process below.

2.3 bpe_train_step2

First, we will integrate it in the same way as cppupdate, which means we need to pass seven parameters, including pair_counts.

void bpe_train_step2(int vocab_size, 
                emhash8::HashMap<std::pair<int, int>, int, pair_hash> & pair_counts, 
                std::unordered_map<std::pair<int, int>, std::vector<std::vector<int>>, pair_hash> & pair_strings, 
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash> & pair_wordids, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges){
    auto start = std::chrono::steady_clock::now();
    while(vocabulary.size() < vocab_size){
        int max_count = -1;
        std::pair<int, int> max_pair;
        std::vector<std::vector<int>> max_strings;
        for(const auto& [pair, count] : pair_counts){
            if(count > max_count){
                max_count = count;
                max_pair = pair;
                max_strings = pair_strings[pair];
            }else if(count == max_count){
                std::vector<std::vector<int>> strings = pair_strings[pair];
                ComparisonResult r1 = three_way_compare(strings[0], max_strings[0]);
                if(r1 == ComparisonResult::Greater){
                    max_count = count;
                    max_pair = pair;
                    max_strings = strings;
                }else if(r1 == ComparisonResult::Equal){
                    ComparisonResult r2 = three_way_compare(strings[1], max_strings[1]);
                    if(r2 == ComparisonResult::Greater){
                        max_count = count;
                        max_pair = pair;
                        max_strings = strings;                        
                    }
                }
            }
        }

        const std::vector<int>& bytes1 = vocabulary[max_pair.first];
        const std::vector<int>& bytes2 = vocabulary[max_pair.second];
        std::vector<int> merge_bytes;
        merge_bytes.reserve(bytes1.size() + bytes2.size());
        merge_bytes.insert(merge_bytes.end(), bytes1.begin(), bytes1.end());
        merge_bytes.insert(merge_bytes.end(), bytes2.begin(), bytes2.end());

        int size = vocabulary.size();
        vocabulary[size] = merge_bytes;

        auto& affected_words = pair_wordids[max_pair];

        
        updated_affected_word_count(max_pair, affected_words, wordid_encodings, wordid_counts,
                                    pair_counts, pair_wordids, size, pair_strings, vocabulary);
  
        merges.push_back({bytes1, bytes2});


    }
    auto end = std::chrono::steady_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
    std::cout << "bpe_train_step2: " << duration.count() << "ms." << std::endl;
}

The code is a complete copy of bpe_train_updater_fine_grained_emhash8.cpp.

2.4 bpe_train_step2_v2

If you’ve carefully reviewed the previous time statistics, you’ll notice that the total time is about 50-60 seconds more than “tokenization time” + “merging time”. This extra time is for calling the BPE_Trainer._count_pairs function:

    @staticmethod    
    def _count_pairs(word_counts, word_encodings, pair_strings, vocabulary, pair_to_words):
        pair_counts = defaultdict(int)
        for word, count in word_counts.items():
            encoding = word_encodings[word]
            for i in range(0, len(encoding) - 1):
                pair = encoding[i], encoding[i + 1]
                pair_counts[pair] += count
                if pair not in pair_strings:
                    pair_strings[pair] = (vocabulary[pair[0]], vocabulary[pair[1]])

                pair_to_words[pair].add(word)

        return pair_counts

This code calculates pair_counts based on word_counts and builds the initial inverted index pair_to_words and pair_strings. Its input parameters are word_counts, word_encodings, and vocabulary. This code can be implemented entirely in C++, which would not only speed it up but also reduce the number of parameters passed between Python and C++.

Therefore, the input parameters for bpe_train_step2_v2 can be reduced by three:

void bpe_train_step2_v2(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges){
    std::unordered_map<std::pair<int, int>, std::vector<std::vector<int>>, pair_hash> pair_strings;
    emhash8::HashMap<std::pair<int, int>, int, pair_hash>  pair_counts;
    std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash>  pair_wordids;
    
    std::pair<int, int> pair;
    for(const auto& [wordid, count] : wordid_counts){
        const auto& word_tokens = wordid_encodings[wordid];
        for(int i = 0; i < word_tokens.size() - 1; ++i){
            pair.first = word_tokens[i];
            pair.second = word_tokens[i + 1];
            pair_counts[pair] += count;
            if (pair_strings.find(pair) == pair_strings.end()) {
                pair_strings[pair] = {vocabulary[pair.first], vocabulary[pair.second]};
            }
            pair_wordids[pair].insert(wordid);
        }
    }
    

    bpe_train_step2(vocab_size, pair_counts, pair_strings, vocabulary,
                    pair_wordids, wordid_counts, wordid_encodings, merges);
}

The bpe_train_step2_v2 function first calculates pair_counts, pair_strings, and pair_wordids from wordid_counts, wordid_encodings, and vocabulary, and then calls bpe_train_step2.

The bpe_train_step2_v3 and subsequent versions have the exact same interface as bpe_train_step2_v2, and the code is copied from the corresponding cppupdate versions, which I won’t detail here.

2.5 Compile and Install

Use the following commands to compile and install:

cd cppstep2
mkdir build && cd build
cmake -D CMAKE_INSTALL_PREFIX=../../lib_bpe_train_step2/  -D CMAKE_BUILD_TYPE=Release ..
cmake --build . -- -j8
# If you are using gcc, you can directly use make -j8
cmake --install .  
# Or make install

Since we don’t want to install this library into system paths like /usr/local (which requires root privileges and is not convenient for Cython integration due to different paths on different systems), I’m installing it to ../../lib_bpe_train_step2/, which is a subdirectory of the project root assignment1-basics-bpe:

$ ls assignment1-basics-bpe

cppstep2/
cppupdate/
cs336_basics/
data/
lib_bpe_train_step2/

After installation, the lib_bpe_train_step2 directory will look like this:

lib_bpe_train_step2$ tree
.
├── include
│   ├── bpe_train_step2.h
│   └── emhash
│       └── hash_table8.hpp
└── lib
    └── libbpe_train_step2.so

Besides bpe_train_step2.h, hash_table8.hpp is also included because bpe_train_step2.h uses it. This is primarily because bpe_train_step2’s declaration needs it. If we only kept bpe_train_step2_v2 and later versions, this header file would not need to be installed (though we’d still need all emhash headers for compilation).

3. Wrapping with Cython into an Extension Module

For those new to wrapping C++ libraries with Cython, you can refer to the Using C++ in Cython guide.

We need two files: bpe_train_step2_wrapper.pxd and bpe_train_step2_wrapper.pyx.

3.1 bpe_train_step2_wrapper.pxd

First, let’s look at the .pxd file:

# distutils: language = c++

# 导入 C++ 标准库类型
from libcpp.utility cimport pair
from libcpp.vector cimport vector
from libcpp.unordered_map cimport unordered_map
from libcpp.unordered_set cimport unordered_set


cdef extern from "../lib_bpe_train_step2/include/bpe_train_step2.h" :
    cppclass pair_hash:
        pass

cdef extern from "../lib_bpe_train_step2/include/emhash/hash_table8.hpp" namespace "emhash8":
    cppclass HashMap[K, V, H]:
        #ValueT& operator[](const KeyT& key) noexcept
        V& operator[](const K& key)


cdef extern from "../lib_bpe_train_step2/include/bpe_train_step2.h":
    void bpe_train_step2(int vocab_size,
                         HashMap[pair[int, int], int, pair_hash] & pair_counts,
                         unordered_map[pair[int, int], vector[vector[int]], pair_hash] & pair_strings,
                         unordered_map[int, vector[int]] & vocabulary,
                         unordered_map[pair[int, int], unordered_set[int], pair_hash] & pair_wordids,
                         const unordered_map[int, long long] & wordid_counts,
                         unordered_map[int, vector[int]] & wordid_encodings,
                         vector[pair[vector[int], vector[int]]] & merges) except +
    
    void bpe_train_step2_v2(int vocab_size,
                         unordered_map[int, vector[int]] & vocabulary,
                         const unordered_map[int, long long] & wordid_counts,
                         unordered_map[int, vector[int]] & wordid_encodings,
                         vector[pair[vector[int], vector[int]]] & merges) except +

The first line tells distutils (or setuptools) that you are building a C++ extension module, not the default C extension.

The next four lines import the C++ standard library types std::vector, std::pair, std::unordered_map, and std::unordered_set which are wrapped by Cython under libcpp. We can use them with cimport, which is a compile-time import. To see which C++ standard libraries are available, refer to this link.

Next, we declare pair_hash with cppclass. Since we only need to know about the pair_hash symbol for the lib_bpe_train_step2 library, we can leave the content empty with pass. Note that the path for cdef extern from is "../lib_bpe_train_step2/include/bpe_train_step2.h", which requires lib_bpe_train_step2 to be installed in the correct location as described earlier.

Next is the declaration of emhash8::HashMap. The "namespace emhash8" after cdef extern from tells Cython that HashMap is in the emhash8 namespace. cppclass HashMap[K, V, H] indicates that HashMap is a template class for Key/Value/Hash function. We’ll need operator[] in the .pyx file, so we also need to declare its prototype here.

Finally, we declare the bpe_train_step2 functions. The syntax is similar to C++, but < > is replaced with [ ], which might look a bit strange.

3.2 bpe_train_step2_wrapper.pyx

Now, we’ll wrap the C++ functions into Python-callable functions. The main task here is parameter conversion, such as converting a Python dict to a C++ std::unordered_map. Let’s look at the most direct implementation first.

3.2.1 py_bpe_train_step2

cpdef py_bpe_train_step2(int vocab_size,
                             pair_counts_py,
                             pair_strings_py,
                             vocabulary_py,
                             pair_wordids_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef HashMap[pair[int, int], int, pair_hash] pair_counts_cpp
    cdef unordered_map[pair[int, int], vector[vector[int]], pair_hash] pair_strings_cpp
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[pair[int, int], unordered_set[int], pair_hash] pair_wordids_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    cdef pair[int, int] pair_key
    cdef vector[vector[int]] strings_value
    cdef vector[int] vector_value
    cdef unordered_set[int] set_value




    for p, count in pair_counts_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        pair_counts_cpp[pair_key] = count

    for p, string in pair_strings_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        strings_value.clear()
        value = [list(item) for item in string] 
        vector_value = value[0]
        strings_value.push_back(vector_value)
        vector_value = value[1]
        strings_value.push_back(vector_value)        
        pair_strings_cpp[pair_key] = strings_value     
    
    for k, v in vocabulary_py.items():
        value = list(v)
        vector_value = value
        vocabulary_cpp[k] = vector_value

    for p, wordids in pair_wordids_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]        
        set_value = wordids
        pair_wordids_cpp[pair_key] = set_value

    for k, v in wordid_counts_py.items():
        wordid_counts_cpp[k] = v

    for k, v in wordid_encodings_py.items():
        vector_value = v
        wordid_encodings_cpp[k] = vector_value

    # 调用 C++ 函数
    bpe_train_step2(vocab_size,
                    pair_counts_cpp,
                    pair_strings_cpp,
                    vocabulary_cpp,
                    pair_wordids_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

Most of the code in this function converts Python variables to C++ variables before calling bpe_train_step2. Let’s look at two typical examples:

    cdef HashMap[pair[int, int], int, pair_hash] pair_counts_cpp
    cdef pair[int, int] pair_key

    for p, count in pair_counts_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        pair_counts_cpp[pair_key] = count

The key for pair_counts_cpp is pair[int,int], which is std::pair<int,int>. We can assign to its first and second members. We then use the operator[] to insert, which is why we needed to declare V& operator[](const K& key) earlier. The code iterates like a Python for loop.

Let’s look at a more complex example:

    cdef unordered_map[pair[int, int], vector[vector[int]], pair_hash] pair_strings_cpp

    for p, string in pair_strings_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        strings_value.clear()
        value = [list(item) for item in string] 
        vector_value = value[0]
        strings_value.push_back(vector_value)
        vector_value = value[1]
        strings_value.push_back(vector_value)        
        pair_strings_cpp[pair_key] = strings_value   

The key for pair_strings_py is a tuple, and the value is also a tuple with two bytes elements. For example:

(111,110): (b'o', b'n')

The corresponding pair_strings_cpp is unordered_map[pair[int, int], vector[vector[int]], pair_hash], so we need to convert the bytes to list[int]. This is done with the following statements:

value = [list(item) for item in string] 
vector_value = value[0]
strings_value.push_back(vector_value)

First, a list comprehension converts tuple[bytes] to list[list[int]]. Then, Python’s list[int] is converted to C++’s vector[int] through assignment, and finally pushed back to strings_value.

3.2.2 py_bpe_train_step2_v2

cpdef py_bpe_train_step2_v2(int vocab_size,
                             vocabulary_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    cdef pair[int, int] pair_key
    cdef vector[vector[int]] strings_value
    cdef vector[int] vector_value
    cdef unordered_set[int] set_value
 
    
    for k, v in vocabulary_py.items():
        value = list(v)
        vector_value = value
        vocabulary_cpp[k] = vector_value


    for k, v in wordid_counts_py.items():
        wordid_counts_cpp[k] = v

    for k, v in wordid_encodings_py.items():
        vector_value = v
        wordid_encodings_cpp[k] = vector_value

    # 调用 C++ 函数
    bpe_train_step2_v2(vocab_size,
                    vocabulary_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

This version is similar to the previous one, but with three fewer parameters.

Note: We return C++ variables merges_cpp and vocabulary_cpp at the end. Their types are:

cdef unordered_map[int, vector[int]] vocabulary_cpp
cdef vector[pair[vector[int], vector[int]]] merges_cpp

When they are returned to Python, Cython automatically converts them to dict[int, list[int]] and list[tuple[list[int], list[int]]]. We’ll need to convert the list[int] back to bytes later.

In fact, not just for return values, but also when we copy a Python variable to a C++ variable or vice versa, Cython automatically performs these common conversions:

Python type =>	C++ type	=> Python type
bytes	std::string	bytes
iterable	std::vector	list
iterable	std::list	list
iterable	std::set	set
iterable	std::unordered_set	set
mapping	std::map	dict
mapping	std::unordered_map	dict
iterable (len 2)	std::pair	tuple (len 2)
complex	std::complex	complex

We can use this feature to simplify the variable conversions and get py_bpe_train_step2_opt:

cpdef py_bpe_train_step2_opt(int vocab_size,
                             vocabulary_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    vocabulary_cpp = vocabulary_py

    wordid_counts_cpp = wordid_counts_py

    wordid_encodings_cpp = wordid_encodings_py
    # 调用 C++ 函数
    bpe_train_step2_v2(vocab_size,
                    vocabulary_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

Here, we use just three assignment statements, and Cython automatically handles the conversion between Python and C++. The calls to other versions, like py_bpe_train_step2_v3, are identical to py_bpe_train_step2_opt:

cpdef py_bpe_train_step2_v3(int vocab_size,
                             vocabulary_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    vocabulary_cpp = vocabulary_py

    wordid_counts_cpp = wordid_counts_py

    wordid_encodings_cpp = wordid_encodings_py
    # 调用 C++ 函数
    bpe_train_step2_v3(vocab_size,
                    vocabulary_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

3.3 Modifying setup.py

project_root = os.path.dirname(os.path.abspath(__file__))


ext_modules = [
    Extension(
        name="cs336_basics.bpe_train_step2_wrapper",
        sources=["cs336_basics/bpe_train_step2_wrapper.pyx"],

        language="c++",
        #extra_compile_args=['-std=c++17', '-O3'],
        extra_compile_args=['-std=c++17'],
        libraries=["bpe_train_step2"],

        library_dirs=[f"{project_root}/lib_bpe_train_step2/lib"],
        runtime_library_dirs=[f"{project_root}/lib_bpe_train_step2/lib"],
        include_dirs=[f"{project_root}/lib_bpe_train_step2/include",
                      f"{project_root}/lib_bpe_train_step2/include/emhash"],
    )
]

setup(
    packages=['cs336_basics'],
    name='bpe_train_step2',
    ext_modules=cythonize(ext_modules),
)

We need to compile bpe_train_step2_wrapper.pyx.

name specifies the module name, allowing Python to import cs336_basics.bpe_train_step2_wrapper. Here, cs336_basics is the package name, and bpe_train_step2_wrapper is the module name.
sources specifies the source files to be compiled.
language specifies the language of the module as C++.
extra_compile_args specifies additional compilation options, in this case '-std=c++17'.
libraries lists the dependent libraries.
library_dirs specifies the library location for compilation.
runtime_library_dirs specifies the library location for runtime.
include_dirs specifies the header file location for compilation.

To avoid hard-coding, project_root is set to the directory where setup.py is located.

These options ultimately result in the following C++ compiler commands (for my GCC environment; they may vary):

c++ -pthread -fno-strict-overflow -Wsign-compare -Wunreachable-code -DNDEBUG -g -O3 -Wall -fPIC -fPIC -Ics336_basics -I......codes/assignment1-basics-bpe/lib_bpe_train_step2/include -I......codes/assignment1-basics-bpe/lib_bpe_train_step2/include/emhash -I......codes/assignment1-basics-bpe/.venv/include -I.......local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/include/python3.12 -c cs336_basics/bpe_train_step2_wrapper.cpp -o build/temp.linux-x86_64-cpython-312/cs336_basics/bpe_train_step2_wrapper.o -std=c++17
c++ -pthread -fno-strict-overflow -Wsign-compare -Wunreachable-code -DNDEBUG -g -O3 -Wall -fPIC -shared -Wl,--exclude-libs,ALL build/temp.linux-x86_64-cpython-312/cs336_basics/bpe_train_step2_wrapper.o -L......codes/assignment1-basics-bpe/lib_bpe_train_step2/lib -L.......local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib -Wl,--enable-new-dtags,-rpath,......codes/assignment1-basics-bpe/lib_bpe_train_step2/lib -lbpe_train_step2 -o build/lib.linux-x86_64-cpython-312/cs336_basics/bpe_train_step2_wrapper.cpython-312-x86_64-linux-gnu.so

3.4 Using it in Python

bpe_v9.py calls py_bpe_train_step2; bpe_v10.py calls py_bpe_train_step2_v2; bpe_v10_v2.py calls py_bpe_train_step2_opt; bpe_v11.py calls py_bpe_train_step2_v3; bpe_v11_bytes.py calls py_bpe_train_step2_v3; bpe_v11_v2.py calls py_bpe_train_step2_v4; bpe_v11_v3.py calls py_bpe_train_step2_v5; bpe_v11_v3_bytes.py calls py_bpe_train_step2_v5; bpe_v11_v4.py calls py_bpe_train_step2_v6; bpe_v11_v4_bytes.py calls py_bpe_train_step2_v6.

Their code is mostly the same. Let’s look at bpe_v11.py:

        vocabulary = {i: bytes([i]) for i in range(N_BYTES)} # every byte
        for i, token in enumerate(special_tokens):
            vocabulary[N_BYTES + i] = token.encode('utf-8')
        size = N_BYTES + len(special_tokens)
        merges = []

        # initial word encodings are utf-8
        word_encodings = {}
        for word in word_counts:
            word_encodings[word] = list(word.encode('utf-8'))

        word_ids = {word:id for id, word in enumerate(word_counts)}

        wordid_counts = {word_ids[word]:count for word, count in word_counts.items()}

        wordid_encodings = {word_ids[word]:encoding for word, encoding in word_encodings.items()}      


        merges_cpp, vocabulary_cpp = py_bpe_train_step2_v3(vocab_size, 
                             vocabulary,
                             wordid_counts,
                             wordid_encodings,
                             merges)


        vocabulary = {k:bytes(v) for k, v in vocabulary_cpp.items()}
        merges = [(bytes(arr[0]), bytes(arr[1])) for arr in merges_cpp]

Before the call, we need to convert string words to integer IDs. This mapping is stored in word_ids. We then use word_ids to convert word_counts to wordid_counts and word_encodings to wordid_encodings. After the call, the returned merges_cpp and vocabulary_cpp need to have their list[int] converted back to bytes.

4. Testing

Version	Data	Total Time (s)	Tokenization Time (s)	Merging Time (s)	Other
bpe_v8_v3	open_web	897/899/951	395/399/395	442/438/493	num_counter=8, num_merger=1
bpe_v9	open_web	831/816/867	400/401/390	prepare & convert: 93/86/94 py_bpe_train_step2: 330/320/374 c++:289/281/326	num_counter=8, num_merger=1
bpe_v10	open_web	816/769/788	390/400/400	prepare & convert: 21/17/19 py_bpe_train_step2: 402/350/367 c++:338/296/309	num_counter=8, num_merger=1
bpe_v10_v2	open_web	767/774/767	400/401/401	prepare & convert: 18/17/17 py_bpe_train_step2: 346/355/347 c++:292/298/294	num_counter=8, num_merger=1
bpe_v10_v2	open_web	498/477/495	120/120/120	prepare & convert: 20/19/18 py_bpe_train_step2: 355/336/354 c++:299/282/298	num_counter=32, num_merger=4
bpe_v11	open_web	350/340/354	120/120/120	prepare & convert: 21/19/19 py_bpe_train_step2: 207/199/212 c++:183/175/190	num_counter=32, num_merger=4
bpe_v11_bytes	open_web	311/307/305	80/80/80	prepare & convert: 18/19/18 py_bpe_train_step2: 211/206/204 c++:189/183/182	num_counter=64, num_merger=8, chunk_size 8mb
bpe_v11_v2	open_web	362/350/338	130/120/120	prepare & convert: 19/18/19 py_bpe_train_step2: 210/210/197 c++:189/190/176	num_counter=32, num_merger=4
bpe_v11_v3	open_web	269/274/270	120/120/120	prepare & convert: 18/19/18 py_bpe_train_step2: 129/133/129 c++: 106/109/106	num_counter=32, num_merger=4
bpe_v11_v3_bytes	open_web	218/219/215	72/74/69	prepare & convert: 21/21/21 py_bpe_train_step2: 123/122/123 c++: 101/100/101	num_counter=64, num_merger=8, chunk_size 8mb
bpe_v11_v4	open_web	258/256/261	116/117/117	prepare & convert: 19/18/19 py_bpe_train_step2: 121/119/123 c++: 98/97/100	num_counter=32, num_merger=4
bpe_v11_v4_bytes	open_web	210/206/207	71/69/70	prepare & convert: 20/18/19 py_bpe_train_step2: 117/117/117 c++: 95/96/95	num_counter=64, num_merger=8, chunk_size 8mb

Comparing bpe_v10 and bpe_v10_v2, the difference is between manually converting parameters in Python to C++ and letting Cython do it automatically. Automatic conversion is not only more convenient but also faster.

Ultimately, with bpe_v11_v4_bytes on 64 cores, the fastest training time is a little over 200 seconds, which is more than 100 times faster than our initial time of over ten hours!

5. Conclusion

This concludes the series. Here is a brief summary and links to the corresponding articles:

Part 0: Introduction Introduces the basic BPE training algorithm and related tasks, as well as the development environment.
Part 1: The Simplest Implementation The simplest implementation of BPE training.
Part 2: Optimized Algorithm Implements incremental updates for pair_counts.
Part 3: Parallel Tokenization and Frequency Counting Uses multiprocessing to implement a multi-process parallel algorithm.
Part 4: A Failed Parallel Optimization An attempt to parallelize the max pair calculation using multiple processes.
Part 5: Implementing the Merge Algorithm in C++ Implements a C++ merge algorithm equivalent to the Python version, and compares two ways of iterating through std::unordered_map.
Part 6: Parallelizing the Max Pair Search with OpenMP Uses OpenMP to find the max pair in pair_counts in parallel.
Part 7: Using Flat Hashmap to Replace std::unordered_map Uses flat hashmap to replace std::unordered_map.
Part 8: Implementing Fine-Grained Updates Implements a fine-grained update algorithm for pair_counts using an inverted index.
Part 9: Using a Heap to Find the Max Pair Uses a heap to find the max pair and improve performance.
Part 10: Using Cython and PyPy for Acceleration Uses Cython and PyPy to accelerate Python code.
Part 11: Wrapping C++ Code with Cython Wraps C++ code using Cython.

动手实现和优化BPE Tokenizer的训练——第11部分：使用cython封装c++代码

Thu, 25 Sep 2025 00:00:00 +0000

本系列文章完成Stanford CS336作业1的一个子任务——实现BPE Tokenizer的高效训练算法。通过一系列优化，我们的算法在OpenWebText上的训练时间从最初的10多个小时优化到小于10分钟。本系列文章解释这一系列优化过程，包括：算法的优化，数据结构的优化，并行(openmp)优化，cython优化，用c++实现关键代码和c++库的cython集成等内容。本文是第十二篇，也是最后一篇，使用cython把之前的c++代码封装成扩展模块供python调用。

1. 目标
2. 把c++代码封装成一个动态库
3. 用cython封装成扩展模块
4. 测试
5. 总结

1. 目标

前面的文章我们已经做了很多探索，bpe训练可以分为两个阶段。第一个阶段是分词和统计，我们通过multiprocessing进行并行计算，实现的bpe_v3.py利用32核可以在120秒内完成。而读取bytes的bpe_v3_bytes_time.py可以使用64核在70秒内完成。第二个阶段是合并，python版本的bpe_v8_v3.py可以在500秒内完成，总时间是十多分钟。c++版本的合并最快的是bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp，它的合并时间是100秒左右。

今天我们的目标就是实现一个最快的版本，通过cython封装c++代码成扩展模块，这样第二个阶段就可以调用这个模块。

2. 把c++代码封装成一个动态库

我们在cppupdate里做了很多实验，现在我们需要把结果最好的一些版本封装成动态库，这样方便后续cython使用。

我们会新建一个cppstep2的c++项目，它的完整代码在cppstep2。

根据cppupdate的实验结果，我选择了bpe_train_updater_fine_grained_emhash8.cpp、bpe_train_updater_fine_grained_emhash8_set.cpp、bpe_train_updater_fine_grained_emhash8_set9.cpp、bpe_train_updater_fine_grained_heap_emhash8_set9.cpp和bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp。

选择emhash而不是absl的原因有二：一是它比较快；二是用它比较简单。我之前是git clone了emhash的完整代码，这里为了简化，我只是复制了需要的3个头文件：hash_set8.hpp, hash_set4.hpp和hash_table8.hpp。另外max_heap.h和max_heap.cpp也从cppupdate项目原封不动的复制了过来。

我们的动态库主要就两个文件：头文件bpe_train_step2.h和实现文件bpe_train_step2.cpp。

2.1 头文件

#pragma once

#include <vector>
#include <unordered_map>
#include <unordered_set>
#include "emhash/hash_table8.hpp"

struct pair_hash {
    template <class T1, class T2>
    std::size_t operator () (const std::pair<T1,T2> &p) const {
        auto h1 = std::hash<T1>{}(p.first);
        auto h2 = std::hash<T2>{}(p.second);

        return h1 ^ (h2 + 0x9e3779b9 + (h1 << 6) + (h1 >> 2));
    }
};

void bpe_train_step2(int vocab_size, 
                emhash8::HashMap<std::pair<int, int>, int, pair_hash> & pair_counts, 
                std::unordered_map<std::pair<int, int>, std::vector<std::vector<int>>, pair_hash> & pair_strings, 
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash> & pair_wordids, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);

void bpe_train_step2_v2(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);

void bpe_train_step2_v3(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);   
                
void bpe_train_step2_v4(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);  

void bpe_train_step2_v5(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges);  

void bpe_train_step2_v6(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges); 

这个头文件定义了bpe_train_step2_v2、bpe_train_step2_v3、bpe_train_step2_v4、bpe_train_step2_v5和bpe_train_step2_v6这5个函数，它们分别对应bpe_train_updater_fine_grained_emhash8.cpp、bpe_train_updater_fine_grained_emhash8_set.cpp、bpe_train_updater_fine_grained_emhash8_set9.cpp、bpe_train_updater_fine_grained_heap_emhash8_set9.cpp和bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp。

bpe_train_step2是完全参考cppudate的实现，它的输入参数是6个。而后面我们会讲到pair_counts是wordid_counts是基于计算出来的，可以从python移植到c++，这样速度更快，而且少传递很多参数。所以bpe_train_step2_v2以及之后的版本都是只有5个参数。

2.2 CMakeLists.txt

我们的目标是编译一个动态库，按照CMake的语法编写如下内容：

cmake_minimum_required(VERSION 3.20)

project(BPE_TRAIN_STEP2 LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)


add_library(bpe_train_step2 SHARED bpe_train_step2.cpp max_heap.cpp)
target_include_directories(bpe_train_step2 PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
target_compile_definitions(bpe_train_step2 PUBLIC MY_LIBRARY_EXPORT)

target_include_directories(bpe_train_step2 PUBLIC
    "${PROJECT_SOURCE_DIR}/emhash"
)

install(TARGETS bpe_train_step2
        EXPORT bpe_train_step2_export_targets
        RUNTIME DESTINATION bin
        LIBRARY DESTINATION lib
        ARCHIVE DESTINATION lib)

install(FILES "bpe_train_step2.h" 
        DESTINATION include
)

我们通过add_library创建bpe_train_step2这个target。接着的target_compile_definitions好像是为了windows的兼容性定义的宏。我不懂windows，也没有windows环境测试，所以不知道这个项目能不能在windows下编译。

然后的target_include_directories把emhash的头文件包含进来。最后是安装的时候把bpe_train_step2这个target(主要是libbpe_train_step2.so)和bpe_train_step2.h复制到合适的位置，后面我会讲到怎么安装。

2.3 bpe_train_step2

首先我们按照cppupdate的方式集成，我们需要传入pair_counts等7个参数。

void bpe_train_step2(int vocab_size, 
                emhash8::HashMap<std::pair<int, int>, int, pair_hash> & pair_counts, 
                std::unordered_map<std::pair<int, int>, std::vector<std::vector<int>>, pair_hash> & pair_strings, 
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash> & pair_wordids, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges){
    auto start = std::chrono::steady_clock::now();
    while(vocabulary.size() < vocab_size){
        int max_count = -1;
        std::pair<int, int> max_pair;
        std::vector<std::vector<int>> max_strings;
        for(const auto& [pair, count] : pair_counts){
            if(count > max_count){
                max_count = count;
                max_pair = pair;
                max_strings = pair_strings[pair];
            }else if(count == max_count){
                std::vector<std::vector<int>> strings = pair_strings[pair];
                ComparisonResult r1 = three_way_compare(strings[0], max_strings[0]);
                if(r1 == ComparisonResult::Greater){
                    max_count = count;
                    max_pair = pair;
                    max_strings = strings;
                }else if(r1 == ComparisonResult::Equal){
                    ComparisonResult r2 = three_way_compare(strings[1], max_strings[1]);
                    if(r2 == ComparisonResult::Greater){
                        max_count = count;
                        max_pair = pair;
                        max_strings = strings;                        
                    }
                }
            }
        }

        const std::vector<int>& bytes1 = vocabulary[max_pair.first];
        const std::vector<int>& bytes2 = vocabulary[max_pair.second];
        std::vector<int> merge_bytes;
        merge_bytes.reserve(bytes1.size() + bytes2.size());
        merge_bytes.insert(merge_bytes.end(), bytes1.begin(), bytes1.end());
        merge_bytes.insert(merge_bytes.end(), bytes2.begin(), bytes2.end());

        int size = vocabulary.size();
        vocabulary[size] = merge_bytes;

        auto& affected_words = pair_wordids[max_pair];

        
        updated_affected_word_count(max_pair, affected_words, wordid_encodings, wordid_counts,
                                    pair_counts, pair_wordids, size, pair_strings, vocabulary);
  
        merges.push_back({bytes1, bytes2});


    }
    auto end = std::chrono::steady_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
    std::cout << "bpe_train_step2: " << duration.count() << "ms." << std::endl;
}

代码完全是复制bpe_train_updater_fine_grained_emhash8.cpp。

2.4 bpe_train_step2_v2

如果读者仔细看过之前的时间统计，会发现总时间会比”统计词频时间” + “合并时间” 要多五六十秒。这个时间就是调用BPE_Trainer._count_pairs的时间：

    @staticmethod    
    def _count_pairs(word_counts, word_encodings, pair_strings, vocabulary, pair_to_words):
        pair_counts = defaultdict(int)
        for word, count in word_counts.items():
            encoding = word_encodings[word]
            for i in range(0, len(encoding) - 1):
                pair = encoding[i], encoding[i + 1]
                pair_counts[pair] += count
                if pair not in pair_strings:
                    pair_strings[pair] = (vocabulary[pair[0]], vocabulary[pair[1]])

                pair_to_words[pair].add(word)

        return pair_counts

这段代码的作用就是根据word_counts统计pair_counts，构建初始的倒排索引pair_to_words以及pair_strings。它的输入参数是word_counts, word_encodings和vocabulary。这段代码完全可以在c++里实现，这样一来可以加速，二来可以减少python和c++之间的参数传递。

因此bpe_train_step2_v2的输入参数就可以减少三个：

void bpe_train_step2_v2(int vocab_size,          
                std::unordered_map<int, std::vector<int>> & vocabulary, 
                const std::unordered_map<int, long long> & wordid_counts, 
                std::unordered_map<int, std::vector<int>> & wordid_encodings, 
                std::vector<std::pair<std::vector<int>, std::vector<int>>> & merges){
    std::unordered_map<std::pair<int, int>, std::vector<std::vector<int>>, pair_hash> pair_strings;
    emhash8::HashMap<std::pair<int, int>, int, pair_hash>  pair_counts;
    std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash>  pair_wordids;
    
    std::pair<int, int> pair;
    for(const auto& [wordid, count] : wordid_counts){
        const auto& word_tokens = wordid_encodings[wordid];
        for(int i = 0; i < word_tokens.size() - 1; ++i){
            pair.first = word_tokens[i];
            pair.second = word_tokens[i + 1];
            pair_counts[pair] += count;
            if (pair_strings.find(pair) == pair_strings.end()) {
                pair_strings[pair] = {vocabulary[pair.first], vocabulary[pair.second]};
            }
            pair_wordids[pair].insert(wordid);
        }
    }
    

    bpe_train_step2(vocab_size, pair_counts, pair_strings, vocabulary,
                    pair_wordids, wordid_counts, wordid_encodings, merges);
}

bpe_train_step2_v2这个函数首先根据wordid_counts、wordid_encodings和vocabulary计算得到pair_counts、pair_strings和pair_wordids，最后还是调用bpe_train_step2。

bpe_train_step2_v3和之后的版本接口与bpe_train_step2_v2完全一样，代码也是复制相应的cppupdate里的代码，这里就不赘述了。

2.5 编译安装

使用如下命令编译安装：

cd cppstep2
mkdir build && cd build
cmake -D CMAKE_INSTALL_PREFIX=../../lib_bpe_train_step2/  -D CMAKE_BUILD_TYPE=Release ..
cmake --build . -- -j8
# 如果是gcc的话可以直接make -j8
cmake --install .  
# 或者make install

由于我们不希望把这个库安装到系统的路径，比如/usr/local下，因为这样的安装需要root权限。而且不同系统的路径不相同，也不便于cython的集成。所以我这里把它安装到../../lib_bpe_train_step2/，也就是项目根目录assignment1-basics-bpe之下的lib_bpe_train_step2：

$ ls assignment1-basics-bpe

cppstep2/
cppupdate/
cs336_basics/
data/
lib_bpe_train_step2/

安装后的lib_bpe_train_step2为：

lib_bpe_train_step2$ tree
.
├── include
│   ├── bpe_train_step2.h
│   └── emhash
│       └── hash_table8.hpp
└── lib
    └── libbpe_train_step2.so

除了bpe_train_step2.h，hash_table8.hpp也被包含了进来，因为bpe_train_step2.h要用到它。这主要是因为bpe_train_step2的声明要用到它。如果我们只保留bpe_train_step2_v2及其之后的版本，那么这个头文件也不需要安装(但是我们编译还是需要emhash下的所有头文件)。

3. 用cython封装成扩展模块

对于cython封装c++库不熟悉的读者可以参考Using C++ in Cython。

我们需要两个文件：bpe_train_step2_wrapper.pxd和bpe_train_step2_wrapper.pyx。

3.1 bpe_train_step2_wrapper.pxd

首先来看pxd文件：

# distutils: language = c++

# 导入 C++ 标准库类型
from libcpp.utility cimport pair
from libcpp.vector cimport vector
from libcpp.unordered_map cimport unordered_map
from libcpp.unordered_set cimport unordered_set


cdef extern from "../lib_bpe_train_step2/include/bpe_train_step2.h" :
    cppclass pair_hash:
        pass

cdef extern from "../lib_bpe_train_step2/include/emhash/hash_table8.hpp" namespace "emhash8":
    cppclass HashMap[K, V, H]:
        #ValueT& operator[](const KeyT& key) noexcept
        V& operator[](const K& key)


cdef extern from "../lib_bpe_train_step2/include/bpe_train_step2.h":
    void bpe_train_step2(int vocab_size,
                         HashMap[pair[int, int], int, pair_hash] & pair_counts,
                         unordered_map[pair[int, int], vector[vector[int]], pair_hash] & pair_strings,
                         unordered_map[int, vector[int]] & vocabulary,
                         unordered_map[pair[int, int], unordered_set[int], pair_hash] & pair_wordids,
                         const unordered_map[int, long long] & wordid_counts,
                         unordered_map[int, vector[int]] & wordid_encodings,
                         vector[pair[vector[int], vector[int]]] & merges) except +
    
    void bpe_train_step2_v2(int vocab_size,
                         unordered_map[int, vector[int]] & vocabulary,
                         const unordered_map[int, long long] & wordid_counts,
                         unordered_map[int, vector[int]] & wordid_encodings,
                         vector[pair[vector[int], vector[int]]] & merges) except +

第一行用来告诉 distutils(deprecated)或setuptools 模块：在这个脚本中，你正在构建一个 C++ 扩展模块，而不是默认的 C 语言扩展。

接下来4行是我们要用到c++标准库的std::vector、std::pair、std::unordered_map和std::unordered_set。cython把这些常见的标准库都封装到libcpp下了，我们只需要用cimport(类似python的import，不过这是编译时而不是运行时的import)导入即可使用。要查看哪些c++标准库可以在cython直接使用可以参考这里。

接下来用cppclass声明编译时要用到的pair_hash，因为我们使用lib_bpe_train_step2只需要知道pair_hash这个符号即可，它的内容我们不需要，所以它的内容就是一行pass。注意使用cdef extern from的路径是”../lib_bpe_train_step2/include/bpe_train_step2.h”，这就要求lib_bpe_train_step2按照前面的步骤安装到了合适的位置。

接下来是声明emhash8::HashMap，在cdef extern from后有一个”namespace emhash8”，这样cython知道HashMap是在emhash8这个namespace下。cppclass HashMap[K, V, H]说明HashMap是一个模板类，分别代表Key/Value/Hash函数。此外后面的pyx文件里我们会用到operator[]，所以我们也需要声明这个重载运算符的原型。

最后就是声明bpe_train_step2这些函数了，cython和c++的语法类似，只不过<>要改成[]，看起来有点别扭。

3.2 bpe_train_step2_wrapper.pyx

接下来就是把c++的函数封装成python可以调用的函数，这里主要做的就是参数的转换，比如把python的dict变成c++的std::unordered_map。我们先看最直接的实现。

3.2.1 py_bpe_train_step2

cpdef py_bpe_train_step2(int vocab_size,
                             pair_counts_py,
                             pair_strings_py,
                             vocabulary_py,
                             pair_wordids_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef HashMap[pair[int, int], int, pair_hash] pair_counts_cpp
    cdef unordered_map[pair[int, int], vector[vector[int]], pair_hash] pair_strings_cpp
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[pair[int, int], unordered_set[int], pair_hash] pair_wordids_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    cdef pair[int, int] pair_key
    cdef vector[vector[int]] strings_value
    cdef vector[int] vector_value
    cdef unordered_set[int] set_value




    for p, count in pair_counts_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        pair_counts_cpp[pair_key] = count

    for p, string in pair_strings_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        strings_value.clear()
        value = [list(item) for item in string] 
        vector_value = value[0]
        strings_value.push_back(vector_value)
        vector_value = value[1]
        strings_value.push_back(vector_value)        
        pair_strings_cpp[pair_key] = strings_value     
    
    for k, v in vocabulary_py.items():
        value = list(v)
        vector_value = value
        vocabulary_cpp[k] = vector_value

    for p, wordids in pair_wordids_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]        
        set_value = wordids
        pair_wordids_cpp[pair_key] = set_value

    for k, v in wordid_counts_py.items():
        wordid_counts_cpp[k] = v

    for k, v in wordid_encodings_py.items():
        vector_value = v
        wordid_encodings_cpp[k] = vector_value

    # 调用 C++ 函数
    bpe_train_step2(vocab_size,
                    pair_counts_cpp,
                    pair_strings_cpp,
                    vocabulary_cpp,
                    pair_wordids_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

这个函数前面的大部分代码都是把python的变量转换成c++的变量，然后调用bpe_train_step2。我们来看两个典型的例子：

    cdef HashMap[pair[int, int], int, pair_hash] pair_counts_cpp
    cdef pair[int, int] pair_key

    for p, count in pair_counts_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        pair_counts_cpp[pair_key] = count

pair_counts_cpp的key是pair[int,int]，也就是std::pair<int,int>，我们可以通过first和second对它赋值。最后通过运算符[]对它进行插入，这也是之前我们需要声明V& operator的原因。我们可以发现上面的代码是类似python的for遍历。

再看一个复杂一点的：

    cdef unordered_map[pair[int, int], vector[vector[int]], pair_hash] pair_strings_cpp

    for p, string in pair_strings_py.items():
        pair_key.first = p[0]
        pair_key.second = p[1]
        strings_value.clear()
        value = [list(item) for item in string] 
        vector_value = value[0]
        strings_value.push_back(vector_value)
        vector_value = value[1]
        strings_value.push_back(vector_value)        
        pair_strings_cpp[pair_key] = strings_value   

pair_strings_py的key是一个tuple，value也是一个tuple，这个tuple有2个元素，每个都是一个bytes。说起来很费劲，我们看一个例子：

(111,110): (b'o', b'n')

我们对应的pair_strings_cpp是unordered_map[pair[int, int], vector[vector[int]], pair_hash]，所以需要把bytes变成list[int]。这是通过下面的语句实现：

value = [list(item) for item in string] 
vector_value = value[0]
strings_value.push_back(vector_value)

首先是用列表推导把tuple[bytes]变成list[[list[int]]]，然后通过赋值把python的list[int]变成c++的vector[int]，最后push_back到strings_value里。

3.2.2 py_bpe_train_step2_v2

cpdef py_bpe_train_step2_v2(int vocab_size,
                             vocabulary_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    cdef pair[int, int] pair_key
    cdef vector[vector[int]] strings_value
    cdef vector[int] vector_value
    cdef unordered_set[int] set_value
 
    
    for k, v in vocabulary_py.items():
        value = list(v)
        vector_value = value
        vocabulary_cpp[k] = vector_value


    for k, v in wordid_counts_py.items():
        wordid_counts_cpp[k] = v

    for k, v in wordid_encodings_py.items():
        vector_value = v
        wordid_encodings_cpp[k] = vector_value

    # 调用 C++ 函数
    bpe_train_step2_v2(vocab_size,
                    vocabulary_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

这个版本和之前差不多，只不过少了3个参数。

注意：最后我们返回的是c++的变量merges_cpp, vocabulary_cpp，它们的类型是：

cdef unordered_map[int, vector[int]] vocabulary_cpp
cdef vector[pair[vector[int], vector[int]]] merges_cpp

返回到python时，cython会自动把它转换成dict[int,list[int]]和list[tuple[list[int], list[int]]]。我们后面需要再把list[int]变成bytes。

其实不只是返回值，我们把一个python变量复制给一个c++变量或者把一个c++变量复制给python时cython也会自动的做这些常见的转换：

Python type =>	C++ type	=> Python type
bytes	std::string	bytes
iterable	std::vector	list
iterable	std::list	list
iterable	std::set	set
iterable	std::unordered_set	set
mapping	std::map	dict
mapping	std::unordered_map	dict
iterable (len 2)	std::pair	tuple (len 2)
complex	std::complex	complex

我们可以用这个特性来简化变量之间的转换，得到py_bpe_train_step2_opt：

cpdef py_bpe_train_step2_opt(int vocab_size,
                             vocabulary_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    vocabulary_cpp = vocabulary_py

    wordid_counts_cpp = wordid_counts_py

    wordid_encodings_cpp = wordid_encodings_py
    # 调用 C++ 函数
    bpe_train_step2_v2(vocab_size,
                    vocabulary_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

这里我们直接用3条复制语句，cython自动就会帮我们在python和c++之间进行转换。其它版本的调用都是和py_bpe_train_step2_opt一样，比如py_bpe_train_step2_v3：

cpdef py_bpe_train_step2_v3(int vocab_size,
                             vocabulary_py,
                             wordid_counts_py,
                             wordid_encodings_py,
                             merges_py):

    # 声明 C++ 容器
    cdef unordered_map[int, vector[int]] vocabulary_cpp
    cdef unordered_map[int, long long] wordid_counts_cpp
    cdef unordered_map[int, vector[int]] wordid_encodings_cpp
    cdef vector[pair[vector[int], vector[int]]] merges_cpp

 
    vocabulary_cpp = vocabulary_py

    wordid_counts_cpp = wordid_counts_py

    wordid_encodings_cpp = wordid_encodings_py
    # 调用 C++ 函数
    bpe_train_step2_v3(vocab_size,
                    vocabulary_cpp,
                    wordid_counts_cpp,
                    wordid_encodings_cpp,
                    merges_cpp)

    return merges_cpp, vocabulary_cpp

3.3 修改setup.py

project_root = os.path.dirname(os.path.abspath(__file__))


ext_modules = [
    Extension(
        name="cs336_basics.bpe_train_step2_wrapper",
        sources=["cs336_basics/bpe_train_step2_wrapper.pyx"],

        language="c++",
        #extra_compile_args=['-std=c++17', '-O3'],
        extra_compile_args=['-std=c++17'],
        libraries=["bpe_train_step2"],

        library_dirs=[f"{project_root}/lib_bpe_train_step2/lib"],
        runtime_library_dirs=[f"{project_root}/lib_bpe_train_step2/lib"],
        include_dirs=[f"{project_root}/lib_bpe_train_step2/include",
                      f"{project_root}/lib_bpe_train_step2/include/emhash"],
    )
]

setup(
    packages=['cs336_basics'],
    name='bpe_train_step2',
    ext_modules=cythonize(ext_modules),
)

我们需要编译bpe_train_step2_wrapper.pyx。

name是指定模块的名字，这样python里可以import cs336_basics.bpe_train_step2_wrapper。这里的cs336_basics是package的名字，bpe_train_step2_wrapper是模块名。
sources指定要编译的源代码
language指定模块是c++语言
extra_compile_args指定额外的编译选项，这里指定了’-std=c++17’
libraries是依赖的库
library_dirs指定编译时库的位置
runtime_library_dirs指定运行时库的位置
include_dirs指定编译时的头文件位置

为了避免硬编码，project_root为setup.py文件所在的目录。

这些选项最终会使得c++编译器的编译和链接命令为(我的gcc环境，不同环境可能会有差异)：

c++ -pthread -fno-strict-overflow -Wsign-compare -Wunreachable-code -DNDEBUG -g -O3 -Wall -fPIC -fPIC -Ics336_basics -I......codes/assignment1-basics-bpe/lib_bpe_train_step2/include -I......codes/assignment1-basics-bpe/lib_bpe_train_step2/include/emhash -I......codes/assignment1-basics-bpe/.venv/include -I.......local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/include/python3.12 -c cs336_basics/bpe_train_step2_wrapper.cpp -o build/temp.linux-x86_64-cpython-312/cs336_basics/bpe_train_step2_wrapper.o -std=c++17
c++ -pthread -fno-strict-overflow -Wsign-compare -Wunreachable-code -DNDEBUG -g -O3 -Wall -fPIC -shared -Wl,--exclude-libs,ALL build/temp.linux-x86_64-cpython-312/cs336_basics/bpe_train_step2_wrapper.o -L......codes/assignment1-basics-bpe/lib_bpe_train_step2/lib -L.......local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib -Wl,--enable-new-dtags,-rpath,......codes/assignment1-basics-bpe/lib_bpe_train_step2/lib -lbpe_train_step2 -o build/lib.linux-x86_64-cpython-312/cs336_basics/bpe_train_step2_wrapper.cpython-312-x86_64-linux-gnu.so

3.4 在python里使用

bpe_v9.py调用的是py_bpe_train_step2函数；bpe_v10.py调用的是py_bpe_train_step2_v2；bpe_v10_v2.py调用的是py_bpe_train_step2_opt；bpe_v11.py调用的是py_bpe_train_step2_v3；bpe_v11_bytes.py调用的是py_bpe_train_step2_v3；bpe_v11_v2.py调用的是py_bpe_train_step2_v4；bpe_v11_v3.py调用的是py_bpe_train_step2_v5；bpe_v11_v3_bytes.py调用的是py_bpe_train_step2_v5；bpe_v11_v4.py调用的是py_bpe_train_step2_v6；bpe_v11_v4_bytes.py调用的是py_bpe_train_step2_v6。

它们的代码基本相同，这里看一下bpe_v11.py：

        vocabulary = {i: bytes([i]) for i in range(N_BYTES)} # every byte
        for i, token in enumerate(special_tokens):
            vocabulary[N_BYTES + i] = token.encode('utf-8')
        size = N_BYTES + len(special_tokens)
        merges = []

        # initial word encodings are utf-8
        word_encodings = {}
        for word in word_counts:
            word_encodings[word] = list(word.encode('utf-8'))

        word_ids = {word:id for id, word in enumerate(word_counts)}

        wordid_counts = {word_ids[word]:count for word, count in word_counts.items()}

        wordid_encodings = {word_ids[word]:encoding for word, encoding in word_encodings.items()}      


        merges_cpp, vocabulary_cpp = py_bpe_train_step2_v3(vocab_size, 
                             vocabulary,
                             wordid_counts,
                             wordid_encodings,
                             merges)


        vocabulary = {k:bytes(v) for k, v in vocabulary_cpp.items()}
        merges = [(bytes(arr[0]), bytes(arr[1])) for arr in merges_cpp]

在调用之前，我们需要把str的word变成int的id，这个映射关系保存在word_ids。然后利用word_ids把word_counts变成wordid_counts，把word_encodings变成wordid_encodings。调用后得到的merges_cpp, vocabulary_cpp需要把list[int]转换成bytes。

4. 测试

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v8_v3	open_web	897/899/951	395/399/395	442/438/493	num_counter=8, num_merger=1
bpe_v9	open_web	831/816/867	400/401/390	prepare & convert: 93/86/94 py_bpe_train_step2: 330/320/374 c++:289/281/326	num_counter=8, num_merger=1
bpe_v10	open_web	816/769/788	390/400/400	prepare & convert: 21/17/19 py_bpe_train_step2: 402/350/367 c++:338/296/309	num_counter=8, num_merger=1
bpe_v10_v2	open_web	767/774/767	400/401/401	prepare & convert: 18/17/17 py_bpe_train_step2: 346/355/347 c++:292/298/294	num_counter=8, num_merger=1
bpe_v10_v2	open_web	498/477/495	120/120/120	prepare & convert: 20/19/18 py_bpe_train_step2: 355/336/354 c++:299/282/298	num_counter=32, num_merger=4
bpe_v11	open_web	350/340/354	120/120/120	prepare & convert: 21/19/19 py_bpe_train_step2: 207/199/212 c++:183/175/190	num_counter=32, num_merger=4
bpe_v11_bytes	open_web	311/307/305	80/80/80	prepare & convert: 18/19/18 py_bpe_train_step2: 211/206/204 c++:189/183/182	num_counter=64, num_merger=8, chunk_size 8mb
bpe_v11_v2	open_web	362/350/338	130/120/120	prepare & convert: 19/18/19 py_bpe_train_step2: 210/210/197 c++:189/190/176	num_counter=32, num_merger=4
bpe_v11_v3	open_web	269/274/270	120/120/120	prepare & convert: 18/19/18 py_bpe_train_step2: 129/133/129 c++: 106/109/106	num_counter=32, num_merger=4
bpe_v11_v3_bytes	open_web	218/219/215	72/74/69	prepare & convert: 21/21/21 py_bpe_train_step2: 123/122/123 c++: 101/100/101	num_counter=64, num_merger=8, chunk_size 8mb
bpe_v11_v4	open_web	258/256/261	116/117/117	prepare & convert: 19/18/19 py_bpe_train_step2: 121/119/123 c++: 98/97/100	num_counter=32, num_merger=4
bpe_v11_v4_bytes	open_web	210/206/207	71/69/70	prepare & convert: 20/18/19 py_bpe_train_step2: 117/117/117 c++: 95/96/95	num_counter=64, num_merger=8, chunk_size 8mb

对比bpe_v10和bpe_v10_v2，它们的差别一是手动在python和c++之间转换参数一是cython自动转换。自动转换不仅更方便，而且更快。

最终我们使用64核bpe_v11_v4_bytes最快的训练时间是200多秒，这比最初的十多个小时快了100多倍！

5. 总结

这就是本系列文章的全部内容。下面是简要的总结和对应文章的链接：

第0部分：简介介绍bpe训练的基本算法和相关任务，并且介绍开发环境。
第1部分：最简单实现 bpe训练最简单的实现。
第2部分：优化算法实现pair_counts的增量更新。
第3部分：并行分词和统计词频使用multiprocessing实现多进程并行算法。
第4部分：一次失败的并行优化尝试用多进程并行计算max pair。
第5部分：用C++实现Merge算法用C++实现和Python等价的merge算法，并且比较std::unordered_map的两种遍历方式。
第6部分：用OpenMP实现并行求最大用OpenMP并行求pair_counts里最大pair。
第7部分：使用flat hashmap替代std::unordered_map 使用flat hashmap来替代std::unordered_map。
第8部分：实现细粒度更新使用倒排索引实现pair_counts的细粒度更新算法。
第9部分：使用堆来寻找最大pair 使用堆来求最大pair，提升性能。
第10部分：使用cython和pypy来加速使用cython和pypy来加速python代码。
第11部分：使用cython封装c++代码使用cython封装c++代码。

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 10: Using Cython and PyPy for Acceleration

Wed, 24 Sep 2025 00:00:00 +0000

This series of articles implements a subtask of Stanford’s CS336 Assignment 1: building an efficient training algorithm for a BPE Tokenizer. Through a series of optimizations, our algorithm’s training time on OpenWebText was reduced from over 10 hours to less than 10 minutes. This series explains these optimizations, including algorithmic improvements, data structure enhancements, parallelization with OpenMP, Cython optimization, and implementing key code in C++ along with its integration via Cython. This article, the eleventh in the series, will cover using Cython and PyPy to accelerate Python code.

Table of Content

1. Problem Analysis
2. Introduction to Cython Principles
3. Optimizing the fine_grained_pair_counter_diff function with Cython
4. Test Results
5. Implementing the entire merge process with Cython
6. Using PyPy
7. Conclusion
Full Series

1. Problem Analysis

In the previous article, we have already performed extreme optimizations on the algorithm. The fastest Python version, bpe_v7_maxheapc_opt_time, has a merge time of over 500 seconds, while the fastest C++ version, bpe_train_updater_fine_grained_heap_emhash8_set9_opt, has a merge time of 100 seconds. Later, we will integrate the C++ code with the previous Python code using Cython. However, today, we want to perform some optimizations at the Python level, as using two languages is quite troublesome.

For the same logic, Python is slower than C++ for many reasons. The main reasons are: Python (more accurately, CPython) interprets code while C++ compiles to machine code; Python is dynamically typed while C++ is statically typed; Python’s GIL prevents true multithreading parallelism; and Python’s memory layout is not CPU-friendly.

The third reason was analyzed when comparing Python’s multiprocessing and C++/OpenMP parallel max search. Furthermore, through algorithmic optimization, we no longer need to scan the entire pair_counts to find the max, so we can ignore this difference.

The fourth reason is also very important. Even in C++, if you use a linked-list memory layout like std::unordered_map, its traversal speed is much slower than absl::flat_hash_map and emhash8::HashMap. In Python, everything is an object, which from a memory layout perspective is essentially a void *. Therefore, although Python’s list looks similar to std::vector, their memory layouts are completely different.

When we create a dict with integer keys and values, like {1: 10, 2: 20}, Python actually does the following in memory:

Creates a hash table structure for the dictionary itself.
Creates a PyLongObject for the key 1.
Creates a PyLongObject for the value 10.
Stores the references (memory addresses) of the key and value objects in a hash table entry.
Repeats the process for key 2 and value 20.

This means that even if your keys and values are just small integers, each integer comes with the overhead of a full Python object (including type information, reference counting, etc.). This design makes dict extremely flexible, allowing it to store data of any type, but at the cost of consuming more memory.

However, at the Python level, it’s difficult to optimize dict. This is the cost of using Python. If you want to optimize memory, the best way is to use other languages (like C/C++) and integrate them via the Python/C API. Our previous max-heap algorithm was integrated via the Python/C API, but this API is very complex. We will later use Cython to integrate the C++ code, as Cython ultimately compiles our C++ code into an extension module. Similarly, NumPy is integrated via the Python/C API.

The remaining two points are the focus of today’s optimization: first, we’ll try to rewrite some key code with Cython; second, we’ll try to use PyPy to replace CPython.

2. Introduction to Cython Principles

Cython is a programming language that combines the ease of use of Python with the high performance of C. It is designed to be a superset of Python, which means you can use it to write regular Python code while also adding C language features, such as static type declarations. The core idea of Cython is to compile your Python code into C language, and then compile it into machine code to form an importable Python module. This process solves the performance bottleneck of the Python interpreter.

For a detailed introduction to Cython, please refer to the official documentation. Additionally, if readers want to learn Cython systematically, they can read the book Cython: A Guide for Python Programmers. Although it’s an old book, most of its content is still useful.

3. Optimizing the `fine_grained_pair_counter_diff` function with Cython

Our optimization is based on bpe_v7_opt. According to our previous analysis, the time for max can be ignored, and most of the remaining time is spent in the _updated_affected_word_count function. This function mainly calls fine_grained_pair_counter_diff and then updates the pair_counts and pair_to_words dictionaries. Updating dictionaries is difficult to optimize with Cython, so we will first optimize the fine_grained_pair_counter_diff function.

We first create a bpe_updater.pyx file, which implements fine_grained_pair_counter_diff with Cython.

cpdef void fine_grained_pair_counter_diff(set affected_words, 
                                          word_encodings, 
                                          word_counts, 
                                          tuple merge_pair, 
                                          diff_pairs, 
                                          int new_id, 
                                          pair_to_words, 
                                          set new_pairs):
    cdef str word
    cdef int wc
    cdef int idx
    cdef int first_idx
    cdef int last_idx
    cdef int i
    cdef int tk_len

    for word in affected_words:
        word_tokens = word_encodings[word]
        wc = word_counts[word]

        # find first and last pairs
        idx = 0
        unaffected_pairs = set()
        tk_len = len(word_tokens)
        #first_idx = -1
        while idx < tk_len - 1:
            if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                first_idx = idx
                break
            idx += 1

        # assert first_idx exists

        idx = tk_len - 2
        while idx > first_idx + 1:
            if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                last_idx = idx
                break
            idx -= 1
        else:
            last_idx = first_idx

        start_idx = max_int(0, first_idx - 1) # inclusive
        end_idx = min_int(last_idx + 3, tk_len) # exclusive

        # unaffected [0, start_idx)


        for i in range(start_idx):
            pair = word_tokens[i], word_tokens[i + 1]
            unaffected_pairs.add(pair)
        # unaffected [end_idx-1, :-1]
        for i in range(end_idx - 1, tk_len - 1):
            pair = word_tokens[i], word_tokens[i + 1]
            unaffected_pairs.add(pair)                

        # TODO avoid slice copy
        affected_tokens = word_tokens[start_idx: end_idx]
        for i in range(len(affected_tokens) - 1):
            old_pair = (affected_tokens[i], affected_tokens[i + 1])
            diff_pairs[old_pair] -= wc 
            if old_pair not in unaffected_pairs:   
                pair_to_words[old_pair].discard(word)
        

        new_tokens = []
        all_new_tokens = []
        for i in range(start_idx):
            all_new_tokens.append(word_tokens[i])
        
        i = 0
        # account for multiple occurrences of the pair
        while i < len(affected_tokens):
            if i < len(affected_tokens) - 1 and (affected_tokens[i], affected_tokens[i + 1]) == merge_pair:
                new_tokens.append(new_id)
                all_new_tokens.append(new_id)
                # jump past pair
                i += 2
            else:
                new_tokens.append(affected_tokens[i])
                all_new_tokens.append(affected_tokens[i])
                i += 1



        for i in range(end_idx, len(word_tokens)):
            all_new_tokens.append(word_tokens[i])
        
        word_encodings[word] = all_new_tokens

        # add new pairs from the updated word
        for i in range(len(new_tokens) - 1):
            new_pair = (new_tokens[i], new_tokens[i + 1])

            diff_pairs[new_pair] += wc
            pair_to_words[new_pair].add(word)

            new_pairs.add(new_pair)

This code is essentially the same as the Python version, except that static type declarations have been added for the loop variables. This allows Cython to compile the loops into C-language versions, bypassing the Python iterator protocol.

Additionally, we can see in the code above:

affected_tokens = word_tokens[start_idx: end_idx]

This slice will copy the affected tokens. Our C++ implementation code avoids this copy by using pointers:

        const int * affected_tokens = word_tokens.data() + start_idx;
        int affected_tokens_len = end_idx - start_idx;

        for(int i = 0; i < affected_tokens_len - 1; ++i){
            std::pair<int, int> old_pair(affected_tokens[i], affected_tokens[i + 1]);
            diff_pairs[old_pair] -= wc;
            if (unaffected_pairs.find(old_pair) == unaffected_pairs.end()) {
                pair_wordids[old_pair].erase(wordid);
            }
        }

Python’s list cannot return a view, so we must manually handle the tedious indices to avoid copying. This gives us fine_grained_pair_counter_diff_v2. We won’t go into the detailed changes here, but interested readers can refer to fine_grained_pair_counter_diff_v2.

Furthermore, based on the optimization from the previous article, we can remove new_pairs, which leads to fine_grained_pair_counter_diff_v3 based on fine_grained_pair_counter_diff_v2.

Next, we need to modify setup.py:

setup(
    packages=['cs336_basics'],
    name='cs336_basics.bpe_updater',
    ext_modules=cythonize("cs336_basics/bpe_updater.pyx"),
)

Then execute python setup.py build_ext -i, which compiles to cs336_basics/bpe_updater.cpython-312-x86_64-linux-gnu.so.

Finally, we use this bpe_updater extension module in our Python code, with the corresponding code in bpe_v8.py, bpe_v8_v2.py, and bpe_v8_v3.py. They call fine_grained_pair_counter_diff, fine_grained_pair_counter_diff_v2, and fine_grained_pair_counter_diff_v3, respectively.

4. Test Results

Since gathering statistics for max and update introduces extra time overhead (e.g., bpe_v7_time is over 10 seconds slower than bpe_v7), I will only report the total merge time here.

Version	Data	Total Time (s)	Word Freq Time (s)	Merge Time (s)	Other
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7	open_web	1051/1017/1023	399/389/398	590/568/568	num_counter=8, num_merger=1
bpe_v7_opt	open_web	1021/1007/980	399/393/393	558/553/528	num_counter=8, num_merger=1
bpe_v8	open_web	934/930/935	390/393/393	479/472/476	num_counter=8, num_merger=1
bpe_v8_v2	open_web	917/908/965	394/392/395	460/455/505	num_counter=8, num_merger=1
bpe_v8_v3	open_web	897/899/951	395/399/395	442/438/493	num_counter=8, num_merger=1

The implementation logic of bpe_v8 and bpe_v7 is identical. We can see that the Cython version bpe_v8 has a merge time that is 17% faster than bpe_v7. bpe_v8_v2 avoids copying affected_tokens = word_tokens[start_idx: end_idx], and its first two runs are over 10 seconds faster than bpe_v8, but the third run is 30 seconds slower. The reason for this is unclear, perhaps due to fluctuations in server load. bpe_v7_opt made the new_pairs deletion optimization over bpe_v7, and it is compared with bpe_v8_v3, which also includes the copy-avoidance optimization from bpe_v8_v2. The average time for bpe_v8_v3 is 457 seconds, which is 16% faster than bpe_v7_opt.

5. Implementing the entire merge process with Cython

To go a step further, can implementing the entire merge process with Cython make it even faster? By refactoring bpe_v8_v3, I implemented bpe_updater_v2.pyx and bpe_v8_v4.py. Here, I encapsulated the merge process into a single function:

cpdef void bpe_train_step2(int vocab_size,
                      pair_counts,
                      pair_strings,
                      vocabulary,
                      pair_to_words,
                      word_counts,
                      word_encodings,
                      merges,
                      pair_heap):

    cdef int size = len(vocabulary)
    while size < vocab_size:
        _merge_a_pair(pair_counts, pair_strings, vocabulary,
                                pair_to_words, word_counts, word_encodings,
                                merges, size, pair_heap)
        size += 1

The test results are as follows:

Version	Data	Total Time (s)	Word Freq Time (s)	Merge Time (s)	Other
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7	open_web	1051/1017/1023	399/389/398	590/568/568	num_counter=8, num_merger=1
bpe_v7_opt	open_web	1021/1007/980	399/393/393	558/553/528	num_counter=8, num_merger=1
bpe_v8	open_web	934/930/935	390/393/393	479/472/476	num_counter=8, num_merger=1
bpe_v8_v2	open_web	917/908/965	394/392/395	460/455/505	num_counter=8, num_merger=1
bpe_v8_v3	open_web	897/899/951	395/399/395	442/438/493	num_counter=8, num_merger=1
bpe_v8_v4	open_web	917/915/982	391/404/400	462/448/506	num_counter=8, num_merger=1

bpe_v8_v4 is even slower than bpe_v8_v3. The reason is that in the fine_grained_pair_counter_diff function, we used Cython’s static type declarations to turn the Python iterator loop into a C loop, making it faster. However, the rest of the code operates on Python dictionaries. Implementing this code in Cython still reverts to Python, which is slower.

6. Using PyPy

PyPy is an alternative implementation of the Python language. Simply put, it’s not a library or a framework, but a complete Python interpreter, just like the CPython we commonly use. PyPy’s biggest highlight is its built-in JIT (Just-In-Time) compiler. This is the biggest difference from the standard CPython interpreter and the reason for its speed. CPython compiles Python code into bytecode, which is then executed one line at a time by a virtual machine. This process is relatively slow. PyPy also first compiles Python code into bytecode. However, while the program is running, PyPy’s JIT compiler monitors which code is frequently executed. It then compiles this “hotspot” code directly into machine code and caches it. The next time it encounters the same code, PyPy executes the high-speed machine code directly instead of re-interpreting the bytecode. PyPy is quite similar to Java’s JVM.

To switch the environment from CPython to PyPy, you don’t need to modify any code, but you do need to reinstall the environment. Since we are using uv to manage our environment, switching is very simple; just run the following commands:

deactivate
UV_PROJECT_ENVIRONMENT=.venv_pypy uv sync --python pypy@3.11
source .venv_pypy/bin/activate

The first command exits the current virtual environment. Then, uv sync is used to create a new PyPy 3.11 environment. To avoid overwriting the original .venv, we use the UV_PROJECT_ENVIRONMENT environment variable to tell uv to create the new virtual environment in the .venv_pypy directory. Finally, we source this environment.

I first tested bpe_v7, and the results are as follows:

Version	Data	Total Time (s)	Word Freq Time (s)	Merge Time (s)	Other
bpe_v7	open_web	1028/1053/1045	395/397/393	total: 575/589/590 max: 6/6/6 update: 569/583/583 make heap: 0.01/0.01 heap_push_time: 102/107/122	num_counter=8, num_merger=1
bpe_v7(pypy)	open_web	2047/1694/1913	1644/1306/1552	403/388/361	num_counter=8, num_merger=1

The results show that PyPy’s overall runtime is much slower than CPython, increasing from over 1000 seconds to over 2000 seconds. However, PyPy’s merge time is 34% faster than CPython. Why is there a speed discrepancy? Through a series of tests, I pinpointed the issue to the speed difference in regex matching, and then I wrote a program, bpe_v2_time_pypy.py. This program only executes the _pretokenize_and_count function and measures its time. The following are the results from testing on the tinystory dataset with a single CPU (num_counter=1) under PyPy 3.11, CPython 3.12, and CPython 3.11:

Version	Data	split_time (s)	match_time (s)
bpe_v2_pypy(pypy 3.11)	tiny_story	40	888
bpe_v2_pypy(cpython 3.12)	tiny_story	3	300
bpe_v2_pypy(cpython 3.11)	tiny_story	3	288

split_time and match_time respectively measure the time for regex.split and regex.finditer:

        for chunk in BPE_Trainer._chunk_documents_streaming(input_path):
            start_time = time.perf_counter()
            blocks = re.split(special_pattern, chunk)
            end_time = time.perf_counter()
            split_time += (end_time - start_time)
            for block in blocks:
                start_time = time.perf_counter()
                for match in re.finditer(pattern, block):
                    text = match.group(0)
                    text_len += len(text)
                end_time = time.perf_counter()
                match_time += (end_time - start_time)

As we can see, the PyPy version is much slower than CPython. To find the reason, I submitted an issue on the regex GitHub repository: regex is much slower in pypy than cpython. According to mattip’s reply, the reason for PyPy’s slowness is that regex uses CPython’s C API, and PyPy can only implement this API through emulation, which is very slow. Not only is it slow, but the results can even be incorrect. According to the regex documentation:

This module is targeted at CPython. It expects that all codepoints are the same width, so it won’t behave properly with pypy outside U+0000..U+007F because pypy stores strings as UTF-8.

It’s very unfortunate that while PyPy’s JIT compilation does speed up the subsequent merge process, the unsupported regex library makes our overall PyPy run slower and even unusable. Of course, we could split the train function into two parts: the first part, _pretokenize_and_count, uses CPython, and the second part, merge, uses PyPy, with inter-process communication between them. However, inter-process communication is complex and data copying also introduces significant overhead.

Therefore, the biggest problem with PyPy is its lack of “compatibility” with CPython, while in reality, many high-performance third-party libraries use CPython’s C API. For example, NumPy: PyPy cannot directly accelerate NumPy’s underlying C code. Instead, it communicates with NumPy’s C extension through a compatibility layer called cpyext. This compatibility layer is designed to allow PyPy to run CPython’s C extensions. This means that when only using NumPy, PyPy can be even slower than CPython.

7. Conclusion

Although PyPy can accelerate Python code execution, we have to give it up because of the lack of support from the third-party library regex. On the other hand, with Cython, we were able to achieve the faster speed without modifying the Python code much (though some code was refactored for module calls). Of course, we can further optimize the Cython code, for example, by using C++’s unordered_map instead of Python’s dict. However, a better approach is to perform these optimizations directly in C++ and then package them into a dynamic library for Cython to call. This will be the topic of our next article.

Full Series

Part 0: Introduction Introduces the basic BPE training algorithm and related tasks, as well as the development environment.
Part 1: The Simplest Implementation The simplest implementation of BPE training.
Part 2: Optimized Algorithm Implements incremental updates for pair_counts.
Part 3: Parallel Tokenization and Frequency Counting Uses multiprocessing to implement a multi-process parallel algorithm.
Part 4: A Failed Parallel Optimization An attempt to parallelize the max pair calculation using multiple processes.
Part 5: Implementing the Merge Algorithm in C++ Implements a C++ merge algorithm equivalent to the Python version, and compares two ways of iterating through std::unordered_map.
Part 6: Parallelizing the Max Pair Search with OpenMP Uses OpenMP to find the max pair in pair_counts in parallel.
Part 7: Using Flat Hashmap to Replace std::unordered_map Uses flat hashmap to replace std::unordered_map.
Part 8: Implementing Fine-Grained Updates Implements a fine-grained update algorithm for pair_counts using an inverted index.
Part 9: Using a Heap to Find the Max Pair Uses a heap to find the max pair and improve performance.
Part 10: Using Cython and PyPy for Acceleration Uses Cython and PyPy to accelerate Python code.
Part 11: Wrapping C++ Code with Cython Wraps C++ code using Cython.

动手实现和优化BPE Tokenizer的训练——第10部分：使用cython和pypy来加速

Wed, 24 Sep 2025 00:00:00 +0000

1. 问题分析
2. cython原理简介
3. 用cython优化fine_grained_pair_counter_diff函数
4. 测试结果
5. 把整个merge过程用cython实现
6. 使用pypy
7. 总结
本系列全部文章

1. 问题分析

上一篇文章我们已经对算法做了极致的优化，目前Python版本最快的是bpe_v7_maxheapc_opt_time，它的merge时间是500多秒；C++版本最快的是bpe_train_updater_fine_grained_heap_emhash8_set9_opt，它的merge时间是100秒。后面我们会把c++版本的代码通过cython和前面的python代码集成。不过今天我们想要在python的层面做一些优化，毕竟使用两种语言还是比较麻烦。

同样的代码逻辑，python比c++要慢，这有很多的原因。最主要的原因有：python(更准确的说是CPython)是通过解释执行代码而c++是编译成机器码执行；python是动态类型而c++是静态类型；python的GIL锁使得python多线程无法真正并行；python的内存布局对于cpu不友好。

第三点原因在前面比较python的多进程和c++/openmp并行求max已经分析过了，而且通过算法的优化，我们不需要扫描整个pair_counts来求max，所以这个差异我们可以忽略。

第四点其实也是很重要的，即使是c++，如果使用std::unordered_map这样的链表这种内存布局，它的遍历速度也比absl::flat_hash_map和emhash8::HashMap要慢很多。在python里，所有的一切都是对象，其实从内存布局来说就是void *，所以python的list的用法虽然看起来和std::vector很像，但是它们内存布局完全不同。

当我们创建一个key/value是整数的dict时，比如 {1: 10, 2: 20}，Python 在内存中实际做了以下几件事：

为字典本身创建一个哈希表结构。
为键 1 创建一个 PyLongObject。
为值 10 创建一个 PyLongObject。
将键对象和值对象的引用（内存地址）存储到哈希表中的一个条目里。
对键 2 和值 20 重复上述过程。

这意味着，即使你的键和值只是一些小小的整数，每个整数背后都有一个完整的 Python 对象开销（包括类型信息、引用计数等）。这种设计使得 dict 极为灵活，可以存储任何类型的数据，但代价就是会消耗更多的内存。

不过在python层面，我们很难对dict进行优化。这是使用python代价，如果想要优化内存，最好的方法是使用其它的语言(比如c/c++)然后通过Python/C API集成。我们之前的大堆算法就是通过Python/C API集成，不过这个API非常复杂。我们后面会通过cython把c++的代码集成，cython最终也是把我们的c++代码编译成扩展模块。类似的numpy就是通过Python/C集成。

剩下的第一点和第二点就是我们今天优化的主题：一是尝试通过cython来重写部分关键代码；二是尝试使用pypy来替代CPython。

2. cython原理简介

Cython是一种编程语言，它结合了Python的易用性和C语言的高性能。它旨在成为Python的超集，这意味着你可以用它来编写普通的Python代码，同时也能添加额外的C语言特性，比如静态类型声明。Cython的核心思想是将你的代码从Python语言编译成C语言，然后再编译成机器码，形成一个可导入的Python模块。这个过程解决了Python解释器在执行速度上的瓶颈。

关于cython的详细介绍请参考官方文档，另外如果读者想系统学习cython也可以阅读书籍Cython: A Guide for Python Programmers，这本书虽然很久远了，但是大部分内容依然又用。

3. 用cython优化fine_grained_pair_counter_diff函数

我们的优化基于bpe_v7_opt，根据之前的分析，max的时间已经可以忽略，剩下大部分时间在函数_updated_affected_word_count。这个函数主要是调用fine_grained_pair_counter_diff然后更新pair_counts和pair_to_words这两个词典。更新词典很难用cython优化，所以我们先优化fine_grained_pair_counter_diff函数。

我们首先创建一个bpe_updater.pyx文件，它用cython实现了fine_grained_pair_counter_diff。

cpdef void fine_grained_pair_counter_diff(set affected_words, 
                                          word_encodings, 
                                          word_counts, 
                                          tuple merge_pair, 
                                          diff_pairs, 
                                          int new_id, 
                                          pair_to_words, 
                                          set new_pairs):
    cdef str word
    cdef int wc
    cdef int idx
    cdef int first_idx
    cdef int last_idx
    cdef int i
    cdef int tk_len

    for word in affected_words:
        word_tokens = word_encodings[word]
        wc = word_counts[word]

        # find first and last pairs
        idx = 0
        unaffected_pairs = set()
        tk_len = len(word_tokens)
        #first_idx = -1
        while idx < tk_len - 1:
            if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                first_idx = idx
                break
            idx += 1

        # assert first_idx exists

        idx = tk_len - 2
        while idx > first_idx + 1:
            if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                last_idx = idx
                break
            idx -= 1
        else:
            last_idx = first_idx

        start_idx = max_int(0, first_idx - 1) # inclusive
        end_idx = min_int(last_idx + 3, tk_len) # exclusive

        # unaffected [0, start_idx)


        for i in range(start_idx):
            pair = word_tokens[i], word_tokens[i + 1]
            unaffected_pairs.add(pair)
        # unaffected [end_idx-1, :-1]
        for i in range(end_idx - 1, tk_len - 1):
            pair = word_tokens[i], word_tokens[i + 1]
            unaffected_pairs.add(pair)                

        # TODO avoid slice copy
        affected_tokens = word_tokens[start_idx: end_idx]
        for i in range(len(affected_tokens) - 1):
            old_pair = (affected_tokens[i], affected_tokens[i + 1])
            diff_pairs[old_pair] -= wc 
            if old_pair not in unaffected_pairs:   
                pair_to_words[old_pair].discard(word)
        

        new_tokens = []
        all_new_tokens = []
        for i in range(start_idx):
            all_new_tokens.append(word_tokens[i])
        
        i = 0
        # account for multiple occurrences of the pair
        while i < len(affected_tokens):
            if i < len(affected_tokens) - 1 and (affected_tokens[i], affected_tokens[i + 1]) == merge_pair:
                new_tokens.append(new_id)
                all_new_tokens.append(new_id)
                # jump past pair
                i += 2
            else:
                new_tokens.append(affected_tokens[i])
                all_new_tokens.append(affected_tokens[i])
                i += 1



        for i in range(end_idx, len(word_tokens)):
            all_new_tokens.append(word_tokens[i])
        
        word_encodings[word] = all_new_tokens

        # add new pairs from the updated word
        for i in range(len(new_tokens) - 1):
            new_pair = (new_tokens[i], new_tokens[i + 1])

            diff_pairs[new_pair] += wc
            pair_to_words[new_pair].add(word)

            new_pairs.add(new_pair)

这个代码和python的版本基本一样，只不过把循环变量都加了静态类型声明，这样cython就可以把它编译成c语言版本的for/while循环，而不需要走python的迭代器协议。

另外我们可以发现上面的代码：

affected_tokens = word_tokens[start_idx: end_idx]

这个slice会复制受影响的tokens，我们在c++实现的代码使用了指针避免复制：

        const int * affected_tokens = word_tokens.data() + start_idx;
        int affected_tokens_len = end_idx - start_idx;

        for(int i = 0; i < affected_tokens_len - 1; ++i){
            std::pair<int, int> old_pair(affected_tokens[i], affected_tokens[i + 1]);
            diff_pairs[old_pair] -= wc;
            if (unaffected_pairs.find(old_pair) == unaffected_pairs.end()) {
                pair_wordids[old_pair].erase(wordid);
            }
        }

python的list没有办法返回一个视图，所以我们只能自己处理繁琐的下标来避免复制，这样我们得到fine_grained_pair_counter_diff_v2。这里就不详细介绍其中的改动了，感兴趣的读者可以阅读fine_grained_pair_counter_diff_v2。

此外根据上一篇文章的优化，我们可以把new_pairs去掉，这样在fine_grained_pair_counter_diff_v2的基础上得到fine_grained_pair_counter_diff_v3。

然后我们需要修改setup.py：

setup(
    packages=['cs336_basics'],
    name='cs336_basics.bpe_updater',
    ext_modules=cythonize("cs336_basics/bpe_updater.pyx"),
)

然后执行python setup.py build_ext -i，这样编译得到cs336_basics/bpe_updater.cpython-312-x86_64-linux-gnu.so。

最后我们在python里使用bpe_updater这个扩展模块，相应的代码在bpe_v8.py、bpe_v8_v2.py和bpe_v8_v3.py。它们分别调用fine_grained_pair_counter_diff、fine_grained_pair_counter_diff_v2和fine_grained_pair_counter_diff_v3。

4. 测试结果

因为统计max和update会带来额外的时间开销，比如bpe_v7_time要比bpe_v7慢10多秒，所以我这里只统计整体的合并时间。

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7	open_web	1051/1017/1023	399/389/398	590/568/568	num_counter=8, num_merger=1
bpe_v7_opt	open_web	1021/1007/980	399/393/393	558/553/528	num_counter=8, num_merger=1
bpe_v8	open_web	934/930/935	390/393/393	479/472/476	num_counter=8, num_merger=1
bpe_v8_v2	open_web	917/908/965	394/392/395	460/455/505	num_counter=8, num_merger=1
bpe_v8_v3	open_web	897/899/951	395/399/395	442/438/493	num_counter=8, num_merger=1

bpe_v8和bpe_v7的实现逻辑完全相同，我们可以对比得出cython版本的bpe_v8要比bpe_v7的合并时间少17%。bpe_v8_v2避免了复制affected_tokens = word_tokens[start_idx: end_idx]，它的前两次运行要比bpe_v8快10多秒，但是第三次运行反而慢了30秒。这个原因不太清楚，也许是服务器的负载波动。bpe_v7_opt对bpe_v7做了new_pairs的删除优化，和它对比的是bpe_v8_v3，但是bpe_v8_v3还多做了bpe_v8_v2的避免复制的优化。bpe_v8_v3的平均时间是457秒，比bpe_v7_opt快了16%。

5. 把整个merge过程用cython实现

更进一步，我们可以把整个merge过程都用cython实现，这能不能加快速度呢？对bpe_v8_v3进行重构我实现了bpe_updater_v2.pyx和bpe_v8_v4.py。这里，我把merge过程封装成一个函数：

cpdef void bpe_train_step2(int vocab_size,
                      pair_counts,
                      pair_strings,
                      vocabulary,
                      pair_to_words,
                      word_counts,
                      word_encodings,
                      merges,
                      pair_heap):

    cdef int size = len(vocabulary)
    while size < vocab_size:
        _merge_a_pair(pair_counts, pair_strings, vocabulary,
                                pair_to_words, word_counts, word_encodings,
                                merges, size, pair_heap)
        size += 1

测试结果如下：

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7	open_web	1051/1017/1023	399/389/398	590/568/568	num_counter=8, num_merger=1
bpe_v7_opt	open_web	1021/1007/980	399/393/393	558/553/528	num_counter=8, num_merger=1
bpe_v8	open_web	934/930/935	390/393/393	479/472/476	num_counter=8, num_merger=1
bpe_v8_v2	open_web	917/908/965	394/392/395	460/455/505	num_counter=8, num_merger=1
bpe_v8_v3	open_web	897/899/951	395/399/395	442/438/493	num_counter=8, num_merger=1
bpe_v8_v4	open_web	917/915/982	391/404/400	462/448/506	num_counter=8, num_merger=1

bpe_v8_v4比bpe_v8_v3还慢。原因在于fine_grained_pair_counter_diff函数我们通过cython的静态类型声明，把python的迭代器循环变成了c的循环，所以速度能够变快。但是其它部分的代码都是操作python的dict，这些代码用cython实现还是会回到python，这反而更慢。

6. 使用pypy

pypy 是 Python 语言的一个替代性实现。简单来说，它不是一个库或框架，而是一个完整的 Python 解释器，就像我们平时用的 CPython 一样。pypy 的最大亮点在于它内置了 JIT (Just-In-Time) 编译器。这是它和标准 CPython 解释器最大的区别，也是它速度快的原因。CPython将 Python 代码编译成字节码，然后由一个虚拟机逐条解释执行。这个过程比较慢。pypy同样先将 Python 代码编译成字节码。但是，当程序运行时，pypy 的 JIT 编译器会监控哪些代码被频繁执行。它会把这些“热点”代码直接编译成机器码，然后缓存起来。下次再遇到同样的代码时，pypy 就会直接执行高速的机器码，而不是重新解释字节码。pypy和Java的JVM是比较类似的。

要把运行环境从CPython切换到pypy不需要修改任何代码，但是需要重新安装环境。由于我们是使用uv来管理环境的，切换环境非常简单，执行如下命令即可：

deactivate
UV_PROJECT_ENVIRONMENT=.venv_pypy uv sync --python pypy@3.11
source .venv_pypy/bin/activate

第一个是先从当前的虚拟环境中退出。然后用uv sync重新创建一个pypy 3.11的环境，为了不覆盖原来的.venv，我们使用环境变量UV_PROJECT_ENVIRONMENT，告诉uv新创建的虚拟环境存放在.venv_pypy目录下。最后source这个环境。

我首先测试了bpe_v7，测试结果如下：

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v7	open_web	1028/1053/1045	395/397/393	total: 575/589/590 max: 6/6/6 update: 569/583/583 make heap: 0.01/0.01 heap_push_time: 102/107/122	num_counter=8, num_merger=1
bpe_v7(pypy)	open_web	2047/1694/1913	1644/1306/1552	403/388/361	num_counter=8, num_merger=1

结果发现pypy的整体运行时间远远慢于CPython，从1000多秒增加到了2000多秒。不过pypy的merge时间却比CPython快了34%。为什么会出现前后速度不一致的情况呢？通过一系列测试我逐渐把问题锁定到了regex匹配的速度差异上，于是写了一个程序bpe_v2_time_pypy.py。这个程序只执行_pretokenize_and_count这个函数并且统计时间。下面是我在pypy3.11、CPython3.12和CPython3.11之下用单个CPU(num_counter=1)在tinystory数据集上测试的结果：

版本	数据	split_time(s)	match_time(s)
bpe_v2_pypy(pypy 3.11)	tiny_story	40	888
bpe_v2_pypy(cpython 3.12) tiny_story	3	300
bpe_v2_pypy(cpython 3.11) tiny_story	3	288

split_time和match_time分别统计regex.split和regex.finditer的时间：

        for chunk in BPE_Trainer._chunk_documents_streaming(input_path):
            start_time = time.perf_counter()
            blocks = re.split(special_pattern, chunk)
            end_time = time.perf_counter()
            split_time += (end_time - start_time)
            for block in blocks:
                start_time = time.perf_counter()
                for match in re.finditer(pattern, block):
                    text = match.group(0)
                    text_len += len(text)
                end_time = time.perf_counter()
                match_time += (end_time - start_time)

可以看到，pypy的版本要比CPython的慢很多。为了找到原因，我在第三方库regex的github上提交了一个issueregex is much slower in pypy than cpython。根据mattip的回复，pypy慢的原因在于regex使用了CPython的C接口，而pypy只能通过模拟来实现类似CPython的这个接口，所以速度很慢。而且不只是慢，它的结果甚至都可能不正确。根据regex的文档：

This module is targeted at CPython. It expects that all codepoints are the same width, so it won’t behave properly with pypy outside U+0000..U+007F because pypy stores strings as UTF-8.

非常遗憾，pypy的JIT编译确实使得后面的merge过程变快了，但是由于第三方库regex不支持，我们用pypy整体更慢甚至无法实现。我们当然可以把train函数分成两部分，第一部分是_pretokenize_and_count，第二部分是merge。第一个部分用CPython，第二个部分用pypy，然后两个部分之间通过进程间通信来完成。但是进程间通信比较复杂，而且数据复制的开销也很大。

所以pypy的最大问题其实是无法和CPython”兼容”，但是现实问题是很多高性能第三方库都会用到CPython的C接口。比如numpy，pypy无法直接加速 numpy 底层的 C 代码。相反，它通过一个名为 cpyext 的兼容层来与 numpy 的 C 扩展进行通信。这个兼容层是为了让 pypy 能够运行 CPython 的 C 扩展而设计的。这样单纯使用numpy的话pypy比CPython还要慢。

7. 总结

通过pypy虽然能够加快python代码的执行，但是由于第三方库regex不支持，我们只能放弃。而通过cython，我们没有怎么修改python代码(不过为了模块调用进行了一定代码的重构)的情况下提升了速度。当然我们还可以进一步优化cython代码，比如使用c++的unordered_map替代python的dict。不过更好的方法是直接在c++语言里完成这些优化，然后把它封装成动态库供cython调用。这就是下一篇文章我们要研究的内容。

本系列全部文章

第0部分：简介介绍bpe训练的基本算法和相关任务，并且介绍开发环境。
第1部分：最简单实现 bpe训练最简单的实现。
第2部分：优化算法实现pair_counts的增量更新。
第3部分：并行分词和统计词频使用multiprocessing实现多进程并行算法。
第4部分：一次失败的并行优化尝试用多进程并行计算max pair。
第5部分：用C++实现Merge算法用C++实现和Python等价的merge算法，并且比较std::unordered_map的两种遍历方式。
第6部分：用OpenMP实现并行求最大用OpenMP并行求pair_counts里最大pair。
第7部分：使用flat hashmap替代std::unordered_map 使用flat hashmap来替代std::unordered_map。
第8部分：实现细粒度更新使用倒排索引实现pair_counts的细粒度更新算法。
第9部分：使用堆来寻找最大pair 使用堆来求最大pair，提升性能。
第10部分：使用cython和pypy来加速使用cython和pypy来加速python代码。
第11部分：使用cython封装c++代码使用cython封装c++代码。

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 9: Using a Heap to Find the Maximum Pair

Sun, 21 Sep 2025 00:00:00 +0000

This series of articles implements a subtask of Stanford’s CS336 Assignment 1: building an efficient training algorithm for a BPE Tokenizer. Through a series of optimizations, our algorithm’s training time on OpenWebText was reduced from over 10 hours to less than 10 minutes. This series explains these optimizations, including algorithmic improvements, data structure enhancements, parallelization with OpenMP, Cython optimization, and implementing key code in C++ along with its integration via Cython. This is the tenth article, where we use the heap data structure to replace the process of finding the maximum pair, thereby improving performance.

Table of Content

1. Problem Analysis
2. Heap
3. Using a Heap to Find the Max
4. Python’s heapq Module
5. Implementing a Max Heap
6. Optimizing maxheap_py with a C Module
7. Using heapq to Implement a Max Heap
8. Performance Testing of Different Max Heap Implementations
9. Using a Max Heap to Find the Max
10. Testing
11. Experiments with a Larger vocab_size
12. Porting the Max Heap to C++
13. Optimizing _updated_affected_word_count
14. Porting the Optimization to C++
Full Series

1. Problem Analysis

In our bpe_v5_time from above, the update time was just over 500 seconds, but the time for finding the max was over 30,000 seconds. How can we optimize this time? We previously tried a parallel algorithm for finding the max, but the Python version was not successful because the Python GIL prevents multiple threads from using multiple CPUs simultaneously, and multiprocessing involves a lot of inter-process communication. We then tried a C++ implementation with OpenMP for parallelism, which reduced the max time to 400 seconds with 32 threads. We then replaced std::unordered_map with a faster flat hashmap, which brought the max time down to under 100 seconds without parallelism.

However, let’s go back to Python. Besides optimizing data structures, is there a way to optimize the algorithm itself? Optimizing data structures means doing the same amount of work faster, while optimizing an algorithm means achieving the same goal with less work. Of course, the two are not completely separate; sometimes, a different algorithm requires designing an appropriate data structure.

To find the maximum value in a set, we must iterate through the entire set, and this time cannot be reduced. This means the time for the first traversal cannot be reduced. However, the second traversal can be optimized because only a portion of the pair counts change between traversals. At the beginning, the affected words are numerous due to high-frequency terms, so relatively more pairs are affected. Later on, as word frequencies decrease, fewer words are affected. The simplest idea is to sort the pair_counts. When some pair counts change, we would only need to re-sort those specific pairs. Since there are only two possibilities for changes in word frequency—counts of old pairs decrease, and counts of new pairs increase—we could also use some heuristic rules. For example, if a pair’s current rank is greater than 32,000 (our target vocabulary size) and its frequency has decreased, we can ignore it.

But there is a better way to find the maximum value in a set: using a heap data structure.

2. Heap

In computer science, “heap” has two completely different meanings: one refers to a data structure, and the other refers to a type of memory allocation. Interestingly, the counterpart to the heap as a memory allocation type is the stack, which also represents both a memory allocation method and a data structure. The stack as a data structure (LIFO) is closely related to the stack as memory allocation, as it leverages the LIFO principle for function calls and returns. However, the heap as a data structure and the heap as memory allocation have no connection; their only link is that someone, for some reason, used the same word for two entirely different concepts.

Here, we’re focusing on the heap as a data structure. It’s typically used to implement priority queues and can also be used for heapsort. I won’t go into a detailed introduction of heaps; readers who are unfamiliar can find plenty of resources in any data structures and algorithms book or online, for example, on Wikipedia.

3. Using a Heap to Find the Max

Typically, when using a heap, we only use three main operations: converting an array into a heap (heapify), popping the top element from the heap (heappop), and pushing an element into the heap (heappush). We first need to use heapify to convert an array into a heap (satisfying the heap property: the root is larger/smaller than every element in its subtrees). Then, we can repeatedly call heappop and heappush, and the array will remain a heap after these operations.

There’s a problem here: after finding the maximum (with heappop), we merge a pair. This decreases the count of some old pairs and introduces some new ones. Adding new pairs is not an issue; we just need to call heappush. But how do we modify the counts of old pairs and keep the structure a valid heap? Let’s look at an example:

         11
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   / \
6  4  5   7

If we change 8 to 10, meaning an element has increased, we need to call a “sift up” operation (siftup) starting from 8:

         11
       /    \
      /      \
     10       4
   /   \     / \
  7     9   3   1 
 / \   / \
6  4  5   7

Conversely, if we change 8 to 6, we need to call a “sift down” operation (siftdown) starting from 8:

         11
       /    \
      /      \
     10       4
   /   \     / \
  7     7   3   1 
 / \   / \
6  4  5   6

But there’s a problem: how do we find the element to be modified? Recall that our main data is pair_counts, which is a dict. However, a heap operates on a list. So we need to copy the elements from pair_counts into a pair_heap list. But a list doesn’t support fast lookups; if we had to sequentially scan the list to find an element for modification, it would be a counterproductive effort (we’ve already found the max).

One solution is to store the index of the element in the list within the pair_counts dictionary, so pair_counts would become:

pair -> (count, index_in_pair_heap)

Then, when we add a new pair to pair_counts, we also add it to the correct position in pair_heap using heappush. This would require the heappush function to not only add an element but also return its index in pair_heap. This way, we can save the index in pair_counts. Later, when the count of a pair changes (in our case, it will only decrease), we can find its position in pair_heap via pair_counts and then call siftdown on that element.

This would require modifying heappush and maintaining the relationship between pair_counts and pair_heap, which makes the code quite complex. Readers who are interested can try to implement this algorithm.

However, I’m using a different approach—lazy modification. With this method, we do nothing when a pair’s count changes. We only check if its count has changed when we call heappop, which we can discover by comparing the heap element’s count with the value in pair_counts. If it has changed (it will only get smaller, which is a very important assumption), we re-insert the new count using heappush. We then continue to heappop until we find an element whose count has not been modified. That element is the current maximum.

This sounds complicated, so let’s walk through an example. Suppose our current heap is:

         11
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   / \
6  4  5   7

The maximum should be 11, but let’s assume that due to a merge, its count has changed to 10. When we pop the pair with 11, we query pair_counts and find its latest count is 10. Since it’s smaller, we can’t be sure it’s the max. So, we first pop 11, which gives us:

          9
       /    \
      /      \
     8        4
   /   \     / \
  7     7   3   1 
 / \   /
6  4  5 

Then we need to re-push 10, resulting in:

         10
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   /  \
6  4  5    7

Next, we pop the current maximum, 10. This time, a query to pair_counts reveals its count is up-to-date, so we’ve found the current maximum pair is 10.

Important Note: The crucial assumption that allows for this lazy update is that the count of an old pair will only decrease. If this assumption were not true, for example, if our heap were:

         11
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   / \
6  4  5   7

And we changed 1 to 12, we would have to siftup 1 immediately, otherwise the max we find would be 11, which is incorrect.

Using this algorithm, we don’t need to maintain pair_heap indices in pair_counts, and we don’t need to call siftdown, which is a private function in Python’s heapq (_siftdown), making its use risky as it might not be available in future versions.

4. Python’s `heapq` Module

Python’s standard library provides the heapq module. The main functions we need are heappush, heappop, and heapify. Readers who are unfamiliar can refer to The Python heapq Module: Using Heaps and Heappushs.

However, there is an issue here: we need a max heap, but Python’s heapq module provides a min heap. Later, when we analyze its code, we’ll see that it has already implemented a max heap internally. But for now, we need to discuss how we can use a min heap to achieve the functionality of a max heap.

A common trick is to reverse the elements. For example, if the heap elements are positive integers, we can store their corresponding negative integers to simulate a max heap. Here’s an example:

import heapq
arr = [3, 5, 1, 2, 6, 8, 7]

arr2 = [-i for i in arr]

arr2
[-3, -5, -1, -2, -6, -8, -7]

heapq.heapify(arr2)
arr2
[-8, -6, -7, -2, -5, -1, -3]

arr3 = [-i for i in arr2]
arr3
[8, 6, 7, 2, 5, 1, 3]

But now we need to put tuples (count, pair_string[pair], pair) into the heap. If count is large, the tuple is large; if count is the same, we compare the pair’s string and choose the larger one. The pair is included as the last element of the tuple for convenience.

Following the above method, we can put -count into the heap. But what about pair_strings[pair]? pair_strings[pair] is a tuple of bytes. If the bytes are of fixed length, since a byte’s range is 0-255, we can reverse it by subtracting each byte from 255. For example:

b1 = b'us'

b2 = b'ua'

b1 > b2
True

c1 = bytes([255 - b for b in b1])
c1
b'\x8a\x8c'

c2 = bytes([255 - b for b in b2])

c1 < c2
True

But if the strings are of variable length, this will cause problems. For example:

b1 = b'us'

b2 = b'usb'

b1 < b2
True

c1 = bytes([255 - b for b in b1])

c2 = bytes([255 - b for b in b2])

c1 > c2
False

b1 has one less character than b2. After reversing by subtracting from 255, the leading b'us' parts are the same, but regardless of reversal, the shorter string is always considered smaller than the longer one.

Therefore, we need a max heap.

5. Implementing a Max Heap

We can slightly modify Python’s built-in heapq module to turn it into a max heap. We’ll copy the source code of heapq and modify it into maxheap_py.py. The heapq module has many functions, but we only need to keep heappush, heappop, and heapify. These three functions, in turn, depend on _siftdown (which I’ve renamed _siftdown_max) and _siftup (renamed _siftup_max).

I won’t go into the full code, but here’s a comparison of _siftdown_max and _siftdown:

def _siftdown_max(heap, startpos, pos):
    'Maxheap variant of _siftdown'
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if parent < newitem:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem

def _siftdown(heap, startpos, pos):
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if newitem < parent:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem

The only difference is the single line: "if parent < newitem" versus "if newitem < parent". This assumes that the newitem and parent values have an overloaded or implemented __lt__ operator, so even when finding the larger element, it’s done by swapping the order and using the __lt__ function.

6. Optimizing `maxheap_py` with a C Module

Readers who are not interested in C module development can skip this section.

If we compare the speed of our version to the CPython version, ours is much slower because the CPython version internally calls a corresponding C module. If we look closely at the heapq source code, we’ll find:

# If available, use C implementation
try:
    from _heapq import *
except ImportError:
    pass
try:
    from _heapq import _heapreplace_max
except ImportError:
    pass
try:
    from _heapq import _heapify_max
except ImportError:
    pass
try:
    from _heapq import _heappop_max
except ImportError:
    pass

This means it will (attempt to) call the corresponding C module implementation. The specific code is in _heapqmodule.c.

I’ve copied this to implement the max heap functionality. The full code is in _maxheapqmodule.c.

Let’s look at just one function, siftdown:

static int
siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
{
    PyObject *newitem, *parent, **arr;
    Py_ssize_t parentpos, size;
    int cmp;

    assert(PyList_Check(heap));
    size = PyList_GET_SIZE(heap);
    if (pos >= size) {
        PyErr_SetString(PyExc_IndexError, "index out of range");
        return -1;
    }

    /* Follow the path to the root, moving parents down until finding
       a place newitem fits. */
    arr = _PyList_ITEMS(heap);
    newitem = arr[pos];
    while (pos > startpos) {
        parentpos = (pos - 1) >> 1;
        parent = arr[parentpos];
        Py_INCREF(newitem);
        Py_INCREF(parent);
        cmp = PyObject_RichCompareBool(newitem, parent, Py_GT);
        Py_DECREF(parent);
        Py_DECREF(newitem);
        if (cmp < 0)
            return -1;
        if (size != PyList_GET_SIZE(heap)) {
            PyErr_SetString(PyExc_RuntimeError,
                            "list changed size during iteration");
            return -1;
        }
        if (cmp == 0)
            break;
        arr = _PyList_ITEMS(heap);
        parent = arr[parentpos];
        newitem = arr[pos];
        arr[parentpos] = newitem;
        arr[pos] = parent;
        pos = parentpos;
    }
    return 0;
}

The two comparison codes have only a one-line difference:

cmp = PyObject_RichCompareBool(newitem, parent, Py_GT);
cmp = PyObject_RichCompareBool(newitem, parent, Py_LT);

As mentioned before, PyObject_RichCompareBool compares the first two arguments based on the third. A return value less than 0 indicates an error; a return value greater than 0 means the third argument’s comparison (e.g., greater than/less than) is true; and a return value of 0 means it’s false.

Therefore, to change a min heap to a max heap, we just need to change the comparison operator from Py_LT (less than) to Py_GT (greater than).

Of course, there are other minor details to modify. For example, I’ve changed the module name to maxheapqc, so the corresponding code needs to be updated:

static struct PyModuleDef _heapqmodule = {
    PyModuleDef_HEAD_INIT,
    "maxheapqc",
    module_doc,
    0,
    heapq_methods,
    heapq_slots,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit_maxheapqc(void)
{
    return PyModuleDef_Init(&_heapqmodule);
}

To compile this C module, we need to modify setup.py:

include_path = sysconfig.get_path('include')
internal_include_path = os.path.join(include_path, 'internal')

print(f"{internal_include_path=}")
project_root = os.path.dirname(os.path.abspath(__file__))

maxheapqc_module = Extension('cs336_basics.maxheapqc', sources=['cs336_basics/_maxheapqmodule.c'],
        extra_compile_args=[
            f'-I{internal_include_path}',
        ]
        )

setup(
    name='maxheapqc',
    version='1.0',
    description='maxheapqc',
    packages=['cs336_basics'],
    ext_modules=[maxheapqc_module]
)

When we compile _maxheapqmodule.c, it depends on:

#include "Python.h"
#include "pycore_list.h" 

For Python.h, setuptools will include it for us. Its location is typically something like /usr/local/include/python3.12/Python.h. But our _maxheapqmodule.c also needs to operate on lists via the C API, which requires pycore_list.h. This file is located at a path like /usr/local/include/python3.12/internal/pycore_list.h. We don’t want to hard-code this path, and the same system might have multiple Python versions and environments like Conda. So we can use sysconfig.get_path('include') to get the current Python interpreter’s path and then find internal within it.

We need to add this path during compilation using the -I flag:

maxheapqc_module = Extension('cs336_basics.maxheapqc', sources=['cs336_basics/_maxheapqmodule.c'],
        extra_compile_args=[
            f'-I{internal_include_path}',
        ]
        )

This -I parameter seems to be the way gcc adds header files. I’m not sure if it’s universal for non-Linux systems or other compilers. If you are using a different compiler, please refer to its manual for adding appropriate header file paths.

Run the following command to compile the C module:

python setup.py build_ext -i

After compilation, you’ll get a file like cs336_basics/maxheapqc.cpython-312-x86_64-linux-gnu.so. We can then write a maxheapq.py to call it.

The code for maxheapq.py is similar to heapq.py; it implements max heap functionality in Python and attempts to load maxheapqc:

# If available, use C implementation
try:
    from cs336_basics.maxheapqc import *
    print("load c!")
except ImportError:
    pass
try:
    from cs336_basics.maxheapqc import _heapreplace_max
except ImportError:
    pass
try:
    from cs336_basics.maxheapqc import _heapify_max
except ImportError:
    pass
try:
    from cs336_basics.maxheapqc import _heappop_max
except ImportError:
    pass

7. Using `heapq` to Implement a Max Heap

In fact, if we read the heapq source code, we’ll find that it already implements _heapify_max, _siftdown_max, and _siftup_max. Among these, _heapify_max has a corresponding C module implementation and is very fast. However, _siftdown_max and _siftup_max are still implemented in Python. We can look here to find which functions have C implementations:

static PyMethodDef heapq_methods[] = {
    _HEAPQ_HEAPPUSH_METHODDEF
    _HEAPQ_HEAPPUSHPOP_METHODDEF
    _HEAPQ_HEAPPOP_METHODDEF
    _HEAPQ_HEAPREPLACE_METHODDEF
    _HEAPQ_HEAPIFY_METHODDEF
    _HEAPQ__HEAPPOP_MAX_METHODDEF
    _HEAPQ__HEAPIFY_MAX_METHODDEF
    _HEAPQ__HEAPREPLACE_MAX_METHODDEF
    {NULL, NULL}           /* sentinel */
};

We can use these three functions from heapq to implement a max heap. The code is in maxheap_heapq.py. Let’s take a look:

def heappush(heap, item):
    heap.append(item)
    heapq._siftdown_max(heap, 0, len(heap)-1)


def heappop(heap):
    """Maxheap version of a heappop."""
    lastelt = heap.pop()  # raises appropriate IndexError if heap is empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        heapq._siftup_max(heap, 0)
        return returnitem
    return lastelt

def heapify(heap):
    heapq._heapify_max(heap)

Here, heapify calls the C interface and is fast, but heappush and heappop are still Python implementations.

8. Performance Testing of Different Max Heap Implementations

I’ve written a simple test script, test_heap_speed.py. The results are as follows:

maxheapq_py: 40.50480842590332
maxheap_heapq: 4.577162265777588
maxheapq: 4.334884405136108

The majority of the time here is spent on heapify. maxheapq_py is implemented in Python, so it’s much slower.

And test_heap_speed2.py primarily tests heappush and heappop. The results are as follows:

maxheapq_py: 1.892250157892704
maxheap_heapq: 1.8967562650796026
maxheapq: 0.3524886401137337

As you can see, maxheapq is much faster than the Python implementations.

9. Using a Max Heap to Find the Max

The complete code is in bpe_v6.py. Let’s just look at the differences between it and bpe_v5.py.

from cs336_basics import maxheap_py as maxheap

    def train(self, input_path, vocab_size, special_tokens, *args):
        ...
        pair_heap = []
        for pair, count in pair_counts.items():
            maxheap.heappush(pair_heap, (count, pair_strings[pair], pair))
            
            
            

First, we import maxheap_py, and for easy switching between different max heap implementations, I’ve aliased it as maxheap. Since all different max heaps have the same interface, switching to maxheapq or maxheap_heapq only requires changing this one line of code.

After the initial pair_counts are calculated, we need to construct a pair_heap. This is done by continuously calling maxheap.heappush. Another implementation method is to use heapify:

        pair_heap = []
        for pair, count in pair_counts.items():
            pair_heap.append((count, pair_strings[pair], pair))
        maxheap.heapify(pair_heap)

Both methods will eventually build a heap, but their results may not be exactly the same. Theoretically, the heapify method is faster, but actual testing shows the time difference is minimal, as pair_counts starts with fewer than 20,000 entries.

Next is the main modification: finding the current maximum pair using heappop:

    @staticmethod
    def _merge_a_pair(pair_counts, pair_strings, vocabulary, pair_to_words, 
                   word_counts, word_encodings, merges, size, pair_heap):
        
        while pair_heap:
            count, string_priority, merge_pair = maxheap.heappop(pair_heap)
            
            # check pair validity
            if merge_pair in pair_counts and pair_counts[merge_pair] == count:
                break
            elif merge_pair in pair_counts:
                # update count (lazily)
                maxheap.heappush(pair_heap, (pair_counts[merge_pair], 
                                               string_priority, 
                                               merge_pair))
        else:
            # no valid pairs found
            return False

The algorithm’s implementation is as described earlier: first, pop the current maximum merge_pair from the top of the heap. Then, check if its count has been updated by comparing it to pair_counts[merge_pair]. If the count hasn’t been updated, then merge_pair is the current maximum, and we break the loop. If it has been updated and the new count is greater than 0 (merge_pair in pair_counts), we re-insert the new count into the heap using heappush. We then continue the loop to find the maximum pair.

The final change is that whenever a new pair is created during a merge, we also need to add it to the heap:

    @staticmethod
    def _updated_affected_word_count(merge_pair, affected_words, word_encodings, 
                                     word_counts, pair_counts, pair_to_words, 
                                     new_id, pair_strings, vocabulary, pair_heap):



        for new_pair in new_pairs:
            if new_pair not in pair_strings:
                pair_strings[new_pair] = (vocabulary[new_pair[0]], vocabulary[new_pair[1]])

            maxheap.heappush(pair_heap, (pair_counts[new_pair], pair_strings[new_pair], new_pair))

10. Testing

To test the time, I also implemented bpe_v6_time.py. The test results are as follows:

Version	Data	Total Time (s)	Word Count Time (s)	Merge Time (s)	Other
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1
bpe_v6_time	open_web	1036/1107/1046	395/395/398	total: 576/641/591 max:6/7/6 update: 570/633/584	num_counter=8, num_merger=1

After using a heap to find the max, the merge time decreased from over 30,000 seconds to just over 570 seconds. The heappush time is over 100 seconds. Can we optimize it with a faster max heap?

By changing maxheap_py to maxheap_heapq or maxheapq, we get bpe_v7.py and bpe_v7_maxheapc.py. The test results are as follows:

Version	Data	Total Time (s)	Word Count Time (s)	Merge Time (s)	Other
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1
bpe_v6_time	open_web	1036/1107/1046	395/395/398	total: 576/641/591 max:6/7/6 update: 570/633/584	num_counter=8, num_merger=1
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7_maxheapc_time	open_web	1069/1017/1011	400/401/399	total: 606/556/555 max: 3/3/3 update: 602/552/552	num_counter=8, num_merger=1

Using a faster heap implementation did not speed things up. I speculate this is because the time for heappush accounts for a relatively small proportion of the total update time, so a minor change in its time doesn’t have a large impact on the overall performance.

11. Experiments with a Larger `vocab_size`

If we compare bpe_v7 with the C++ version bpe_train_updater_fine_grained_emhash8:

Version	Data	Total Time (sec)	Word Count Time (sec)	Merge Time (sec)	Other
bpe_v7_time	open_web	1028/1053/1045	395/397/393	total: 575/589/590 max: 6/6/6 update: 569/583/583 make heap: 0.01/0.01 heap_push_time: 102/107/122	num_counter=8, num_merger=1

Program	Hash Function	Total Time (sec)	Update Time (sec)	Max Time (sec)	Other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60

The total merge time for bpe_v7 is around 600 seconds, while bpe_train_updater_fine_grained_emhash8 is around 260 seconds. But what if we increase the number of merges? Since the vocabularies of modern large models are getting larger, let’s test the results for vocab_size of 64,000 and 96,000:

Version	Data	Total Time (s)	Word Count Time (s)	Merge Time (s)	Other
bpe_v7_time	open_web	1028/1053/1045	395/397/393	total: 575/589/590 max: 6/6/6 update: 569/583/583 make heap: 0.01/0.01 heap_push_time: 102/107/122	num_counter=8, num_merger=1
bpe_v7_time2	open_web	1111/1174/1100	393/420/403	total: 655/686/639 max: 9/10/10 update: 645/675/628 make heap: 0.01/0.01/0.01 heap_push_time: 128/157/108	vocab_size=64000
bpe_v7_time2	open_web	1129/1130/1123	394/406/393	total: 675/666/670 max: 13/12/12 update: 661/653/657 make heap: 0.01/0.01/0.01 heap_push_time: 143/152/120	vocab_size=96000

As you can see, because the later parts deal with low-frequency words and pairs, the time barely changes. Now let’s look at bpe_train_updater_fine_grained_emhash8:

Program	Hash Function	Total Time (sec)	Update Time (sec)	Max Time (sec)	Other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8	Boost Hash	413/414/406	233/227/223	179/187/182	64k
bpe_train_updater_fine_grained_emhash8	Boost Hash	664/593/606	305/255/269	358/338/337	96k

The time for bpe_train_updater_fine_grained_emhash8 gradually increases. This is because as the number of pairs in pair_counts grows, the time to find the max will only increase, not decrease.

12. Porting the Max Heap to C++

We can also port this algorithm to C++. The C++ standard library’s <algorithm> header has functions like std::make_heap, std::push_heap, and std::pop_heap, which can provide the same functionality as Python. The C++ heap is a max heap by default, so we just need to overload operator<. We can define:

struct HeapItem{
    int count;
    std::vector<std::vector<int>> pair_string;
    std::pair<int,int> pair;

    bool operator<(const HeapItem& other) const;
};

You can refer to test_max_heap.cpp for how to use the standard library.

I didn’t use the standard library here; instead, I rewrote the Python version in C++ and tested its speed. It seems to be slightly faster than the standard library. The full code for the max heap is in max_heap.cpp.

Based on bpe_train_updater_fine_grained, I’ve implemented bpe_train_updater_fine_grained_heap.cpp. Based on bpe_train_updater_fine_grained_emhash8_set, I’ve implemented bpe_train_updater_fine_grained_heap_emhash8_set.cpp. Based on bpe_train_updater_fine_grained_emhash8_set9, I’ve implemented bpe_train_updater_fine_grained_heap_emhash8_set9.cpp.

The implementation logic is completely consistent with the Python version. Readers who are interested can read the code themselves. Below are the comparison test results:

Program	Hash Function	Total Time (sec)	Update Time (sec)	Max Time (sec)
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8_set	Boost Hash	192/192/194	117/117/117	75/75/77
bpe_train_updater_fine_grained_emhash8_set9	Boost Hash	168/170/171	107/108/109	61/62/61
bpe_train_updater_fine_grained_heap	Boost Hash	200/228/211	194/220/208	2/3/2
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	199/200/210	193/195/203	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set	Boost Hash	139/122/128	136/118/123	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set9	Boost Hash	111/121/122	109/116/116	2/2/2

As you can see, the time to find the max using a heap is less than 2 seconds. Now let’s look at its performance with vocabularies of 64,000 and 96,000:

Program	Hash Function	Total Time (sec)	Update Time (sec)	Max Time (sec)	Other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8	Boost Hash	413/414/406	233/227/223	179/187/182	64k
bpe_train_updater_fine_grained_emhash8	Boost Hash	664/593/606	305/255/269	358/338/337	96k
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	199/200/210	193/195/203	2/2/2
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	224/228/227	216/218/217	3/3/3	64k
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	249/255/240	239/238/230	5/6/4	96k

By comparing them, we can see that for the bpe_train_updater_fine_grained_heap_emhash8 algorithm, the update time only increases by about 20 seconds when the vocabulary increases from 32k to 64k, and the max time is almost unchanged. In contrast, the max time for bpe_train_updater_fine_grained_emhash8 doubles.

13. Optimizing `_updated_affected_word_count`

If you carefully read the _updated_affected_word_count function and fine_grained_pair_counter_diff, you can see there’s redundancy between new_pairs and diff_pairs. Recall that diff_pairs records the change in a pair’s count: a value greater than zero indicates an increase (new pair), and a value less than zero indicates a decrease (old pair). new_pairs represents all affected pairs (including pairs that may not have changed, which happens when the merge_pair appears multiple times with other pairs in between). However, based on our previous algorithm, we only need to add new pairs to the heap with heappush, and old pairs can be updated lazily. So, we can remove new_pairs, and any pair in diff_pairs with a count > 0 must be a new pair.

This leads to bpe_v7_maxheapc_opt_time.py. Other versions can be modified similarly.

The main code changes are:

    @staticmethod
    def _updated_affected_word_count(merge_pair, affected_words, word_encodings, 
                                     word_counts, pair_counts, pair_to_words, 
                                     new_id, pair_strings, vocabulary, pair_heap):
        # we may update/delete words when iterate it.
        affected_words = affected_words.copy()
        diff_pairs = defaultdict(int)

        BPE_Trainer.fine_grained_pair_counter_diff(affected_words, word_encodings, word_counts, merge_pair, diff_pairs, 
                             new_id, pair_to_words)
        for pair, count in diff_pairs.items():
            if count == 0: continue
            pair_counts[pair] += count
            if count > 0: # new pair
                pair_strings[pair] = (vocabulary[pair[0]], vocabulary[pair[1]])
                maxheap.heappush(pair_heap, (pair_counts[pair], pair_strings[pair], pair))
            
            if pair_counts[pair] <= 0: # should not less than 0!
                del pair_counts[pair]
                pair_to_words.pop(pair, None)

BPE_Trainer.fine_grained_pair_counter_diff no longer needs to take new_pairs as an output parameter. When iterating through diff_pairs, we check the count; if count > 0, it’s a new pair, and we save its pair_strings and add it to the heap with heappush.

Version	Data	Total Time (s)	Word Count Time (s)	Merge Time (s)	Other
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1
bpe_v6_time	open_web	1036/1107/1046	395/395/398	total: 576/641/591 max:6/7/6 update: 570/633/584	num_counter=8, num_merger=1
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7_maxheapc_time	open_web	1069/1017/1011	400/401/399	total: 606/556/555 max: 3/3/3 update: 602/552/552	num_counter=8, num_merger=1
bpe_v7_maxheapc_opt_time	open_web	984/965/1000	394/394/403	total: 532/514/538 max: 0.8/0.8/0.9 update: 531/513/537	num_counter=8, num_merger=1

bpe_v7_maxheapc_opt_time is about 7% faster than bpe_v7_maxheapc_time.

14. Porting the Optimization to C++

The optimized versions are bpe_train_updater_fine_grained_heap_opt.cpp, bpe_train_updater_fine_grained_heap_emhash8_opt.cpp, bpe_train_updater_fine_grained_heap_emhash8_set_opt.cpp, and bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp.

Test results on the OpenWeb dataset:

Program	Hash Function	Total Time (sec)	Update Time (sec)	Max Time (sec)
bpe_train_updater_fine_grained_heap	Boost Hash	200/228/211	194/220/208	2/3/2
bpe_train_updater_fine_grained_heap_opt	Boost Hash	190/192/190	187/189/187	0/0/0
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	199/200/210	193/195/203	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_opt	Boost Hash	188/213/192	185/209/189	0/0/0
bpe_train_updater_fine_grained_heap_emhash8_set	Boost Hash	139/122/128	136/118/123	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set_opt	Boost Hash	110/130/110	108/128/108	0/0/0
bpe_train_updater_fine_grained_heap_emhash8_set9	Boost Hash	111/121/122	109/116/116	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set9_opt	Boost Hash	105/102/104	102/100/101	0/0/0

bpe_train_updater_fine_grained_heap_emhash8_set9_opt is 11% faster in update time than bpe_train_updater_fine_grained_heap_emhash8_set9. bpe_train_updater_fine_grained_heap_emhash8_set_opt is 8% faster than bpe_train_updater_fine_grained_heap_emhash8_set. This conclusion is consistent with the Python version.

Full Series

Part 0: Introduction Introduces the basic BPE training algorithm and related tasks, as well as the development environment.
Part 1: The Simplest Implementation The simplest implementation of BPE training.
Part 2: Optimized Algorithm Implements incremental updates for pair_counts.
Part 3: Parallel Tokenization and Frequency Counting Uses multiprocessing to implement a multi-process parallel algorithm.
Part 4: A Failed Parallel Optimization An attempt to parallelize the max pair calculation using multiple processes.
Part 5: Implementing the Merge Algorithm in C++ Implements a C++ merge algorithm equivalent to the Python version, and compares two ways of iterating through std::unordered_map.
Part 6: Parallelizing the Max Pair Search with OpenMP Uses OpenMP to find the max pair in pair_counts in parallel.
Part 7: Using Flat Hashmap to Replace std::unordered_map Uses flat hashmap to replace std::unordered_map.
Part 8: Implementing Fine-Grained Updates Implements a fine-grained update algorithm for pair_counts using an inverted index.
Part 9: Using a Heap to Find the Max Pair Uses a heap to find the max pair and improve performance.
Part 10: Using Cython and PyPy for Acceleration Uses Cython and PyPy to accelerate Python code.
Part 11: Wrapping C++ Code with Cython Wraps C++ code using Cython.

动手实现和优化BPE Tokenizer的训练——第9部分：使用堆来寻找最大pair

Sun, 21 Sep 2025 00:00:00 +0000

本系列文章完成Stanford CS336作业1的一个子任务——实现BPE Tokenizer的高效训练算法。通过一系列优化，我们的算法在OpenWebText上的训练时间从最初的10多个小时优化到小于10分钟。本系列文章解释这一系列优化过程，包括：算法的优化，数据结构的优化，并行(openmp)优化，cython优化，用c++实现关键代码和c++库的cython集成等内容。本文是第十篇，使用堆(heap)这个数据结构来替代求最大pair，提升性能。

1. 问题分析
2. 堆(heap)
3. 使用堆求max
4. Python的heapq模块
5. 实现一个大堆
6. 用c模块优化maxheap_py
7. 用heapq实现大堆
8. 不同的大堆实现的性能测试
9. 用大堆求max
10. 测试
11. 更大vocab_size的实验
12. 把最大堆移植到c++
13. 优化_updated_affected_word_count
14. 把上面的优化移植到c++
本系列全部文章

1. 问题分析

在上面我们的bpe_v5_time中,update time只有500多秒，但是max的时间是30000多秒。怎么才能优化这个时间呢？之前我们尝试的办法是用并行算法求max，Python的版本并不成功，因为Python的GIL使得多线程无法同时使用多个CPU，而多进程又存在大量的通信。我们又尝试了在c++里用OpenMP实现了并行，在32线程的时候max的时间降低到400秒。接着我们又用更快的flat hashmap替代了std::unordered_map，使得max的时间在不用并行算法就能把max时间降低到100秒之内。

不过我们还是回到Python，除了优化数据结构，我们还有没有办法能够优化算法呢？优化数据结构就是让同样的工作做得更快，而优化算法是用更少的工作达到同样的目的。当然这两者也不是完全分开的，有的时候为了实现某种不同的算法就需要设计合适的数据结构来完成。

为了找到一个集合中的最大值，我们当然需要遍历整个集合，这个时间不可能更少。也就是说第一次遍历整个集合的时间是不可能变少的。但是第二次遍历时我们就有可能优化了，因为两次遍历之间只有部分pair的计数是发生变化的。刚开始都是高频词，因此affected_words很多，所以受到影响的pair相对较多。而到了后面，词频没有那么高了，affected_words就少了。最简单的想法就是我们可以给pair_counts排序，然后某些pair的计数发生了变化，那么就可以只对这些pair重新排序。因为我们的词频发生变化只有两种可能：老的pair词频减少，新增的pair词频增加。所以我们还可以用一些启发式规则，比如某个pair当前的序号是大于32000(我们最终需要的词表大小)，并且它的频次变低了，那么我们可以忽略它。

不过我们还有更好的方法来找到一个集合的最大值，那就是使用堆这种数据结构。

2. 堆(heap)

在计算机领域，堆有两个完全不同的意思：一是表示一种数据结构；二是一种内存分配方式。有趣的是在表示内存分配方式时和堆对应的是栈(stack)，而栈除了表示内存分配方式之外也表示一种数据结构。不过作为数据结构的栈和内存分配的栈还是存在紧密联系的，比如数据结构的栈是一种后进先出的数据结构(和先入先出的队列对应)，而内存分配的栈(函数)正是利用了这种后进先出的特点，使得函数的调用和返回顺序符合我们的期望。但是数据结构的堆和内存分配的堆就没有什么联系了，它们唯一的联系就是当初起名字的人不知怎的偶然用了相同的词表示了完全不同两个概念。

我们这里关注的是作为数据结构的堆，它通常用于实现优先队列，也可以用来实现堆排序。我这里不再详细介绍堆，如果不熟悉的读者随便找一本数据结构和算法的书，或者上网搜索都有很多介绍，比如wiki。

3. 使用堆求max

一般使用堆的时候我们只是使用如下3个操作：把一个数组变成一个堆(heapify)、从堆顶弹出一个元素(heappop)和把一个元素加入堆(heappush)。我们首先需要通过heapify把一个数组变成堆(满足堆的定义：树根比它的子树中的每一个都大/小)，然后不断的调用heappop/heappush，这两个操作完成之后的数组依然还是堆。

这里有一个问题，那就是每次找完最大之后(heappop)，我们会合并pair，这会使得某些老的pair的计数减少，同时一些新的pair出现(原来没有)。增加新的pair没有问题，我们只需要调用heappush就可以了。但是怎么修改老的pair的计数并且使得修改之后还是堆呢？我们来看一个例子：

         11
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   / \
6  4  5   7

如果我们把8改成10，也就是一个元素变大了，那么需要从8开始往上调用siftup：

         11
       /    \
      /      \
     10       4
   /   \     / \
  7     9   3   1 
 / \   / \
6  4  5   7

而反正如果我们把8改成了6，那么需要从8开始往下调用siftdown：

         11
       /    \
      /      \
     10       4
   /   \     / \
  7     7   3   1 
 / \   / \
6  4  5   6

但是这里有一个问题，那就是我们怎么找到这个修改的元素。请读者回顾一下，我们的主要数据是pair_counts，这是一个dict，但是堆需要操作的是一个数组(list)，因此我们需要把pair_counts中的元素复制一份到pair_heap这个list里。但是list无法支持快速查找，如果为了修改某个元素就需要顺序扫描一遍list，那就得不偿失了(max都求出来了)。

一种解决办法是在pair_counts里记录它在list里的下标，这样pair_counts变成：

pair -> (count, index_in_pair_heap)

那么新增一个pair到pair_counts的同时也把这pair通过heappush添加到pair_heap的合适位置，这就要求heappush函数不但可以把一个元素添加到pair_heap数组，而且还需要返回它在pair_heap的下标。这样我们才能把这个下标保存到pair_counts里。等到后面某个pair的计数发生改变(我们这里只会减少)，我们就可以通过pair_counts找到它在pair_heap的位置，然后对这个位置的元素调用siftdown。

这就需要修改heappush，而且还需要维护pair_counts和pair_heap的关系，这会使得代码变得很复杂。感兴趣的读者可以尝试实现一下这种算法。

不过我这里使用的是另外一种方法——lazy的修改。这种方法当某个pair的计数发生改变时，我们什么也不做。只有当我们调用heappop时需要检查一下它的计数是否修改了，这可以通过和pair_counts里的对比来发现。如果修改了(只会变小，这个假设非常重要)，那么我们把新的计数重新通过heappush加进去。然后不断的heappop，直到某个元素没有被修改，那么它就是当前的最大值。

说起来比较复杂，我通过一个例子来看。比如当前的堆为：

         11
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   / \
6  4  5   7

最大的应该是11，但是我们假设因为合并11变成了8，那么当我们弹出11这个pair时，我们查询pair_counts得知它的最新计数是10，因为它变小了，所以我们不能确定它是否最大。所以我们首先把11弹出，得到：

          9
       /    \
      /      \
     8        4
   /   \     / \
  7     7   3   1 
 / \   /
6  4  5 

然后还需要把10重新heappush进去，得到：

         10
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   /  \
6  4  5    7

接着，我们再次弹出当前最大10，这个时候查询pair_counts得知它的计数是最新的，因此就找到了当前的最大pair是10。

注意：我们可以lazy更新的一个重要假设是老pair的计数只会变小。如果这个假设不成立，比如当前我们的堆是：

         11
       /    \
      /      \
     9        4
   /   \     / \
  7     8   3   1 
 / \   / \
6  4  5   7

我们把1变成了12，那么我们就必须马上对1进行siftup，否则我们找到的最大是11，但这是不正确的。

使用这种算法，我们不需要在pair_counts里维护pair_heap的下标，而且我们也不需要调用siftdown，这在Python的heapq里是一个私有函数_siftdown，使用私有函数是有风险的，也许新的版本就没有这个函数了。

4. Python的heapq模块

Python的标准库提供了heapq。我们需要用到的主要函数是heappush、heappop和heapify。不熟悉的读者可以参考The Python heapq Module: Using Heaps and Priority Queues。

不过这里有一个问题，我们需要的是一个大堆(max heap)，但是Python的heapq模块提供的是小堆(min heap)。我们后面分析其代码时会发现它的内部已经实现了大堆。不过这里我们需要讨论如果只有小堆我们能不能用它来实现大堆的功能。

一种常见的技巧是把元素逆转来实现大堆。比如如果堆的元素是正整数，那么我们可以存入它对应的负整数，这样就可以得到大堆了，比如下面的例子：

import heapq
arr = [3, 5, 1, 2, 6, 8, 7]

arr2 = [-i for i in arr]

arr2
[-3, -5, -1, -2, -6, -8, -7]

heapq.heapify(arr2)
arr2
[-8, -6, -7, -2, -5, -1, -3]

arr3 = [-i for i in arr2]
arr3
[8, 6, 7, 2, 5, 1, 3]

但是现在我们需要放入堆的是(count, pair_string[pair], pair)。如果count大，那么这个tuple就大；如果count相同，我们再比较这个pair的字符串，选择大的那个。tuple最后一个放入pair是为了使用方便。

如果按照上面的办法，count可以放入-count。pair_strings[pair]呢？pair_strings[pair]里是个tuple，这个tuple的每个元素都是bytes。如果是定长的bytes，因为一个byte的范围是0~255，那么我们可以用255减去这个数来逆转它。比如：

b1 = b'us'

b2 = b'ua'

b1 > b2
True

c1 = bytes([255 - b for b in b1])
c1
b'\x8a\x8c'

c2 = bytes([255 - b for b in b2])

c1 < c2
True

但是如果字符串不定长，就会出问题，比如：

b1 = b'us'

b2 = b'usb'

b1 < b2
True

c1 = bytes([255 - b for b in b1])

c2 = bytes([255 - b for b in b2])

c1 > c2
False

b1比b2少一个字符串。用255减去逆转之后前面的b’us’是相同的，但是不管是否逆转，位数少的都是小于位数多的。

所以我们需要一个大堆。

5. 实现一个大堆

我们只需要稍微修改一下Python自带的heapq模块就可以把它改成大堆，所以我们把heapq的源代码复制一遍然后修改成maxheap_py.py。heapq的函数很多，我们只需要保留heappush、heappop和heapify。另外这3个函数又依赖_siftdown(我把它改名为_siftdown_max)和_siftup(改名_siftup_max)。

完整代码我就不展开，这里就对比一下_siftdown_max和_siftdown：

def _siftdown_max(heap, startpos, pos):
    'Maxheap variant of _siftdown'
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if parent < newitem:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem

def _siftdown(heap, startpos, pos):
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if newitem < parent:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem

它们唯一的区别就是”if newitem < parent”和”if parent < newitem”这两行代码。这里假设newitem和parent值实现/重载了小于运算符，所以即使想求较大的元素也是通过调换顺序通过小于函数来实现的。

6. 用c模块优化maxheap_py

对于c模块开发不感兴趣的读者可以跳过本节内容。

但是如果我们比较自己版本速度，它会把CPython的版本慢很多，原因在于CPython的版本内部是调用了对应的c模块。如果我们仔细阅读heap源代码，我们会发现：

# If available, use C implementation
try:
    from _heapq import *
except ImportError:
    pass
try:
    from _heapq import _heapreplace_max
except ImportError:
    pass
try:
    from _heapq import _heapify_max
except ImportError:
    pass
try:
    from _heapq import _heappop_max
except ImportError:
    pass

也就是它会(尝试)调用对应的c模块实现。具体的代码在_heapqmodule.c里。

我复制它实现了大堆的功能，完整代码在_maxheapqmodule.c。

我们只看其中一个函数siftdown：

static int
siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
{
    PyObject *newitem, *parent, **arr;
    Py_ssize_t parentpos, size;
    int cmp;

    assert(PyList_Check(heap));
    size = PyList_GET_SIZE(heap);
    if (pos >= size) {
        PyErr_SetString(PyExc_IndexError, "index out of range");
        return -1;
    }

    /* Follow the path to the root, moving parents down until finding
       a place newitem fits. */
    arr = _PyList_ITEMS(heap);
    newitem = arr[pos];
    while (pos > startpos) {
        parentpos = (pos - 1) >> 1;
        parent = arr[parentpos];
        Py_INCREF(newitem);
        Py_INCREF(parent);
        cmp = PyObject_RichCompareBool(newitem, parent, Py_GT);
        Py_DECREF(parent);
        Py_DECREF(newitem);
        if (cmp < 0)
            return -1;
        if (size != PyList_GET_SIZE(heap)) {
            PyErr_SetString(PyExc_RuntimeError,
                            "list changed size during iteration");
            return -1;
        }
        if (cmp == 0)
            break;
        arr = _PyList_ITEMS(heap);
        parent = arr[parentpos];
        newitem = arr[pos];
        arr[parentpos] = newitem;
        arr[pos] = parent;
        pos = parentpos;
    }
    return 0;
}

两种对比代码只有一行区别：

cmp = PyObject_RichCompareBool(newitem, parent, Py_GT);
cmp = PyObject_RichCompareBool(newitem, parent, Py_LT);

前面介绍过PyObject_RichCompareBool，它就是根据第三个参数比较前两个参数。如果返回值小于0，表示异常；如果返回值大于0，则表示第三个参数，第一个比第二个大/小是true；如果返回值是0，则表示第一个参数比第二个参数大/小是false。

因此我们把小堆改成大堆只需要把比较运算符从Py_LT(小于)变成Py_GT(大于)就行了。

当然，其它还有一些细节需要修改，比如我把模块名字改成了maxheapqc，那么就需要修改相应的代码：

static struct PyModuleDef _heapqmodule = {
    PyModuleDef_HEAD_INIT,
    "maxheapqc",
    module_doc,
    0,
    heapq_methods,
    heapq_slots,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit_maxheapqc(void)
{
    return PyModuleDef_Init(&_heapqmodule);
}

另外为了编译这个C模块，需要修改setup.py：

include_path = sysconfig.get_path('include')
internal_include_path = os.path.join(include_path, 'internal')

print(f"{internal_include_path=}")
project_root = os.path.dirname(os.path.abspath(__file__))

maxheapqc_module = Extension('cs336_basics.maxheapqc', sources=['cs336_basics/_maxheapqmodule.c'],
        extra_compile_args=[
            f'-I{internal_include_path}',
        ]
        )

setup(
    name='maxheapqc',
    version='1.0',
    description='maxheapqc',
    packages=['cs336_basics'],
    ext_modules=[maxheapqc_module]
)

我们在编译_maxheapqmodule.c时，需要依赖：

#include "Python.h"
#include "pycore_list.h" 

对于Python.h，setuptools会帮我们包含这个头文件，它的位置通常类似于：/usr/local/include/python3.12/Python.h。但是我们的_maxheapqmodule.c还需要通过c接口操作list，这需要pycore_list.h，它的位置在类似于/usr/local/include/python3.12/internal/pycore_list.h的地方。我们不想硬编码这个位置，而且同一个系统里可能安装了很多版本的python，还有很多conda这样的环境。所以我们可以通过sysconfig.get_path(‘include’)得到当前Python解释器的路径，然后再在里面找internal。

我们在编译时需要通过-I添加这个路径：

maxheapqc_module = Extension('cs336_basics.maxheapqc', sources=['cs336_basics/_maxheapqmodule.c'],
        extra_compile_args=[
            f'-I{internal_include_path}',
        ]
        )

这个-I参数好像是gcc添加头文件的方式，我不知道在非linux系统非gcc编译器下能否通用，如果读者使用的是其它编译器，请参考编译器的手册添加合适的头文件路径。

运行下面的命令编译c模块：

python setup.py build_ext -i

编译后可以得到类似cs336_basics/maxheapqc.cpython-312-x86_64-linux-gnu.so。接下来我们就可以写一个maxheapq.py来调用它了。

maxheapq.py的代码和heapq.py类似，它用Python实现了大堆的功能，并且会尝试加载maxheapqc：

# If available, use C implementation
try:
    from cs336_basics.maxheapqc import *
    print("load c!")
except ImportError:
    pass
try:
    from cs336_basics.maxheapqc import _heapreplace_max
except ImportError:
    pass
try:
    from cs336_basics.maxheapqc import _heapify_max
except ImportError:
    pass
try:
    from cs336_basics.maxheapqc import _heappop_max
except ImportError:
    pass

7. 用heapq实现大堆

其实如果我们阅读了heapq的代码，我们会发现它已经实现了_heapify_max、_siftdown_max和_siftup_max。而其中_heapify_max是有对应的c模块代码，速度比较快。但是_siftdown_max和_siftup_max还是用Python实现的。我们可以查看这里来找到哪些函数有c语言实现：

static PyMethodDef heapq_methods[] = {
    _HEAPQ_HEAPPUSH_METHODDEF
    _HEAPQ_HEAPPUSHPOP_METHODDEF
    _HEAPQ_HEAPPOP_METHODDEF
    _HEAPQ_HEAPREPLACE_METHODDEF
    _HEAPQ_HEAPIFY_METHODDEF
    _HEAPQ__HEAPPOP_MAX_METHODDEF
    _HEAPQ__HEAPIFY_MAX_METHODDEF
    _HEAPQ__HEAPREPLACE_MAX_METHODDEF
    {NULL, NULL}           /* sentinel */
};

我们可以利用heapq的这3个函数来实现大堆。代码在maxheap_heapq.py。我们来看一下：

def heappush(heap, item):
    heap.append(item)
    heapq._siftdown_max(heap, 0, len(heap)-1)


def heappop(heap):
    """Maxheap version of a heappop."""
    lastelt = heap.pop()  # raises appropriate IndexError if heap is empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        heapq._siftup_max(heap, 0)
        return returnitem
    return lastelt

def heapify(heap):
    heapq._heapify_max(heap)

这里heapify会调用c接口，速度比较快，但是heappush和heappop还是Python实现。

8. 不同的大堆实现的性能测试

我写了一个简单的测试代码test_heap_speed.py。结果如下：

maxheapq_py: 40.50480842590332
maxheap_heapq: 4.577162265777588
maxheapq: 4.334884405136108

这里主要的时间用在heapify，maxheapq_py是Python实现，所以比较慢。

而test_heap_speed2.py主要测试heappush和heappop，结果如下：

maxheapq_py: 1.892250157892704
maxheap_heapq: 1.8967562650796026
maxheapq: 0.3524886401137337

可以看到maxheapq要比Python的实现快。

9. 用大堆求max

完整的代码在bpe_v6.py。我们只看一下它和bep_v5.py的区别。

from cs336_basics import maxheap_py as maxheap

    def train(self, input_path, vocab_size, special_tokens, *args):
        ...
        pair_heap = []
        for pair, count in pair_counts.items():
            maxheap.heappush(pair_heap, (count, pair_strings[pair], pair))
            
            
            

首先需要import maxheap_py，为了便于切换不同的大堆实现，我把它重命名为maxheap。因为所有不同大堆的接口都相同，所以要切换为maxheapq或者maxheap_heapq只需要修改这一行代码就可以了。

然后在第一次统计了pair_counts之后，我们需要构造一个pair_heap，这里是不停调用maxheap.heappush。另外一种实现方法是使用heapify：

        pair_heap = []
        for pair, count in pair_counts.items():
            pair_heap.append((count, pair_strings[pair], pair))
        maxheap.heapify(pair_heap)

两种方法最终都会构建出一个堆，但是它们的结果可能不完全相同。理论上用heapify的方法更快一点，不过实际测试发现它们的时间差很小，因为pair_counts一开始不到两万。

接下来是主要的修改之处，通过heappop求当前最大的pair：

    @staticmethod
    def _merge_a_pair(pair_counts, pair_strings, vocabulary, pair_to_words, 
                   word_counts, word_encodings, merges, size, pair_heap):
        
        while pair_heap:
            count, string_priority, merge_pair = maxheap.heappop(pair_heap)
            
            # check pair validity
            if merge_pair in pair_counts and pair_counts[merge_pair] == count:
                break
            elif merge_pair in pair_counts:
                # update count (lazily)
                maxheap.heappush(pair_heap, (pair_counts[merge_pair], 
                                               string_priority, 
                                               merge_pair))
        else:
            # no valid pairs found
            return False

这个算法的实现和前面描述的一样：首先从堆顶弹出当前最大的merge_pair，然后通过对比pair_counts[merge_pair]来检查它的count有没有更新，如果没有更新，那么merge_pair就是当前的最大pair，break出循环。如果有了更新而且新的count>0(merge_pair in pair_counts)，则需要把新的count重新插入到堆里。然后再循环去找最大的pair。

最后一处修改就是每次merge产生新的pair时，我们需要把它也加到堆里：

    @staticmethod
    def _updated_affected_word_count(merge_pair, affected_words, word_encodings, 
                                     word_counts, pair_counts, pair_to_words, 
                                     new_id, pair_strings, vocabulary, pair_heap):



        for new_pair in new_pairs:
            if new_pair not in pair_strings:
                pair_strings[new_pair] = (vocabulary[new_pair[0]], vocabulary[new_pair[1]])

            maxheap.heappush(pair_heap, (pair_counts[new_pair], pair_strings[new_pair], new_pair))

10. 测试

为了测试时间，我也实现了bpe_v6_time.py。测试结果如下：

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1
bpe_v6_time	open_web	1036/1107/1046	395/395/398	total: 576/641/591 max:6/7/6 update: 570/633/584	num_counter=8, num_merger=1

使用堆来求最大之后，合并时间从30000多秒减少到了570多秒。heappush的时间是100多秒，我们还能通过其它更快的大堆来优化它吗？

把maxheap_py改成maxheap_heapq或者maxheapq就得到bpe_v7.py和bpe_v7_maxheapc.py，测试结果如下：

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1
bpe_v6_time	open_web	1036/1107/1046	395/395/398	total: 576/641/591 max:6/7/6 update: 570/633/584	num_counter=8, num_merger=1
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7_maxheapc_time	open_web	1069/1017/1011	400/401/399	total: 606/556/555 max: 3/3/3 update: 602/552/552	num_counter=8, num_merger=1

使用更快的堆实现并没有加快速度。我猜测是因为heappush的时间占总update时间的比例比较低，因此它的微小时间变化对于整体影响不大。

11. 更大vocab_size的实验

如果我们拿bpe_v7和c++版本的bpe_train_updater_fine_grained_emhash8相比：

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v7_time	open_web	1028/1053/1045	395/397/393	total: 575/589/590 max: 6/6/6 update: 569/583/583 make heap: 0.01/0.01 heap_push_time: 102/107/122	num_counter=8, num_merger=1

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60

bpe_v7的merge总时间是600秒左右，而bpe_train_updater_fine_grained_emhash8的总时间是260秒左右。但是如果我们把合并的次数增多呢？因为现在的大模型的词典越来越大，所以我们测试一下vocab_size是64000和96000的结果：

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v7_time	open_web	1028/1053/1045	395/397/393	total: 575/589/590 max: 6/6/6 update: 569/583/583 make heap: 0.01/0.01 heap_push_time: 102/107/122	num_counter=8, num_merger=1
bpe_v7_time2	open_web	1111/1174/1100	393/420/403	total: 655/686/639 max: 9/10/10 update: 645/675/628 make heap: 0.01/0.01/0.01 heap_push_time: 128/157/108	vocab_size=64000
bpe_v7_time2	open_web	1129/1130/1123	394/406/393	total: 675/666/670 max: 13/12/12 update: 661/653/657 make heap: 0.01/0.01/0.01 heap_push_time: 143/152/120	vocab_size=96000

可以看到由于后面处理的都是低频词和低频的pair，所以几乎时间没有什么变化。我们再看bpe_train_updater_fine_grained_emhash8：

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8	Boost Hash	413/414/406	233/227/223	179/187/182	64k
bpe_train_updater_fine_grained_emhash8	Boost Hash	664/593/606	305/255/269	358/338/337	96k

可以看到bpe_train_updater_fine_grained_emhash8的时间逐渐增大，这是因为随着pair_counts中pair的增多，max的时间只会增加不会减少。

12. 把最大堆移植到c++

我们也可以把这个算法移植到c++，c++的标准库的里有std::make_heap、std::push_heap和std::pop_heap等函数，可以实现和python相同的功能，而且c++的堆是大堆，我们只需要实现operator<的重载即可，我们可以定义：

struct HeapItem{
    int count;
    std::vector<std::vector<int>> pair_string;
    std::pair<int,int> pair;

    bool operator<(const HeapItem& other) const;
};

标准库的用法可以参考test_max_heap.cpp。

我这里没有使用标准库，而是把Python版本的用c++重写了，测试了一下速度，似乎比标准库的还要稍微快一点点。大堆的完整代码在max_heap.cpp。

基于bpe_train_updater_fine_grained实现了bpe_train_updater_fine_grained_heap.cpp，基于bpe_train_updater_fine_grained_emhash8_set实现了bpe_train_updater_fine_grained_heap_emhash8_set.cpp，基于bpe_train_updater_fine_grained_emhash8_set9实现了bpe_train_updater_fine_grained_heap_emhash8_set9.cpp。

代码实现逻辑和Python版本完全一致，感兴趣的读者可以自行阅读。下面是对比实验结果：

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8_set	Boost Hash	192/192/194	117/117/117	75/75/77
bpe_train_updater_fine_grained_emhash8_set9	Boost Hash	168/170/171	107/108/109	61/62/61
bpe_train_updater_fine_grained_heap	Boost Hash	200/228/211	194/220/208	2/3/2
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	199/200/210	193/195/203	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set	Boost Hash	139/122/128	136/118/123	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set9	Boost Hash	111/121/122	109/116/116	2/2/2

可以看到，使用堆来求max的时间不到2秒。我们再来看一下它在合并词典64000和96000时的结果：

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8	Boost Hash	413/414/406	233/227/223	179/187/182	64k
bpe_train_updater_fine_grained_emhash8	Boost Hash	664/593/606	305/255/269	358/338/337	96k
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	199/200/210	193/195/203	2/2/2
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	224/228/227	216/218/217	3/3/3	64k
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	249/255/240	239/238/230	5/6/4	96k

通过对比我们可以发现，使用堆的bpe_train_updater_fine_grained_heap_emhash8算法词典从32k增加到64k时update时间值增加了20秒左右，max时间几乎没变；而bpe_train_updater_fine_grained_emhash8的max时间增加了一倍。

13. 优化_updated_affected_word_count

如果仔细阅读_updated_affected_word_count函数和fine_grained_pair_counter_diff，可以发现其中new_pairs和diff_pairs存在冗余。回顾一下，diff_pairs记录了pair的count的变化，如果大于零表示增加(新pair)；如果小于零表示减少(老pair)。而new_pairs表示受到影响的所有pair(还包括可能没有变化的pair，这种情况发生在merge_pair出现多次时夹在中间的那些没有变化的pair)。但是根据之前的算法，我们只需要把新pair通过heappush加入堆就行，老的pair可以lazy的方式更新。所以我们可以去掉new_pairs，而diff_pairs里count>0的pair就一定是新pair。

这样就得到bpe_v7_maxheapc_opt_time.py，其它版本也可以参考这个修改。

主要改动的代码是：

    @staticmethod
    def _updated_affected_word_count(merge_pair, affected_words, word_encodings, 
                                     word_counts, pair_counts, pair_to_words, 
                                     new_id, pair_strings, vocabulary, pair_heap):
        # we may update/delete words when iterate it.
        affected_words = affected_words.copy()
        diff_pairs = defaultdict(int)

        BPE_Trainer.fine_grained_pair_counter_diff(affected_words, word_encodings, word_counts, merge_pair, diff_pairs, 
                             new_id, pair_to_words)
        for pair, count in diff_pairs.items():
            if count == 0: continue
            pair_counts[pair] += count
            if count > 0: # new pair
                pair_strings[pair] = (vocabulary[pair[0]], vocabulary[pair[1]])
                maxheap.heappush(pair_heap, (pair_counts[pair], pair_strings[pair], pair))
            
            if pair_counts[pair] <= 0: # should not less than 0!
                del pair_counts[pair]
                pair_to_words.pop(pair, None)

BPE_Trainer.fine_grained_pair_counter_diff不再需要传入new_pairs这个出参。在遍历diff_pairs时我们判断一下count，如果count>0则说明是一个新加入的pair，我们保存它的pair_strings然后通过heappush加入堆中。

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1
bpe_v6_time	open_web	1036/1107/1046	395/395/398	total: 576/641/591 max:6/7/6 update: 570/633/584	num_counter=8, num_merger=1
bpe_v7_time	open_web	1062/1035/1036	392/397/395	total: 606/573/577 max:6/6/6 update: 599/567/571	num_counter=8, num_merger=1
bpe_v7_maxheapc_time	open_web	1069/1017/1011	400/401/399	total: 606/556/555 max: 3/3/3 update: 602/552/552	num_counter=8, num_merger=1
bpe_v7_maxheapc_opt_time	open_web	984/965/1000	394/394/403	total: 532/514/538 max: 0.8/0.8/0.9 update: 531/513/537	num_counter=8, num_merger=1

bpe_v7_maxheapc_opt_time比bpe_v7_maxheapc_time要快7%。

14. 把上面的优化移植到c++

优化后的版本是bpe_train_updater_fine_grained_heap_opt.cpp、bpe_train_updater_fine_grained_heap_emhash8_opt.cpp、bpe_train_updater_fine_grained_heap_emhash8_set_opt.cpp和bpe_train_updater_fine_grained_heap_emhash8_set9_opt.cpp。

在OpenWeb数据集上的测试结果为：

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater_fine_grained_heap	Boost Hash	200/228/211	194/220/208	2/3/2
bpe_train_updater_fine_grained_heap_opt	Boost Hash	190/192/190	187/189/187	0/0/0
bpe_train_updater_fine_grained_heap_emhash8	Boost Hash	199/200/210	193/195/203	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_opt	Boost Hash	188/213/192	185/209/189	0/0/0
bpe_train_updater_fine_grained_heap_emhash8_set	Boost Hash	139/122/128	136/118/123	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set_opt	Boost Hash	110/130/110	108/128/108	0/0/0
bpe_train_updater_fine_grained_heap_emhash8_set9	Boost Hash	111/121/122	109/116/116	2/2/2
bpe_train_updater_fine_grained_heap_emhash8_set9_opt	Boost Hash	105/102/104	102/100/101	0/0/0

bpe_train_updater_fine_grained_heap_emhash8_set9_opt比bpe_train_updater_fine_grained_heap_emhash8_set9的update时间快11%。bpe_train_updater_fine_grained_heap_emhash8_set_opt比bpe_train_updater_fine_grained_heap_emhash8_set快8%。这个结论和python版本是一致的。

本系列全部文章

第0部分：简介介绍bpe训练的基本算法和相关任务，并且介绍开发环境。
第1部分：最简单实现 bpe训练最简单的实现。
第2部分：优化算法实现pair_counts的增量更新。
第3部分：并行分词和统计词频使用multiprocessing实现多进程并行算法。
第4部分：一次失败的并行优化尝试用多进程并行计算max pair。
第5部分：用C++实现Merge算法用C++实现和Python等价的merge算法，并且比较std::unordered_map的两种遍历方式。
第6部分：用OpenMP实现并行求最大用OpenMP并行求pair_counts里最大pair。
第7部分：使用flat hashmap替代std::unordered_map 使用flat hashmap来替代std::unordered_map。
第8部分：实现细粒度更新使用倒排索引实现pair_counts的细粒度更新算法。
第9部分：使用堆来寻找最大pair 使用堆来求最大pair，提升性能。
第10部分：使用cython和pypy来加速使用cython和pypy来加速python代码。
第11部分：使用cython封装c++代码使用cython封装c++代码。

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 8: Implementing Fine-Grained Updates

Fri, 19 Sep 2025 00:00:00 +0000

This series of articles implements a subtask of Stanford’s CS336 Assignment 1: building an efficient training algorithm for a BPE Tokenizer. Through a series of optimizations, our algorithm’s training time on OpenWebText was reduced from over 10 hours to less than 10 minutes. This series explains these optimizations, including algorithmic improvements, data structure enhancements, parallelization with OpenMP, Cython optimization, and implementing key code in C++ along with its integration via Cython. This is the ninth article, which focuses on optimizing the update process for data structures like pair_counts.

Table of Content

1. Problem Analysis
2. Fine-Grained Frequency Update Algorithm
3. Code Implementation
4. Testing
5. Porting to C++
6. Using emhash8::HashSet/emhash9::HashSet to Replace std::unordered_set
Full Series

1. Problem Analysis

We previously left one optimization point unaddressed: when updating pair_counts, many pairs’ counts don’t actually need to change, but for the sake of simple implementation, we first delete them and then add them back. More specifically, let’s look at a previous example.

Suppose the word frequency count is:

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

Then we count the pair frequencies:

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 11, st: 11, ne: 6, ew: 6}

Based on the merge rules, we choose to merge ('s', 't'). Since they are merged, we need to update the pair frequency counts. Let’s look at our current update algorithm (the code is in BPE_Trainer._updated_affected_word_count, though a more accurate name would be _updated_pair_count_of_affected_word).

To get the new pair frequencies, we use pair_to_words to find the words containing the pair ('s', 't'):

'widest'
'newest'

We first subtract all pairs within these two words, which means subtracting:

{wi: 3, id: 3, de: 3, es: 3, st: 3}
{ne: 6, ew: 6, we: 6, es: 6, st: 6}

This gives us:

{lo: 7, ow: 7, we: 2, er: 2, wi: 0, id: 0, de: 0, es: 2, st: 2, ne: 0, ew: 0}

Next, the tokens for these two words become:

word_encodings['widest'] = ['w','i','d','e','st']
word_encodings['newest'] = ['n','e','s','e','st']

Therefore, based on the new word_encodings, we calculate the new pair frequencies:

{wi: 3, id: 3, de: 3, est: 3}
{ne: 6, ew: 6, we: 6, est: 6}

Adding these frequencies back gives us:

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 2, st: 2, ne: 6, ew: 6, est: 9}

As you can see, the frequencies for pairs like 'wi', 'id', and 'de' were subtracted and then added back, leading to unnecessary updates. Worse still, after 'wi' is deleted, its frequency becomes 0, so it needs to be deleted from pair_counts and pair_to_words.

                    if pair_counts[old_pair] <= 0:
                        # we accounted for all occurrences of this pair
                        del pair_counts[old_pair]
                        pair_to_words.pop(old_pair)

According to Internals of sets and dicts, CPython’s dict uses an open-addressing collision resolution strategy. Deleting a key from a dict just marks it (unless a rehash occurs), and adding the same key back will reuse the original address. So the impact on the dict itself is minimal, but the key and value objects still need to be reconstructed (since the dict stores pointers to them). If a separate chaining method like std::unordered_map is used, deleting and reinserting will mess up the original order of the linked list (the new element is inserted at the head). For example, if a bucket’s elements were a->b->c, deleting b and reinserting it might make the new list b->a->c.

However, if we analyze carefully, we can identify which deletions are unnecessary. For instance, in our example, we are merging ('s', 't'). In 'widest', the only pairs affected by the merge are ('e', 's') and ('s', 't') (there’s no pair after 't'). The pairs ['w', 'i'], ['i', 'd'], and ['d', 'e'] are unaffected. Therefore, we only need to remove 'es' and 'st' and add ('e', 'st').

2. Fine-Grained Frequency Update Algorithm

Based on the previous analysis, we can implement a fine-grained frequency update algorithm. The example we analyzed was a case where the pair to be merged appeared only once in a word. But what if the pair appears multiple times? One approach would be to fall back to the previous algorithm. The possibility of a pair appearing multiple times in a single word is not high. If it appears only once, it can be handled simply; if it appears multiple times, we can fall back to the old, coarse-grained update algorithm—first delete, then insert. I didn’t think of this method before writing the article, which led me to a more complex solution. I created a complicated algorithm that can handle a pair appearing multiple times in one word. Although this algorithm seems perfect, the implementation is much more difficult, and the payoff is not that great. But since I’ve already implemented it, I’ll continue with it.

For example, consider a hypothetical word "abcststefstbbef". We can first find all indices where ‘st’ appears. To make it easier to see, we’ll separate ‘st’ from other tokens with a space.

abc st st ef st bbef

Based on the previous analysis, the pairs before the first ‘st’ ('ab', 'bc') are unaffected, and the pairs after the last ‘st’ ('bb', 'be') are also unaffected. What about the pairs in the middle? Let’s write them out:

cs st ts st te ef fs st tb

These pairs first need to be deleted, and then we generate new tokens for the affected section:

c st st e f st b

Then we add the new pairs:

('c','st'), ('st','st'), ('st','e'), ('e','f'), ('f','st'), ('st','b')

Besides decreasing the counts of the affected old pairs and increasing the counts of the new pairs, we also need to modify the inverted index pair_to_words. My earliest code would simply delete the affected pairs from it. But this has a bug. For example, 'ef' is affected and needs its count reduced. However, 'ef' also appears at the end of the word, not just in the middle of the three ‘st’ occurrences. If we simply delete 'ef' from pair_to_words, the inverted index becomes incorrect. A “true” inverted index not only records that the pair ('e', 'f') appeared in the word 'abcststefstbbef' but also how many times it appeared. This would mean that the value of pair_to_words should be a dict mapping words to their pair frequencies, not a set. So we only know that the pair ('e', 'f') appeared in 'abcststefstbbef', but we don’t know how many times. This is where the problem lies. For example, in the case above, if we simply delete the word 'abcststefstbbef' from pair_to_words[('e','f')]’s value set, the inverted index would be wrong. However, we must delete 'abcststefstbbef' from the value of pair_to_words[('t','e')], because the pair ('t', 'e') appears only once in the word 'abcststefstbbef'.

One solution is to change the pair_to_words code to a {pair: {word: word_freq}} structure. This is a clearer way to do it. However, at the time, I felt that this change would require too many modifications, and I wanted to limit the changes to a single function. So I used another method (which upon reflection is not ideal). This method is: find all the unaffected pairs and put them into a set. When modifying the inverted index, if the pair to be deleted is also in the set of unaffected pairs, we don’t delete it. Logically, this algorithm is correct, but it makes the code exceptionally complex. And for a case that rarely occurs, it complicates the code and slows down the common cases, which is not a worthwhile trade-off.

Nonetheless, scanning all the unaffected pairs is relatively fast, and many later improved versions are based on this current algorithm, so modifying it would be a hassle. For the time being, I will not change it in this article.

3. Code Implementation

The complete code is in bpe_v5.py. Let’s look at the changes. First, the modification to _updated_affected_word_count:

    @staticmethod
    def _updated_affected_word_count(merge_pair, affected_words, word_encodings, 
                                     word_counts, pair_counts, pair_to_words, 
                                     new_id, pair_strings, vocabulary):
        # we may update/delete words when iterate it.
        affected_words = affected_words.copy()
        diff_pairs = defaultdict(int)

        new_pairs = set() 
        BPE_Trainer.fine_grained_pair_counter_diff(affected_words, word_encodings, word_counts, merge_pair, diff_pairs, 
                             new_id, pair_to_words, new_pairs)
        for pair, count in diff_pairs.items():
            if count == 0: continue
            pair_counts[pair] += count
            if pair_counts[pair] <= 0: # should not less than 0!
                del pair_counts[pair]
                pair_to_words.pop(pair, None)


        for new_pair in new_pairs:
            if new_pair not in pair_strings:
                pair_strings[new_pair] = (vocabulary[new_pair[0]], vocabulary[new_pair[1]])

It becomes very simple. It calls a new function, BPE_Trainer.fine_grained_pair_counter_diff, to calculate the changes in pair counts after a merge. The main output of this function is diff_pairs. For example, a hypothetical output might be:

('a','s'): -3
('t','f'): -5
('st','b'):10

This indicates that due to the merge of ('s','t'), the count of the old pair ('a','s') decreased by 3, and ('t','f') decreased by 5; while the new pair ('st','b') increased by 10. The for loop below then uses this diff to update pair_counts. Additionally, the BPE_Trainer.fine_grained_pair_counter_diff function also returns the new_pairs that were created after the merge, and we need to update pair_strings with these new pairs as well. Let’s focus on the BPE_Trainer.fine_grained_pair_counter_diff function.

    @staticmethod
    def fine_grained_pair_counter_diff(affected_words, word_encodings, word_counts, merge_pair, diff_pairs, new_id, pair_to_words, new_pairs):
        for word in affected_words:
            word_tokens = word_encodings[word]
            wc = word_counts[word]

            # find first and last pairs
            idx = 0
            unaffected_pairs = set()
            while idx < len(word_tokens) - 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    first_idx = idx
                    break
                idx += 1
            else:
                print(f"bug {merge_pair}, {word}, {word_tokens}")
                raise
            # assert first_idx exists

            idx = len(word_tokens) - 2
            while idx > first_idx + 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    last_idx = idx
                    break
                idx -= 1
            else:
                last_idx = first_idx

            start_idx = max(0, first_idx - 1) # inclusive
            end_idx = min(last_idx + 3, len(word_tokens)) # exclusive

            # unaffected [0, start_idx)
            
            for i in range(start_idx):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair)
            # unaffected [end_idx-1, :-1]
            for i in range(end_idx - 1, len(word_tokens) - 1):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair)                

            affected_tokens = word_tokens[start_idx: end_idx]
            for i in range(len(affected_tokens) - 1):
                old_pair = (affected_tokens[i], affected_tokens[i + 1])
                diff_pairs[old_pair] -= wc 
                if old_pair not in unaffected_pairs:   
                    pair_to_words[old_pair].discard(word)
                    

            new_tokens = []
            all_new_tokens = []
            for i in range(start_idx):
                all_new_tokens.append(word_tokens[i])
            
            i = 0

            # account for multiple occurrences of the pair
            while i < len(affected_tokens):
                if i < len(affected_tokens) - 1 and (affected_tokens[i], affected_tokens[i + 1]) == merge_pair:
                    new_tokens.append(new_id)
                    all_new_tokens.append(new_id)
                    # jump past pair
                    i += 2
                else:
                    new_tokens.append(affected_tokens[i])
                    all_new_tokens.append(affected_tokens[i])
                    i += 1
            

            for i in range(end_idx, len(word_tokens)):
                all_new_tokens.append(word_tokens[i])
            
            word_encodings[word] = all_new_tokens

            # add new pairs from the updated word
            for i in range(len(new_tokens) - 1):
                new_pair = (new_tokens[i], new_tokens[i + 1])

                diff_pairs[new_pair] += wc
                pair_to_words[new_pair].add(word)

                new_pairs.add(new_pair)

This code is quite complex, especially the index handling. Let’s look at it section by section.

Finding the first and last occurrences of the merged pair

We first need to find the first and last occurrence of the pair. If it only appears once, these two positions are the same. The range we care about is from the first element of the first occurrence to the last element of the last occurrence.

            idx = 0
            unaffected_pairs = set()
            while idx < len(word_tokens) - 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    first_idx = idx
                    break
                idx += 1
            else:
                print(f"bug {merge_pair}, {word}, {word_tokens}")
                raise
            # assert first_idx exists

            idx = len(word_tokens) - 2
            while idx > first_idx + 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    last_idx = idx
                    break
                idx -= 1
            else:
                last_idx = first_idx

The first while loop finds the first occurrence of the pair and stores its index in first_idx. The else block should not be executed (unless our inverted index has a bug), so it’s a defensive exception. The second while loop works backward from the end to find the last occurrence, storing its index in last_idx. If it doesn’t find it, it means the pair only appeared once, so last_idx will be the same as first_idx.

Using the previous example, we find that first_idx is 3 and last_idx is 9. As shown below:

"abcststefstbbef"
 012345678901234
    |     | 
first_idx |
       last_idx  

Next, we find all the unaffected pairs:

            start_idx = max(0, first_idx - 1) # inclusive
            end_idx = min(last_idx + 3, len(word_tokens)) # exclusive

            # unaffected [0, start_idx) 
            for i in range(start_idx):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair)
                
            # unaffected [end_idx-1, :-1]
            for i in range(end_idx - 1, len(word_tokens) - 1):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair) 

The index calculations here are particularly tricky, so I won’t go into details. The final result is that the following unaffected pairs are stored in unaffected_pairs:

ab bc bb be ef

Next is to decrease the frequency of the affected pairs and update the inverted index:

            affected_tokens = word_tokens[start_idx: end_idx]
            for i in range(len(affected_tokens) - 1):
                old_pair = (affected_tokens[i], affected_tokens[i + 1])
                diff_pairs[old_pair] -= wc 
                if old_pair not in unaffected_pairs:   
                    pair_to_words[old_pair].discard(word)

For our example, assuming wc is 10, the code will execute:

diff_pairs[('c','s')] = -10
pair_to_words[('c','s')].discard(word)

For ('e','f'), pair_to_words[('e','f')].discard(word) will not be executed because ('e','f') is in unaffected_pairs.

Next, we calculate the new tokens:

            new_tokens = []
            all_new_tokens = []
            for i in range(start_idx):
                all_new_tokens.append(word_tokens[i])
            
            i = 0

            # account for multiple occurrences of the pair
            while i < len(affected_tokens):
                if i < len(affected_tokens) - 1 and (affected_tokens[i], affected_tokens[i + 1]) == merge_pair:
                    new_tokens.append(new_id)
                    all_new_tokens.append(new_id)
                    # jump past pair
                    i += 2
                else:
                    new_tokens.append(affected_tokens[i])
                    all_new_tokens.append(affected_tokens[i])
                    i += 1
            

            for i in range(end_idx, len(word_tokens)):
                all_new_tokens.append(word_tokens[i])

Note: Previously we only had new_tokens, but now we also have all_new_tokens. new_tokens is for the affected part, while all_new_tokens is for the entire word. So for our "abcststefstbbef" example, the results are:

new_tokens = ['c','st','st','e','f','st','b']
all_new_tokens = ['a','b'] + new_tokens + ['b','e','f']

Finally, we calculate the frequencies of the new pairs:

            word_encodings[word] = all_new_tokens

            # add new pairs from the updated word
            for i in range(len(new_tokens) - 1):
                new_pair = (new_tokens[i], new_tokens[i + 1])

                diff_pairs[new_pair] += wc
                pair_to_words[new_pair].add(word)

                new_pairs.add(new_pair)

The first line updates word_encodings with all_new_tokens. For the example above:

word_encodings['abcststefstbbef'] = ['a','b', 'c','st','st','e','f','st','b', 'b','e','f']

Then it finds all the new pairs in new_tokens:

('c','st')
('st','st')
('st','e')
('e','f')
('f','st')
('st','b')

Here, “new” pairs don’t necessarily mean they didn’t exist before. For example, ('e','f') did exist, but it was deleted and added back, which could result in diff_pairs[('e','f')]==0. This is why the code checks if diff_pairs is zero; if it is, it means there was no change. However, the code above will first delete the ('e','f') pair with diff_pairs[old_pair] -= wc and then add it back later. But this only happens if a pair appears more than twice in a single word. If a pair appears only once, like in the word 'abcstbbef', the deleted pairs are only [('c','s'), ('s','t'), ('t','b')], and it’s impossible for them to be added back.

4. Testing

The test code is in bpe_v5_time.py.

Version	Data	Total Time (s)	Word Freq Time (s)	Merge Time (s)	Other
bpe_v1_time	tinystory	2187/2264/2247	622/642/628	count_pair:944/995/984 max:167/173/174/174 update:453/453/460
bpe_v2_time	tinystory	757/738/746	639/621/627	118/117/118
bpe_v2_time	openweb	35358/34265/35687	2870/2949/2930	32437/31264/32708
bpe_v3_time	tinystory	250/90/90	80/90/90	170/0/0	num_counter=8, num_merger=1 skip
bpe_v3_time	openweb	391/401/32352	390/400/410	total:31883 max:31187 update:695	num_counter=8, num_merger=1 skip
bpe_v5_time	tinystory	221/236/222	90/90/90	total:130/146/132 max:127/143/129 update:3/3/3	num_counter=8, num_merger=1
bpe_v5_time	openweb	34333/34853/35804	401/390	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1

Compared to bpe_v3_time, using fine-grained updates on openweb reduces the time from 695 seconds to a little over 500 seconds.

5. Porting to C++

Similarly, I ported this fine-grained update algorithm to the C++ versions, including bpe_train_updater_fine_grained, bpe_train_updater_fine_grained_absl, and bpe_train_updater_fine_grained_emhash8.

I won’t go into detail about the code, but interested readers can compare the C++ versions with the Python version.

The experimental results for these versions are:

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v7	Boost hash	1268/1196/1215	548/473/481	719/723/734	export OMP_NUM_THREADS=16 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2_hash	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270
bpe_train_updater_omp_v2_hash2	Absl hash	1170/1074/1071	545/456/449	625/617/621
bpe_train_updater_opt_absl	Absl hash	1072/1012/1022	423/378/384	648/633/637
bpe_train_updater_emhash8	Boost hash	479/485/485	398/401/401	80/83/83
bpe_train_updater_opt_emhash8	Boost hash	469/474/479	389/395/399	79/78/79
bpe_train_updater_opt_emhash8_hash	my hash	2316/1951/1983	2250/1888/1918	66/63/64
bpe_train_updater_fine_grained	Boost Hash	8773/8873/7641	220/219/233	8552/8653/7408
bpe_train_updater_fine_grained_absl	Absl hash	845/845/856	204/201/203	641/643/653
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60

After implementing the fine-grained update, the total time for bpe_train_updater_fine_grained_emhash8 is only 260 seconds, with an update time of just 200 seconds. In contrast, bpe_train_updater_emhash8 had a total time of 480 seconds, an update time of nearly 400 seconds, and a max time of 79 seconds. This shows that the fine-grained update reduces unnecessary deletions and insertions, cutting the update time by half. These unnecessary operations also tend to make the data structure messy, which is why the max time was slightly reduced as well.

6. Using `emhash8::HashSet`/`emhash9::HashSet` to Replace `std::unordered_set`

Besides pair_counts, another frequently updated data structure is the inverted index pair_wordids:

std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash> pair_wordids;

Here, in addition to the outer std::unordered_map, the value itself is an std::unordered_set<int>. We can replace std::unordered_set with emhash8::HashSet or emhash9::HashSet. The complete code is available in bpe_train_updater_fine_grained_emhash8_set.cpp and bpe_train_updater_fine_grained_emhash8_set9.cpp.

Using hash_set4.hpp (emhash9::HashSet) and hash_table8.hpp (emhash8::HashMap) together can cause some compilation warnings. However, according to this issue, we can ignore them. The experimental results are as follows:

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v7	Boost hash	1268/1196/1215	548/473/481	719/723/734	export OMP_NUM_THREADS=16 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2_hash	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270
bpe_train_updater_omp_v2_hash2	Absl hash	1170/1074/1071	545/456/449	625/617/621
bpe_train_updater_opt_absl	Absl hash	1072/1012/1022	423/378/384	648/633/637
bpe_train_updater_emhash8	Boost hash	479/485/485	398/401/401	80/83/83
bpe_train_updater_opt_emhash8	Boost hash	469/474/479	389/395/399	79/78/79
bpe_train_updater_opt_emhash8_hash	my hash	2316/1951/1983	2250/1888/1918	66/63/64
bpe_train_updater_fine_grained	Boost Hash	8773/8873/7641	220/219/233	8552/8653/7408
bpe_train_updater_fine_grained_absl	Absl hash	845/845/856	204/201/203	641/643/653
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8_set	Boost Hash	192/192/194	117/117/117	75/75/77
bpe_train_updater_fine_grained_emhash8_set9	Boost Hash	168/170/171	107/108/109	61/62/61

As you can see, by replacing pair_wordids with the faster emhash8::HashMap/emhash8::HashSet/emhash9::HashSet, the time was further reduced from over 260 seconds to around 170 seconds.

Full Series

Part 0: Introduction Introduces the basic BPE training algorithm and related tasks, as well as the development environment.
Part 1: The Simplest Implementation The simplest implementation of BPE training.
Part 2: Optimized Algorithm Implements incremental updates for pair_counts.
Part 3: Parallel Tokenization and Frequency Counting Uses multiprocessing to implement a multi-process parallel algorithm.
Part 4: A Failed Parallel Optimization An attempt to parallelize the max pair calculation using multiple processes.
Part 5: Implementing the Merge Algorithm in C++ Implements a C++ merge algorithm equivalent to the Python version, and compares two ways of iterating through std::unordered_map.
Part 6: Parallelizing the Max Pair Search with OpenMP Uses OpenMP to find the max pair in pair_counts in parallel.
Part 7: Using Flat Hashmap to Replace std::unordered_map Uses flat hashmap to replace std::unordered_map.
Part 8: Implementing Fine-Grained Updates Implements a fine-grained update algorithm for pair_counts using an inverted index.
Part 9: Using a Heap to Find the Max Pair Uses a heap to find the max pair and improve performance.
Part 10: Using Cython and PyPy for Acceleration Uses Cython and PyPy to accelerate Python code.
Part 11: Wrapping C++ Code with Cython Wraps C++ code using Cython.

动手实现和优化BPE Tokenizer的训练——第8部分：实现细粒度更新

Fri, 19 Sep 2025 00:00:00 +0000

本系列文章完成Stanford CS336作业1的一个子任务——实现BPE Tokenizer的高效训练算法。通过一系列优化，我们的算法在OpenWebText上的训练时间从最初的10多个小时优化到小于10分钟。本系列文章解释这一系列优化过程，包括：算法的优化，数据结构的优化，并行(openmp)优化，cython优化，用c++实现关键代码和c++库的cython集成等内容。本文是第九篇，优化pair_counts等数据结构的更新过程。

1. 问题分析
2. 细粒度的频率更新算法
3. 代码实现
4. 测试
5. 移植到c++
6. 使用emhash8::HashSet/emhash9::HashSet替代std::unordered_set
本系列全部文章

1. 问题分析

我们之前遗留一个优化点，那就是在更新pair_counts时很多pair的计数本来是不需要变的，但是为了实现简单，我们先删除再加入。更具体的，我们来看一下之前的例子。

假设词频统计的结果为：

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

然后统计pair的频率：

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 11, st: 11, ne: 6, ew: 6}

然后根据合并规则，我们选择合并(‘s’, ‘t’)。由于它们发生了合并，我们需要更新pair的频率统计。先看看目前我们的更新算法(代码在BPE_Trainer._updated_affected_word_count，其实更准确的名字应该叫_updated_pair_count_of_affected_word)。

为了获得新的pair的频率，我们根据pair_to_words找到包含(‘s’, ‘t’)这个pair的word有：

'widest'
'newest'

我们首先把这两个词里的pair都先减掉，也就是减去：

{wi: 3, id: 3, de: 3, es: 3, st: 3}
{ne: 6, ew: 6, we: 6, es: 6, st: 6}

这样得到：

{lo: 7, ow: 7, we: 2, er: 2, wi: 0, id: 0, de: 0, es: 2, st: 2, ne: 0, ew: 0}

接着这两个词的tokens变成了：

word_encodings['widest'] = ['w','i','d','e','st']
word_encodings['newest'] = ['n','e','s','e','st']

因此我们根据新的word_encodings计算新的pair的频次：

{wi: 3, id: 3, de: 3, est: 3}
{ne: 6, ew: 6, we: 6, est: 6}

把这些频次加回去得到：

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 2, st: 2, ne: 6, ew: 6, est: 9}

可以发现’wi’,’id’和’de’这些pair的频次被减掉又加回去，这就带来没有必要的更新，更糟糕的是，比如’wi’删除后频率为0，那么就需要从pair_counts和pair_to_words里删除。

                    if pair_counts[old_pair] <= 0:
                        # we accounted for all occurrences of this pair
                        del pair_counts[old_pair]
                        pair_to_words.pop(old_pair)

根据Internals of sets and dicts，CPython的dict使用的是开放地址的冲突解决方法。从dict里删除一个key只是做一个记号(除非rehash)，然后再增加同样的key会复用原来的地址。所以对dict本身的影响不大，但是key和value还是需要重新构造对象(dict里保存的是key和value的指针)。如果是std::unordered_map这样的分离链接（Separate Chaining）方法，那么删除再插入就会把链表原来的顺序打乱(链表是插入在表头)，比如某个桶原来的元素是a->b->c，现在把b删除再插入，新的链表就可能变成b->a->c。

但是如果我们仔细分析，我们就能知道哪些删除是不必要的。比如我们的例子，我们合并的是(‘s’,’t’)，而在’widest’，受到合并影响的除了(‘s’,’t’)之外，只有这个pair之前和之后的token，因此’widest’里的pair受到影响的只有‘es’,’st’，而[‘wi’,’id’,’de’]不受影响，因此我们只需要删除’es’,’st’，然后增加’est’即可。

2. 细粒度的频率更新算法

根据前面的分析，我们可以实现细粒度的频率更新算法。前面分析的例子是一个要合并的pair在word里只出现一次的情况。但是如果要合并的pair出现多次呢？一种办法是fallback到之前的算法。因为一个word重复出现pair的可能性不大。如果只出现一次，那么可以比较简单的处理；如果出现多次，则fallback到之前的粗粒度的更新算法——先删除再插入。不过我在写文章之前没有想到这种方法，而是把问题搞得比较复杂。弄了一个复杂的算法能够处理pair在一个word里重复出现的算法。虽然这个算法看起来比较完美，但是实现的难度大了很多，而且收益也没有那么大。不过既然已经实现了，那么就继续吧。

比如一个假想的词”abcststefstbbef”，我们首先可以找到所有出现st的下标，为了看起来简单，我们用空格把st与其它token分开。

abc st st ef st bbef

根据前面的分析，第一个st之前的pair(‘ab’,’bc’)不受影响，最后一个st之后的pair(‘bb’)也不受影响。那么中间的pair呢？我们先写出来：

cs st ts st te ef fs st tb

这些pair首先得删除，然后我们把受到影响的部分生成新的token：

c st st e f st b

我们再加入新的pair：

('c','st'), ('st','st'), ('st','e'), ('e','f'), ('f','st'), ('st','b')

除了减少受影响的老pair的频次计数，增加新pair的计数，我们还需要修改倒排索引pair_to_words。我最早写的代码是把中间受影响的pair从中删除。但是这是有bug的。比如ef会受到影响，因此需要减少计数。但是ef除了在3个st的中间出现，它还出现在最后。如果我们直接从pair_to_words里删除ef，那么倒排索引就和实际不一致了。”真正”的倒排索引不仅仅记录pair(‘e’,’f’)在词’abcststefstbbef’里出现了，而且还需要记录出现的次数。也就是说pair_to_words的key是pair，value也是dict，value这个dict是word/pair_freq。而我们现在的pair_to_words的value是一个set。因此我们只知道pair(‘e’,’f’)在词’abcststefstbbef’里出现了，但是不知道出现几次。这就出现问题了。比如上面的例子我们之间从pair_to_words[(‘e’,’f’)]里删除词’abcststefstbbef’，那么倒排索引就不对了，但是我们又必须从pair_to_words[(‘t’,’e’)]的value里删除’abcststefstbbef’，因为在单词’abcststefstbbef’里pair(‘t’,’e’)只出现一次。

一种解决办法就是修改pair_to_words相关代码，把它变成{pair: {word: word_freq}}这样的结构。这是比较清晰的改动方法。不过我当时觉得这么改需要修改的地方太多，我想把修改局限在一个函数里。因此我使用了另外一个办法(现在仔细思考其实并不太好)。这个办法是：找到不受影响的所有pair，把它放到一个set里。在修改倒排索引时如果发现要删除的pair在不受影响的pair也出现，那么不删除。从逻辑的角度这个算法也是正确的，但是它使得代码变得复杂。

不过总体来说扫描一遍不受影响的pair速度也比较快，而且后面的很多改进版本都基于当前算法，要改动比较麻烦。所以在写这篇文章的时候我就暂时不修改了。

3. 代码实现

完整的代码在bpe_v5.py。我们看一下改动的地方。首先是对_updated_affected_word_count的修改：

    @staticmethod
    def _updated_affected_word_count(merge_pair, affected_words, word_encodings, 
                                     word_counts, pair_counts, pair_to_words, 
                                     new_id, pair_strings, vocabulary):
        # we may update/delete words when iterate it.
        affected_words = affected_words.copy()
        diff_pairs = defaultdict(int)

        new_pairs = set() 
        BPE_Trainer.fine_grained_pair_counter_diff(affected_words, word_encodings, word_counts, merge_pair, diff_pairs, 
                             new_id, pair_to_words, new_pairs)
        for pair, count in diff_pairs.items():
            if count == 0: continue
            pair_counts[pair] += count
            if pair_counts[pair] <= 0: # should not less than 0!
                del pair_counts[pair]
                pair_to_words.pop(pair, None)


        for new_pair in new_pairs:
            if new_pair not in pair_strings:
                pair_strings[new_pair] = (vocabulary[new_pair[0]], vocabulary[new_pair[1]])

它调用一个新的函数BPE_Trainer._updated_affected_word_count变得非常简单，它调用一个新的函数BPE_Trainer.fine_grained_pair_counter_diff来计算合并pair后pair计数的变化。这个函数的主要输出是diff_pairs，比如假想的一个输出：

('a','s'): -3
('t','f'): -5
('st','b'):10

则它说明由于合并(‘s’,’t’)，导致老的pair(‘a’,’s’)的计数减少了3，(‘t’,’f’)减少了5；而新的pair(‘st’,’b’)增加了10。所以下面的for循环就用这个diff去更新pair_counts。另外BPE_Trainer.fine_grained_pair_counter_diff函数也会返回合并后新增的new_pairs，这些新出现的pair，我们也需要更新pair_strings。下面我们重点来看BPE_Trainer.fine_grained_pair_counter_diff函数。

    @staticmethod
    def fine_grained_pair_counter_diff(affected_words, word_encodings, word_counts, merge_pair, diff_pairs, new_id, pair_to_words, new_pairs):
        for word in affected_words:
            word_tokens = word_encodings[word]
            wc = word_counts[word]

            # find first and last pairs
            idx = 0
            unaffected_pairs = set()
            while idx < len(word_tokens) - 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    first_idx = idx
                    break
                idx += 1
            else:
                print(f"bug {merge_pair}, {word}, {word_tokens}")
                raise
            # assert first_idx exists

            idx = len(word_tokens) - 2
            while idx > first_idx + 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    last_idx = idx
                    break
                idx -= 1
            else:
                last_idx = first_idx

            start_idx = max(0, first_idx - 1) # inclusive
            end_idx = min(last_idx + 3, len(word_tokens)) # exclusive

            # unaffected [0, start_idx)
            
            for i in range(start_idx):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair)
            # unaffected [end_idx-1, :-1]
            for i in range(end_idx - 1, len(word_tokens) - 1):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair)                

            affected_tokens = word_tokens[start_idx: end_idx]
            for i in range(len(affected_tokens) - 1):
                old_pair = (affected_tokens[i], affected_tokens[i + 1])
                diff_pairs[old_pair] -= wc 
                if old_pair not in unaffected_pairs:   
                    pair_to_words[old_pair].discard(word)
                    

            new_tokens = []
            all_new_tokens = []
            for i in range(start_idx):
                all_new_tokens.append(word_tokens[i])
            
            i = 0

            # account for multiple occurrences of the pair
            while i < len(affected_tokens):
                if i < len(affected_tokens) - 1 and (affected_tokens[i], affected_tokens[i + 1]) == merge_pair:
                    new_tokens.append(new_id)
                    all_new_tokens.append(new_id)
                    # jump past pair
                    i += 2
                else:
                    new_tokens.append(affected_tokens[i])
                    all_new_tokens.append(affected_tokens[i])
                    i += 1
            

            for i in range(end_idx, len(word_tokens)):
                all_new_tokens.append(word_tokens[i])
            
            word_encodings[word] = all_new_tokens

            # add new pairs from the updated word
            for i in range(len(new_tokens) - 1):
                new_pair = (new_tokens[i], new_tokens[i + 1])

                diff_pairs[new_pair] += wc
                pair_to_words[new_pair].add(word)

                new_pairs.add(new_pair)

这段代码比较复杂，尤其是下标的处理。我们分段来仔细看一下。

找到合并pair的第一次和最后一次出现

我们首先需要找到pair的第一次和最后一次出现的位置，如果只有一次，那么这两个位置是相同的。pair第一次出现的第一个元素和第二次出现第二个元素就是我们关注的范围。

            idx = 0
            unaffected_pairs = set()
            while idx < len(word_tokens) - 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    first_idx = idx
                    break
                idx += 1
            else:
                print(f"bug {merge_pair}, {word}, {word_tokens}")
                raise
            # assert first_idx exists

            idx = len(word_tokens) - 2
            while idx > first_idx + 1:
                if word_tokens[idx] == merge_pair[0] and word_tokens[idx+1] == merge_pair[1]:
                    last_idx = idx
                    break
                idx -= 1
            else:
                last_idx = first_idx

第一个while循环找到pair出现的第一次位置，记录在first_idx里。else是不应该运行的(除非我们的倒排索引出现问题)，这是防御性的抛出异常。第二个while循环反过来从最后的位置开始找，记录在last_idx里，如果找不到则说明这个pair只出现了一次，则last_idx==first_idx。

用之前的例子来分析，我们找到的first_idx是3，last_idx是9。如下所示：

"abcststefstbbef"
 012345678901234
    |     | 
first_idx |
       last_idx  

接下来是找到所有不受影响的pair：

            start_idx = max(0, first_idx - 1) # inclusive
            end_idx = min(last_idx + 3, len(word_tokens)) # exclusive

            # unaffected [0, start_idx) 
            for i in range(start_idx):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair)
                
            # unaffected [end_idx-1, :-1]
            for i in range(end_idx - 1, len(word_tokens) - 1):
                pair = word_tokens[i], word_tokens[i + 1]
                unaffected_pairs.add(pair) 

这里的下标计算特别要小心，我就不细讲了，总之最终的结果是找到如下不受影响的pair保存的unaffected_pairs里：

ab bc bb be ef

接下来就是减少受影响的pair的频次和更新倒排索引：

            affected_tokens = word_tokens[start_idx: end_idx]
            for i in range(len(affected_tokens) - 1):
                old_pair = (affected_tokens[i], affected_tokens[i + 1])
                diff_pairs[old_pair] -= wc 
                if old_pair not in unaffected_pairs:   
                    pair_to_words[old_pair].discard(word)

比如我们的例子，假设wc是10，则会执行：

diff_pairs[('c','s')] = -10
pair_to_words[('c','s')].discard(word)

而对于(‘e’,’f’)，pair_to_words[(‘e’,’f’)].discard(word)不会执行，因为(‘e’,’f’)在unaffected_pairs里。

接下来是计算新的tokens：

            new_tokens = []
            all_new_tokens = []
            for i in range(start_idx):
                all_new_tokens.append(word_tokens[i])
            
            i = 0

            # account for multiple occurrences of the pair
            while i < len(affected_tokens):
                if i < len(affected_tokens) - 1 and (affected_tokens[i], affected_tokens[i + 1]) == merge_pair:
                    new_tokens.append(new_id)
                    all_new_tokens.append(new_id)
                    # jump past pair
                    i += 2
                else:
                    new_tokens.append(affected_tokens[i])
                    all_new_tokens.append(affected_tokens[i])
                    i += 1
            

            for i in range(end_idx, len(word_tokens)):
                all_new_tokens.append(word_tokens[i])

注意：之前我们只有一个new_tokens，但是现在我们还有一个all_new_tokens。new_tokens是计算受影响的那部分，而all_new_tokens是整个word。因此对于我们之前”abcststefstbbef”的例子，它们的结果是：

new_tokens = ['c','st','st','e','f','st','b']
all_new_tokens = ['a','b'] + new_tokens + ['b','e','f']

最后是计算新的pair的频次：

            word_encodings[word] = all_new_tokens

            # add new pairs from the updated word
            for i in range(len(new_tokens) - 1):
                new_pair = (new_tokens[i], new_tokens[i + 1])

                diff_pairs[new_pair] += wc
                pair_to_words[new_pair].add(word)

                new_pairs.add(new_pair)

第一行是更新word_encodings为all_new_tokens。比如前面的例子：

word_encodings['abcststefstbbef'] = ['a','b', 'c','st','st','e','f','st','b', 'b','e','f']

接着是找到new_tokens里的所有新pair：

('c','st')
('st','st')
('st','e')
('e','f')
('f','st')
('st','b')

这里所谓的新并不一定就是之前没有出现的pair，比如(‘e’,’f’)之前其实就有，只不过被删除又加进来，因此可能出现diff_pairs[(‘e’,’f’)]==0的情况，这也是前面代码会判断diff_pairs是否为0，如果为0，则说明它并没有改变。但是上面的代码会在前面diff_pairs[old_pair] -= wc先删除(‘e’,’f’)这个pair，后面又再加进来。不过这种情况只有在一个word里包含pair两次以上才会出现。如果pair只出现一次，比如单词是’abcstbbef’，那么被删除的只有[(‘c’,’s’), (‘s’,’t’), (‘t’,’b’)]，这些如果被删除就不可能再被加进来。

4. 测试

测试的代码是bpe_v5_time.py。

版本	数据	总时间(s)	统计词频时间(s)	合并时间(s)	其它
bpe_v1_time	tinystory	2187/2264/2247	622/642/628	count_pair:944/995/984 max:167/173/174/174 update:453/453/460
bpe_v2_time	tinystory	757/738/746	639/621/627	118/117/118
bpe_v2_time	openweb	35358/34265/35687	2870/2949/2930	32437/31264/32708
bpe_v3_time	tinystory	250/90/90	80/90/90	170/0/0	num_counter=8, num_merger=1 skip
bpe_v3_time	openweb	391/401/32352	390/400/410	total:31883 max:31187 update:695	num_counter=8, num_merger=1 skip
bpe_v5_time	tinystory	221/236/222	90/90/90	total:130/146/132 max:127/143/129 update:3/3/3	num_counter=8, num_merger=1
bpe_v5_time	openweb	34333/34853/35804	401/390/401	total:33879/34401/35347 max:33353/33820/34816 update:525/579/530	num_counter=8, num_merger=1

和bpe_v3_time相比，使用细粒度的更新在openweb上的时间从695秒下降到500多秒。

5. 移植到c++

同样的，我把这个细粒度更新的算法移植到了c++的版本，包括bpe_train_updater_fine_grained、bpe_train_updater_fine_grained_absl和bpe_train_updater_fine_grained_emhash8。

代码就不详细介绍，感兴趣的读者可以对照Python的版本阅读c++版本的代码。

这些版本的实验结果为：

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v7	Boost hash	1268/1196/1215	548/473/481	719/723/734	export OMP_NUM_THREADS=16 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2_hash	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270
bpe_train_updater_omp_v2_hash2	Absl hash	1170/1074/1071	545/456/449	625/617/621
bpe_train_updater_opt_absl	Absl hash	1072/1012/1022	423/378/384	648/633/637
bpe_train_updater_emhash8	Boost hash	479/485/485	398/401/401	80/83/83
bpe_train_updater_opt_emhash8	Boost hash	469/474/479	389/395/399	79/78/79
bpe_train_updater_opt_emhash8_hash	my hash	2316/1951/1983	2250/1888/1918	66/63/64
bpe_train_updater_fine_grained	Boost Hash	8773/8873/7641	220/219/233	8552/8653/7408
bpe_train_updater_fine_grained_absl	Absl hash	845/845/856	204/201/203	641/643/653
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60

使用了细粒度更新后，bpe_train_updater_fine_grained_emhash8的总时间只有260秒，其中更新时间只有200秒。对比bpe_train_updater_emhash8，它的总时间是480秒，update时间接近400秒，max时间79秒。这说明细粒度更新减少了不必要的删除和插入，因此update时间变为原来的一半，而且这些不必要的删除和插入也会让数据结构变得混乱，因此max的时间也有少量减少。

6. 使用emhash8::HashSet/emhash9::HashSet替代std::unordered_set

除了pair_counts，还有一个比较频繁更新的就是倒排索引pair_wordids:

std::unordered_map<std::pair<int, int>, std::unordered_set<int>, pair_hash> pair_wordids;

这里除了外面的std::unordered_map，value本身也是一个std::unordered_set。我们可以用emhash8::HashSet/emhash9::HashSet替代std::unordered_set。完整的代码在[bpe_train_updater_fine_grained_emhash8_set.cpp](https://github.com/fancyerii/assignment1-basics-bpe/blob/main/cppupdate/bpe_train_updater_fine_grained_emhash8_set.cpp)和[bpe_train_updater_fine_grained_emhash8_set9.cpp](https://github.com/fancyerii/assignment1-basics-bpe/blob/main/cppupdate/bpe_train_updater_fine_grained_emhash8_set9.cpp)。

同时使用hash_set4.hpp(emhash9::HashSet) and hash_table8.hpp(emhash8::HashMap)在编译时会出现一些警告，不过根据这个issue，我们可以忽略它们。实验结果如下：

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v7	Boost hash	1268/1196/1215	548/473/481	719/723/734	export OMP_NUM_THREADS=16 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2_hash	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270
bpe_train_updater_omp_v2_hash2	Absl hash	1170/1074/1071	545/456/449	625/617/621
bpe_train_updater_opt_absl	Absl hash	1072/1012/1022	423/378/384	648/633/637
bpe_train_updater_emhash8	Boost hash	479/485/485	398/401/401	80/83/83
bpe_train_updater_opt_emhash8	Boost hash	469/474/479	389/395/399	79/78/79
bpe_train_updater_opt_emhash8_hash	my hash	2316/1951/1983	2250/1888/1918	66/63/64
bpe_train_updater_fine_grained	Boost Hash	8773/8873/7641	220/219/233	8552/8653/7408
bpe_train_updater_fine_grained_absl	Absl hash	845/845/856	204/201/203	641/643/653
bpe_train_updater_fine_grained_emhash8	Boost Hash	261/259/261	200/198/200	61/60/60
bpe_train_updater_fine_grained_emhash8_set	Boost Hash	192/192/194	117/117/117	75/75/77
bpe_train_updater_fine_grained_emhash8_set9	Boost Hash	168/170/171	107/108/109	61/62/61

可以看到pair_wordids用更快的emhash8::HashMap/emhash8::HashSet/emhash9::HashSet替换会后，时间进一步从260多秒降到了170秒。

本系列全部文章

第0部分：简介介绍bpe训练的基本算法和相关任务，并且介绍开发环境。
第1部分：最简单实现 bpe训练最简单的实现。
第2部分：优化算法实现pair_counts的增量更新。
第3部分：并行分词和统计词频使用multiprocessing实现多进程并行算法。
第4部分：一次失败的并行优化尝试用多进程并行计算max pair。
第5部分：用C++实现Merge算法用C++实现和Python等价的merge算法，并且比较std::unordered_map的两种遍历方式。
第6部分：用OpenMP实现并行求最大用OpenMP并行求pair_counts里最大pair。
第7部分：使用flat hashmap替代std::unordered_map 使用flat hashmap来替代std::unordered_map。
第8部分：实现细粒度更新使用倒排索引实现pair_counts的细粒度更新算法。
第9部分：使用堆来寻找最大pair 使用堆来求最大pair，提升性能。
第10部分：使用cython和pypy来加速使用cython和pypy来加速python代码。
第11部分：使用cython封装c++代码使用cython封装c++代码。

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 7: Using Flat Hash Map instead of std::unordered_map

Thu, 18 Sep 2025 00:00:00 +0000

This series of articles implements a subtask of Stanford’s CS336 Assignment 1: building an efficient training algorithm for a BPE Tokenizer. Through a series of optimizations, our algorithm’s training time on OpenWebText was reduced from over 10 hours to less than 10 minutes. This series explains these optimizations, including algorithmic improvements, data structure enhancements, parallelization with OpenMP, Cython optimization, and implementing key code in C++ along with its integration via Cython. This is the eighth article in the series, focusing on using a flat hash map to replace the C++ standard library’s std::unordered_map for improved performance.

Table of Content

1. Problem Analysis
2. Optimizing the -= Operation
3. Flat Hash Map Survey
4. absl::flat_hash_map
5. absl::flat_hash_map Test
6. emhash8::HashMap
7. emhash8::HashMap Test Results
8. Summary
Full Series

1. Problem Analysis

In the previous article, we used OpenMP parallelization to reduce the time for the second merge step from over 6000 seconds to less than 1000 seconds with 32 threads (bpe_train_updater_omp_v7). Out of that 1000 seconds, over 500 seconds were spent on updating data like pair_counts, and another 400+ seconds were for finding the max pair.

Before continuing the optimization, let’s review the main data read/write operations in the function bpe_train_step2. We’ll use the serial version of the code, bpe_train_updater_omp_v3.cpp, as an example.

The main code for finding the max pair iterates through pair_counts. It only needs to read pair_strings[pair] when the count of the current pair is the same as the current maximum.

The code for updating variables like pair_counts is as follows:

        const int token_size = word_tokens.size();
        for(int i = 0; i < token_size - 1; ++i){
            std::pair<int, int> old_pair(word_tokens[i], word_tokens[i + 1]);
            pair_counts[old_pair] -= wc;
            if(pair_counts[old_pair] <= 0){
                pair_counts.erase(old_pair);
                pair_wordids.erase(old_pair);
            }else{
                pair_wordids[old_pair].erase(wordid);
            }
        }
        
        
        for(int i = 0; i < new_tokens_size - 1; ++i){
            std::pair<int, int> new_pair(new_tokens[i], new_tokens[i + 1]);
            pair_counts[new_pair] += wc;
            pair_wordids[new_pair].insert(wordid);
            if (pair_strings.find(new_pair) == pair_strings.end()) {
                pair_strings[new_pair] = {vocabulary[new_pair.first], vocabulary[new_pair.second]};
            }
        }

The most critical operations here are the write operations: pair_counts[old_pair] -= wc; and pair_counts[new_pair] += wc;. Additionally, when pair_counts[old_pair] <= 0, the code deletes old_pair from both pair_counts and pair_wordids.

There is also a read operation if(pair_counts[old_pair] <= 0). This operation can actually be optimized away, which we’ll discuss later. But even without optimization, the time it takes is very short. First, unordered_map lookups are fast (updates require memory allocation, insertion into a linked list and bucket, and may cause a rehash). Second, this lookup happens right after an update, so the relevant data should still be in the CPU’s cache.

First, let’s look at iterating through pair_counts. In the serial version, we iterate through the unordered_map’s internal singly-linked lists. In the parallel version, we also traverse the linked list contents, but we do so in chunks using the bucket interface. From an implementation perspective, unordered_map is a linked list, which allows for O(1) time complexity for element insertion and deletion. However, our main write operation is updating the value corresponding to a key. The memory layout of a linked list is very fragmented, and its traversal is much slower than a contiguous memory layout. Therefore, for our scenario, we can use a hash map with a contiguous layout, a so-called flat hash map. Its main feature is that it stores all data in a single, contiguous block of memory, unlike a standard hash map which disperses data using linked lists or pointers.

A standard hash map (e.g., C++’s std::unordered_map) typically uses a hash function to map keys to different “buckets,” where each bucket might be a linked list to store elements with the same hash value. This results in elements being scattered and non-contiguous in memory.

A flat hash map, on the other hand, is completely different. Its core principles are:

Contiguous Memory Layout: It uses one large array to store all key-value pairs. Since the data is stored contiguously, this significantly improves Cache-Friendliness. When a processor accesses one element, it usually loads nearby elements into the cache as well, making subsequent accesses much faster.
Open Addressing: To handle hash collisions (when two different keys hash to the same location), a flat hash map does not use linked lists. Instead, it employs a technique called Probing. When a position is already occupied, it follows a predefined rule (e.g., linear or quadratic probing) to find the next available empty slot for the data.

Pros and Cons

Pros
- Excellent Cache Performance: This is its biggest advantage. In modern CPUs, cache access is much faster than main memory access. Because data is contiguous, a flat hash map’s access pattern is ideal for the CPU cache, especially for traversal or bulk lookups.
- Low Memory Overhead: Since it doesn’t need to store extra pointers for each element, its memory footprint is smaller, especially when storing a large number of small elements.
- Ideal for High-Performance Scenarios: In fields with extremely low latency requirements, such as game development, network programming, and high-performance computing, the flat hash map is a great choice.
Cons
- Complex Deletion: Due to the probing mechanism, elements cannot simply be removed. A special “tombstone” marker is needed to ensure that subsequent lookups in the chain are not broken. This can lead to fragmentation in the array and requires periodic cleanup (re-hashing).
- Clustering Issues: If the hash function is not ideal or if collisions are frequent, data can cluster together in the array, forming “blocks.” This makes the probing process longer and can degrade performance.
- High Rehash Cost: When the underlying array needs to be resized, a larger new array must be allocated and all elements copied over, a process with a relatively high cost.

However, for our specific scenario, a flat hash map is very suitable: our main operation is traversal, and deletions are infrequent.

2. Optimizing the `-=` Operation

Before replacing std::unordered_map with a flat hash map, let’s do a small optimization to save one query to pair_counts.

In C++, unordered_map’s -= is an expression that returns the value after the operation; whereas in Python, a dict’s -= is a statement with no return value. For example, the following code compiles in C++:

int x = 1;
int y = (x += 1);

But similar code in Python:

x = 1
y = (x += 1)

will cause a syntax error:

    y = (x += 1)
           ^
SyntaxError: invalid syntax

So in Python, we can only query pair_counts[old_pair] again after the update pair_counts[old_pair] -= wc;. But in C++, we can directly get the value after the -= operation.

We can apply this optimization to all versions (e.g., bpe_train_updater_omp_v3/bpe_train_updater_omp_v5/bpe_train_updater_omp_v7, etc.), but since we only discovered it after the last version, we’ll first compare it with the original bpe_train_updater.cpp. If it proves effective, we’ll apply it to the final version. The complete code is in bpe_train_updater_opt.cpp, with only one change:

            int& c = (pair_counts[old_pair] -= wc);
            if(c <= 0){
                pair_counts.erase(old_pair);
                pair_wordids.erase(old_pair);
            }else{
                pair_wordids[old_pair].erase(wordid);
            }

Here’s the test comparison:

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_opt	Boost hash	8153/7634/7362	411/396/389	7741/7238/6973

Comparing the update times, we see an average decrease from 450s to 398.67s, an improvement of 11.41%. While this may not be a huge gain for the overall runtime, every little bit helps.

3. Flat Hash Map Survey

We’ve already introduced the principles of the flat hash map and concluded it’s a good fit for our scenario. Of course, we don’t need to reinvent the wheel, so the next step is to research available open-source flat hash map libraries. A search reveals many third-party libraries, and the article Comprehensive C++ Hashmap Benchmarks 2022 tests several C++ hash map libraries, many of which are flat.

Based on the summary of this article, I’ve tried absl::flat_hash_map and emhash8::HashMap.

4. absl::flat_hash_map

The first reason for trying this one is that many articles recommend it, and it’s open-source from Google. Google supposedly uses it internally, which is quite impressive.

absl::flat_hash_map is part of the abseil-cpp library. To use absl::flat_hash_map, you have to pull in a huge number of other dependencies, which seems to be a common issue with Google’s open-source libraries. You only want to use one small feature, but since it’s part of a larger library, you have to include a bunch of dependencies. Of course, this isn’t entirely Google’s fault; it’s a fundamental problem with C++. Unlike Python or Java, C++ lacks a cross-platform ABI, which means to use a C++ library, you generally have to recompile all of its dependencies. In the early days, C++ also lacked a central repository like Maven or PyPI (though things like Conan exist now, they’re not widely used), and compiling every library from scratch was a hassle. As a result, many people just used their own code or copied and modified others’ code, leading to fragmented branches that couldn’t be easily updated. A mature library has many dependencies. For example, you probably need logging for any code you write, so you need a dependency for that. You need to parse command-line arguments, another dependency. You need string manipulation, but the C++ standard library strings aren’t great, so you need to create your own or rely on one. You need regular expressions, XML/JSON/YAML parsing, etc. The C++ standard library lacks all of these. In contrast, in Python or Java, these features are either in the standard library or have converged on a few winning open-source libraries after several rounds of competition, leading to a more unified ecosystem. But C++ can’t converge, so every company has to build its own set of tools, and some companies even rebuild tools that the standard library already has because they don’t like them. This results in them open-sourcing almost an entire suite of libraries, making it impossible to separate a single small feature.

There are also some very small libraries, usually header-only, that you can just copy directly into your project. Even with these, because they don’t have a package manager like PyPI or Maven, once you use them, you basically can’t (or don’t dare to) upgrade them. This makes them only suitable for libraries that will never be modified after their release. But even small libraries need bug fixes. You can’t just diff and apply patches yourself.

Anyway, let’s get back to the main point. We first include absl::flat_hash_map in our code, which requires the full abseil-cpp code. We’ve already got the code via git submodule, so we just need to modify CMakeLists.txt:

add_subdirectory(abseil-cpp)

target_link_libraries(bpe_train_updater_omp_v2 PRIVATE absl::flat_hash_map absl::hash)

Then, include the header file where you need to use it:

#include "absl/container/flat_hash_map.h"

The absl::flat_hash_map interface is similar to std::unordered_map, but it lacks the bucket interface, so parallel traversal via buckets is not possible.

The full code using absl::flat_hash_map is in bpe_train_updater_omp_v2.cpp. The code is almost identical to bpe_train_updater.cpp, except that std::unordered_map has been replaced with absl::flat_hash_map.

Previously, compiling cppupdate took only tens of seconds. After introducing this dependency, it takes several minutes even with 8 threads.

5. `absl::flat_hash_map` Test

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270

We found that compared to bpe_train_updater, the total time for bpe_train_updater_omp_v2 dropped from around 8000 seconds to just over 2000 seconds. The max time is less than 300 seconds, which is more than 20 times faster than the 7000 seconds of bpe_train_updater and even faster than the 32-thread bpe_train_updater_omp_v7. This shows that the contiguous memory layout of the flat hash map is indeed well-suited for traversal. However, the update time increased from over 400 seconds to over 2000 seconds, which doesn’t seem to align with the benchmarks from other people online.

As we mentioned before, the hash function for a hash map is crucial. absl::flat_hash_map comes with its own hash function, which we need to test. We also have a very simple hash function that we implemented previously, which we can compare. The code using the simple hash is in bpe_train_updater_omp_v2_hash.cpp, and the code using the native absl hash is in bpe_train_updater_omp_v2_hash2.cpp. Additionally, we incorporated the -= optimization into bpe_train_updater_opt_absl.

Using the native absl hash is very simple; it’s used by default if you don’t specify a hash function:

absl::flat_hash_map<std::pair<int, int>, int> pair_counts;

Here are the test results:

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v7	Boost hash	1268/1196/1215	548/473/481	719/723/734	export OMP_NUM_THREADS=16 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2_hash	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270
bpe_train_updater_omp_v2	my hash	6276/6349/6613	6030/6104/6364	245/245/249
bpe_train_updater_omp_v2_hash2	Absl hash	1170/1074/1071	545/456/449	625/617/621
bpe_train_updater_opt_absl	Absl hash	1072/1012/1022	423/378/384	648/633/637

When using a simple hash (my hash), absl::flat_hash_map’s max time is still around 200 seconds, but the update time increases to over 6000 seconds. However, when using the native absl hash function, the max time increases to over 600 seconds (still far less than bpe_train_updater’s 7000+ seconds), but the update time is only around 500 seconds. The performance of this version is excellent, even surpassing the 16-thread speed of bpe_train_updater_omp_v7. The bpe_train_updater_opt_absl version, with the -= optimization, further reduces the update time to around 400 seconds.

6. emhash8::HashMap

You can refer to this issue for the principles behind emhash. I haven’t researched it in depth myself, so interested readers can explore it on their own or ask questions on the project’s GitHub page.

emhash is much lighter than absl::flat_hash_map. We only need to use emhash8, so we just need to include the header files. We can do this by adding a target_include_directories line to CMakeLists.txt:

target_include_directories(bpe_train_updater_emhash8 PUBLIC
    "${PROJECT_SOURCE_DIR}/emhash"
)

The change is also very simple: just replace std::unordered_map with emhash8::HashMap. The full code is in bpe_train_updater_emhash8.cpp.

7. `emhash8::HashMap` Test Results

For comparison, we also implemented an -= optimized version, bpe_train_updater_opt_emhash8.cpp, and an -= optimized version with a custom hash function, bpe_train_updater_opt_emhash8_hash.cpp.

Here are the results:

program	hash function	total time(sec)	update time(sec)	max time(sec)	other
bpe_train_updater	Boost hash	7171/7856/9248	392/480/478	6779/7376/8770
bpe_train_updater_omp_v7	Boost hash	907/908/955	514/503/554	391/403/400	export OMP_NUM_THREADS=32 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v7	Boost hash	1268/1196/1215	548/473/481	719/723/734	export OMP_NUM_THREADS=16 export OMP_SCHEDULE=”dynamic,1000”
bpe_train_updater_omp_v2_hash	Boost hash	2201/2392/2281	1931/2120/2010	269/272/270
bpe_train_updater_omp_v2	my hash	6276/6349/6613	6030/6104/6364	245/245/249
bpe_train_updater_omp_v2_hash2	Absl hash	1170/1074/1071	545/456/449	625/617/621
bpe_train_updater_opt_absl	Absl hash	1072/1012/1022	423/378/384	648/633/637
bpe_train_updater_emhash8	Boost hash	479/485/485	398/401/401	80/83/83
bpe_train_updater_opt_emhash8	Boost hash	469/474/479	389/395/399	79/78/79
bpe_train_updater_opt_emhash8_hash	my hash	2316/1951/1983	2250/1888/1918	66/63/64

The speed of emhash8::HashMap is even faster than absl::flat_hash_map, with a total time of less than 500 seconds. The max time is even less than 100 seconds. Using the -= optimization shaves off another ten seconds or so. At the same time, we see that performance also degrades if we use a simple hash function.

8. Summary

By using a flat hash map, the max time was drastically reduced, which proves that the flat, contiguous memory structure is indeed suitable for our scenario. As a result, we achieved a very high speed without even needing parallel optimization. The lesson here is that optimizing data structures and algorithms can sometimes yield better results than parallelization. We’ve accomplished the same task with fewer resources and at a faster speed.

With this, our work on optimizing the second step with C++ is temporarily concluded. Next, we will return to Python to continue our optimizations.

Full Series

Part 0: Introduction Introduces the basic BPE training algorithm and related tasks, as well as the development environment.
Part 1: The Simplest Implementation The simplest implementation of BPE training.
Part 2: Optimized Algorithm Implements incremental updates for pair_counts.
Part 3: Parallel Tokenization and Frequency Counting Uses multiprocessing to implement a multi-process parallel algorithm.
Part 4: A Failed Parallel Optimization An attempt to parallelize the max pair calculation using multiple processes.
Part 5: Implementing the Merge Algorithm in C++ Implements a C++ merge algorithm equivalent to the Python version, and compares two ways of iterating through std::unordered_map.
Part 6: Parallelizing the Max Pair Search with OpenMP Uses OpenMP to find the max pair in pair_counts in parallel.
Part 7: Using Flat Hashmap to Replace std::unordered_map Uses flat hashmap to replace std::unordered_map.
Part 8: Implementing Fine-Grained Updates Implements a fine-grained update algorithm for pair_counts using an inverted index.
Part 9: Using a Heap to Find the Max Pair Uses a heap to find the max pair and improve performance.
Part 10: Using Cython and PyPy for Acceleration Uses Cython and PyPy to accelerate Python code.
Part 11: Wrapping C++ Code with Cython Wraps C++ code using Cython.

李理的博客

《C++ High Performance: Master the art of optimizing the functioning of your C++ code , Second Edition》学习

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 11: Wrapping C++ Code with Cython

1. Goal

2. Wrapping the C++ Code into a Dynamic Library

2.1 Header File

2.2 CMakeLists.txt

2.3 bpe_train_step2

2.4 bpe_train_step2_v2

2.5 Compile and Install

3. Wrapping with Cython into an Extension Module

3.1 bpe_train_step2_wrapper.pxd

3.2 bpe_train_step2_wrapper.pyx

3.2.1 py_bpe_train_step2

3.2.2 py_bpe_train_step2_v2

3.3 Modifying setup.py

3.4 Using it in Python

4. Testing

5. Conclusion

动手实现和优化BPE Tokenizer的训练——第11部分：使用cython封装c++代码

1. 目标

2. 把c++代码封装成一个动态库

2.1 头文件

2.2 CMakeLists.txt

2.3 bpe_train_step2

2.4 bpe_train_step2_v2

2.5 编译安装

3. 用cython封装成扩展模块

3.1 bpe_train_step2_wrapper.pxd

3.2 bpe_train_step2_wrapper.pyx

3.2.1 py_bpe_train_step2

3.2.2 py_bpe_train_step2_v2

3.3 修改setup.py

3.4 在python里使用

4. 测试

5. 总结

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 10: Using Cython and PyPy for Acceleration

1. Problem Analysis

2. Introduction to Cython Principles

3. Optimizing the fine_grained_pair_counter_diff function with Cython

4. Test Results

5. Implementing the entire merge process with Cython

6. Using PyPy

7. Conclusion

Full Series

动手实现和优化BPE Tokenizer的训练——第10部分：使用cython和pypy来加速

1. 问题分析

2. cython原理简介

3. 用cython优化fine_grained_pair_counter_diff函数

4. 测试结果

5. 把整个merge过程用cython实现

6. 使用pypy

7. 总结

本系列全部文章

Implementing and Optimizing a BPE Tokenizer from Scratch—Part 9: Using a Heap to Find the Maximum Pair

1. Problem Analysis

2. Heap

3. Using a Heap to Find the Max

4. Python’s heapq Module

5. Implementing a Max Heap

6. Optimizing maxheap_py with a C Module

7. Using heapq to Implement a Max Heap

8. Performance Testing of Different Max Heap Implementations

9. Using a Max Heap to Find the Max

10. Testing

11. Experiments with a Larger vocab_size

12. Porting the Max Heap to C++

13. Optimizing _updated_affected_word_count

14. Porting the Optimization to C++

Full Series

动手实现和优化BPE Tokenizer的训练——第9部分：使用堆来寻找最大pair

1. 问题分析

2. 堆(heap)

3. 使用堆求max

4. Python的heapq模块

5. 实现一个大堆

6. 用c模块优化maxheap_py

7. 用heapq实现大堆

8. 不同的大堆实现的性能测试

9. 用大堆求max

3. Optimizing the `fine_grained_pair_counter_diff` function with Cython

4. Python’s `heapq` Module

6. Optimizing `maxheap_py` with a C Module

7. Using `heapq` to Implement a Max Heap

11. Experiments with a Larger `vocab_size`

13. Optimizing `_updated_affected_word_count`

6. Using `emhash8::HashSet`/`emhash9::HashSet` to Replace `std::unordered_set`

2. Optimizing the `-=` Operation

5. `absl::flat_hash_map` Test

7. `emhash8::HashMap` Test Results