Repository navigation
Require a huge amount of memory #6
Description
Hi,
We are trying to train a super BPE for an upcoming LLM, it will be a corpus of 12T in 30 languages.
We are trying to train the tokenizers but the second step of the training is require an immense amount of RAM,
Currently for 60GB of data the job is killed on the 1.5TB of RAM node due to OOM error.
Specifically, it's the step : tokenize_words that fails
I'm having doubt about the fact that 60GB of data require this amount of memory,
Is there any memory leak or do you have a clear ideas on the reason why the step 2 require much more RAM than the step one ?
Best Regards,
Activity
alisawuffles commentedon May 9, 2025
Hi, thanks for the question! Stage 2 does require more memory, we discuss this briefly at the end of §2.2 in the paper. Basically, the dictionary of pair counts is really large because there are many more unique pretokens w/o whitespace pretokenization (≈ the set of all substrings of all documents) compared to with (≈ the set of words across all documents). In our experiments we trained on 10 GB of data on 100 CPUs and 1T of memory, and I think we were able to train on up to 30 GB of data before encountering OOM. I would recommend reducing the size of the dataset and trying again.
iliailiaai commentedon May 2, 2026
Hi! I’m currently working on a research project around LLM pretraining, and I’m trying to integrate SuperBPE into it.
At the moment I have 128 GB of RAM available, which only lets me train the Stage 2 tokenizer on about 2 GB of data before running into memory issues. That is a bit smaller than I’d like for my experiments.
As one of the authors, could you suggest any promising directions for optimizing Stage 2 memory usage? Are there particular data structures, parts of the training loop, or algorithmic changes that you think would be worth exploring?
I’d be very grateful for any pointers or intuition you could share.
PythonNut commentedon May 3, 2026
If you look into the HuggingFace training code, you'll find that the system is heavily based around the idea that pretokenization allows one to "deduplicate" the training data at the pretoken level, since the data can be losslessly represented by a map of unique prerokens to their counts. When pretokenization is disabled, this (beneficial) effect goes away.
The HF BPE training itself assumes the post-dedupe data size is small and aggressively trades memory usage for reduced runtime. IIRC, some of the main data structures that consume a lot of memory are: the pair count map which stores one integer count per pair in the dataset (and this hash table is even bigger than it's contents due to low load factor stemming from its dynamic size) and the back pointer arrays which store the offsets of every pair's occurrences in the data. It's been a while since we've run the memory profiling so I'm sure there are other major contributors.
We'd definitely be interested in any progress in this direction!
alisawuffles commentedon Jun 6, 2026
This may be of interest :) Faster Superword Tokenization