In this article, we will tackle a programming problem:
How to globally and randomly shuffle hundreds of gigabytes (GB) of text files under limited memory?
The background of this problem is quite clear: modern pre-trained models often involve tens or even hundreds of GBs of corpora. To ensure the model trains effectively, it is necessary to perform a global random shuffle of the training data. However, for many people, a corpus of several hundred GBs is often larger than their available RAM. Therefore, how to achieve a global random shuffle under limited memory is a problem worth investigating.
Existing Tools
Assuming our files are stored line by line (i.e., one line represents
one sample), our goal is to randomly shuffle the file by lines. If we
have only one file and its size is significantly smaller than the
available memory, we can use the shuf command that comes
with Linux:
shuf input.txt -o output.txtThe reason I emphasize that the file size must be significantly
smaller than the memory is that shuf loads the entire file
into memory before shuffling, which requires a sufficiently large amount
of RAM. To address this, an improved version called
terashuf (link) exists. It works
by partitioning the file and using hard drive space instead of memory,
allowing us to shuffle files larger than the available RAM.
It seems that terashuf already fully meets our needs.
Theoretically, this is true, but sometimes we may have personalized
requirements, such as mixing multiple files together for a random
shuffle or splitting the shuffled output into multiple files, etc.
Therefore, it is best if we can implement it ourselves in Python to
handle more complex customization needs.
Shuffling Algorithm
Now let’s look at what the algorithm for global shuffling under limited memory actually looks like. Generally, the steps are as follows:
Assume the file has a total of mn lines; split it into m files, each containing n lines.
Randomly shuffle the contents of each n-line file internally. Since n can be arbitrarily specified, this step can be completed in memory.
Read the first line of each file (obtaining m lines of data), and randomly write these m lines into the output file.
Sequentially read the 2^{nd}, \dots, n^{th} lines of each file, repeating step 3.
Simply put, you first shuffle vertically and then shuffle horizontally to obtain a sufficiently shuffled result that approximates a global shuffle, as shown in the figure below:
Note that this algorithm only guarantees a result that is as random as possible, but it cannot guarantee that every possible permutation can occur. For example, it is impossible for the first two samples of the output to be exactly the first two samples of the first file. To truly achieve an equal probability for all permutations, we would need to sample each step according to the remaining number of lines in each file, rather than reading the k-th line of every file simultaneously before moving to the (k+1)-th line. However, sampling according to probability at every step increases the sampling cost, and for practical use, the difference in effect is negligible, so it is not strictly necessary.
In practice, when splitting files by n lines, the last file might have fewer than n lines. If we don’t mind this small detail, we can still follow the process above; once the last file is exhausted, it can simply return empty lines without causing the process to crash. Of course, this would cause the samples from the last file to be positioned earlier in the shuffle. For perfectionist readers who cannot accept this, one could consider introducing rejection sampling when reading the last file, with a rejection rate of 1 - \frac{\text{remaining lines in the last file}}{\text{remaining lines in other files}}.
Reference implementation (without rejection sampling):
Performance Testing
If the task is just a combination of merging, shuffling, and
splitting, it can actually be achieved using shell commands combined
with terashuf. The approximate command would be:
cat corpus/*.json | TMPDIR=/root/tmp MEMORY=20 ./terashuf | split -l 100000 -a 5 -d - corpus-Upon comparison, in the same environment, for files with a total size
of approximately 280GB, the total time for shuffling using
terashuf was roughly 2.7
hours, while the Python code written by the author took
approximately 3.5 hours. It seems the
Python code is not significantly slower and is quite acceptable; after
all, terashuf is written in C++, and it is no shame for
Python to be slower than C++.
In most cases, the bottleneck of this global shuffling algorithm is the disk I/O speed; therefore, multi-processing or multi-threading is generally not particularly effective.
Summary
This article briefly introduced the idea of using hard drive space to perform global shuffling of large files under limited memory and provided a Python implementation. Readers in need can derive code for more complex requirements based on this.
When reprinting, please include the original address of this article: https://kexue.fm/archives/8662
For more detailed reprinting matters, please refer to: “Scientific Space FAQ”