Require a huge amount of memory

[jgcb00]

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,

[alisawuffles]

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.