"Solution" to memory hogging in train_new_from_iterator with a hack about tokenizers

[ArthurZucker]

Thanks!
Probably because the new tokenizer does not use the normalizer / uses only a pre-tokenizer? Not super sure but thanks for sharing! 🤗

[purefall]

Hi, I am also having a similar issue while training a tokenizer using on Refined Web Dataset.

Similar to issue, I am using the train_from_iterator() function and training a SentencePieceBPETokenizer(), and during the Pre-processing sequences I ran out-of-memory.

I don't believe switching to another tokenizer serves my needs and, it would be great if you could provide some insights on what I can do if I want to train on large datasets.

[morphpiece]

Can you give a short script to reproduce the problem?

[purefall]

Sure the code below throws a out-of-memory error, process_files is the path list of arrow files downloaded from HF of RefinedWeb Dataset:

from datasets import load_dataset
data_files = {"train": process_files}

dataset = load_dataset("arrow", data_files=data_files, split="train", streaming=True)

# Build an iterator over this dataset
def batch_iterator(input_sentence_size=None, batch_size=1000):
    for elem in dataset.iter(batch_size=batch_size):
        yield elem["content"]


tokenizer = SentencePieceBPETokenizer()
tokenizer.train_from_iterator(
    iterator=batch_iterator(),
    vocab_size=30_000,
    min_frequency=5,
    show_progress=True,
    limit_alphabet=500,
)

[morphpiece]

I just ran the script and didn't observe any OOM.

I used a 30 vCPU/240 GB AMD-Epyc server. The memory consumption increased linearly during count_pairs stage and the max was about 60%.

Caveats : I used FineWeb-Edu-10BT sample, saved as parquet files instead of arrow.

Hope that helps.

[ArthurZucker]

There is an ongoing PR for arrow support, if that helps we should def merge it!