Issues with tests/test_train_bpe.py::test_train_bpe_speed

[prateekgargX]

I have been trying to optimise my tokenizer using multiprocessing. After two days of debugging, I still see no speedup on this testcase, if anything it slows down. Without multiprocessing, it executes in 1.5 seconds and with multiprocessing it takes around 1.9-2.5 seconds based on the number of processes I use. But I see a lot of gains on TinyStories. I was able to train 10K vocab with wallclock time of 15 minutes. This was on Macbook Air M2 with 16gb RAM.

[marcelroed]

A couple of suggestions to get better performance with multiprocessing:

1. Minimize the amount of information communicated between processes. This means you should only input the filename and the byte range from this file to read to your process. When you return, you should not return all the words, just a dictionary with the local counts for the current process.
2. Minimize the number of processes you use. It's better to communicate something small once, then communicate when you're done. I've seen a few student solutions where processes are given one sequence at a time, where the time of your solution will be dominated by process overhead.
3. Once you've counted a token in your process-local pretokenization, free it from memory. For a memory efficient solution (which you will need to successfully train on OWT), you need to make sure you don't need to load the entire dataset at once. If you memmap your files and just add to your counter as tokens are encountered with find_iter, this should be quite easy.
4. Don't use multiprocessing for the token merge process, only for pretokenization and counting. There's a lot of contention for resources in token merging, so you should drop the multiprocessing here.

Hope this helps!

[prateekgargX]

thanks for the suggestions! but regarding bytes_dict not being local, I am using fresh dictionaries per process as seen in process_chunk_worker and then merging them in the main code:

    bytes_dict = {}
    with Pool(num_processes) as p:

        local_bytes_dicts = p.imap_unordered(
            worker_wrapper,
            [
                (input_path, special_tokens, ele)
                for ele in zip(boundaries[:-1], boundaries[1:])
            ],
        )
        for src_dict in local_bytes_dicts:
            for key, value in src_dict.items():
                bytes_dict[key] = value + bytes_dict.get(key, 0)

[marcelroed]

Ah, Github was hiding the worker function behind a scroll-bar. If the issue is not the communication overhead I suspect you might be running out of memory and your machine is falling back to swap.

[marcelroed]

How many processes are you using, and how many cores do you have?

[prateekgargX]

I have 8 cores (4 performance and 4 efficiency). I tried using 4 and 8 threads, didn't make much of a difference. Falling to back to swap is true because out of 16gb of RAM, 14gb is used by browser and tokenizer used 11gb of RAM

[marcelroed]

Yes, then you just need to reduce your memory usage. mmap and find_iter will get you there.