Globally Shuffling Hundreds of Gigabytes of Files with Limited Memory (Python)

有限内存下全局打乱几百G文件（Python） 苏剑林 · 2021-09-08 · Machine Learning programming Python

In this post we're going to tackle a programming problem:

How do you globally shuffle a text file that's hundreds of gigabytes in size, using only limited memory?

The motivation for this problem should be pretty clear: nowadays pretraining corpora routinely run into the tens or even hundreds of gigabytes, and to help models pretrain better, it's important to perform a global random shuffle of the training corpus. But for many people, a corpus of hundreds of gigabytes is often larger than available memory, so figuring out how to do a global shuffle under limited memory is a genuinely worthwhile problem to study.

Existing Tools

Suppose our file is stored line by line — that is, each line represents one sample — and what we want to do is randomly shuffle the file by line. If we only have a single file, and this file is clearly smaller than our available memory, then we can use Linux's built-in shuf command:

shuf input.txt -o output.txt

The reason I emphasize that the file size must be clearly smaller than memory is that shuf loads the entire file into memory before shuffling it, which requires us to have enough memory to hold it. To address this, an improved version called terashuf (see link) was created, which splits the file and uses disk space in place of memory, allowing us to shuffle files larger than our available memory.

At first glance it seems like terashuf should already fully satisfy our needs. In theory that's true, but sometimes we may have all sorts of customized requirements — for instance, mixing multiple files together and shuffling them jointly, or splitting the shuffled output into multiple files, and so on. So it's best if we can implement this ourselves in Python, so as to accommodate more complex custom requirements.

The Shuffling Algorithm

Now let's take a look at what the algorithm for global shuffling under limited memory actually looks like. Roughly, the steps are as follows:

1. Suppose the file has $mn$ lines in total; split it into $m$ files, each with $n$ lines.
2. Shuffle each $n$-line file internally at random. Since $n$ is arbitrary here, this step can be done entirely in memory.
3. Read the first line of each file (giving us $m$ lines of data), and write these $m$ lines to the output file in random order.
4. Read line $2,\cdots,n$ of each file in turn, and repeat step 3.

In plain terms, we first shuffle "vertically" once, and then shuffle "horizontally" once, which gives us a sufficiently well-mixed result, close to a global shuffle, as shown in the figure below:

Left: original data; middle: each column shuffled vertically within itself; right: on top of the vertical shuffle, each row shuffled horizontally

Note that this algorithm can only guarantee a result that is as thoroughly mixed as possible — it cannot guarantee that every possible ordering is equally likely to occur. For example, the first two samples of the output can never happen to be exactly the first two samples of the first file. To truly achieve a scheme where every ordering is equally probable, we would need to sample proportionally to each file's remaining number of lines at every single step, rather than reading line $k$ from every file simultaneously and then reading line $k+1$ from every file. But sampling proportionally at every single step increases the sampling cost, and in practice makes little difference to the final result, so it's not really necessary.

In practice, when splitting the file into chunks of $n$ lines each, the last file may end up with fewer than $n$ lines. If we don't care about this minor detail, we can just proceed with the process described above as usual: once the last file has been fully read, it simply keeps returning empty lines, and the whole process won't error out. Of course, this does mean that samples from the last file will be somewhat biased toward the beginning of the output — for readers with a touch of OCD who can't accept this, one option is to introduce rejection sampling when reading from the last file, with a rejection rate of $1-\frac{\text{last text remaining lines}}{\text{remaining lines of each remaining text}}$.

Reference implementation (without rejection sampling):

GitHub: https://github.com/bojone/shuffle

Performance Test

If all we need is a combination of merging, shuffling, and splitting, this can actually be achieved with shell commands plus terashuf, roughly as follows:

cat corpus/*.json | TMPDIR=/root/tmp MEMORY=20 ./terashuf | split -l 100000 -a 5 -d - corpus-

By comparison, under the same environment, for files totaling roughly 280GB, shuffling with terashuf took about 2.7 hours in total, while my own Python code took about 3.5 hours to run. It seems the Python code isn't all that much slower, which is quite acceptable — after all, terashuf is written in C++, and there's no shame in Python being slower than C++.

In most cases, the bottleneck for this global shuffling algorithm is disk I/O speed, so multiprocessing or multithreading generally doesn't help all that much.

Summary

This post gave a brief introduction to the idea of using disk space to perform a global shuffle of large files under limited memory, and provided a Python implementation. Readers who need more complex functionality are welcome to derive more sophisticated code from this as a starting point.