Dataset load_from_disk is too slow

[avacaondata]

@lhoestq

Describe the bug

It's not normal that I have to wait 7-8 hours for a dataset to be loaded from disk, as there are no preprocessing steps, it's only loading it with load_from_disk. I have 96 cpus, however only 1 is used for this, which is inefficient. Moreover, its usage is at 1%... This is happening in the context of a language model training, therefore I'm wasting 100$ each time I have to load the dataset from disk again (because the spot instance was stopped by aws and I need to relaunch it for example).

Steps to reproduce the bug

Just get the oscar in spanish (around 150GGB) and try to first save in disk and then load the processed dataset. It's not dependent on the task you're doing, it just depends on the size of the text dataset.

Expected results

I expect the dataset to be loaded in a normal time, by using the whole machine for loading it, I mean if you store the dataset in multiple files (.arrow) and then load it from multiple files, you can use multiprocessing for that and therefore don't waste so much time.

Environment info

• datasets version: 1.8.0
• Platform: Ubuntu 18
• Python version: 3.8

I've seen you're planning to include a streaming mode for load_dataset, but that only saves the downloading and processing time, that's not being a problem for me, you cannot save the pure loading from disk time, therefore that's not a solution for my use case or for anyone who wants to use your library for training a language model.

[lhoestq]

Hi ! It looks like an issue with the virtual disk you are using.

We load datasets using memory mapping. In general it makes it possible to load very big files instantaneously since it doesn't have to read the file (it just assigns virtual memory to the file on disk).
However there happens to be issues with virtual disks (for example on spot instances), for which memory mapping does a pass over the entire file, and this takes a while. We are discussing about this issue here: #2252

Memory mapping is something handled by the OS so we can't do much about it, though we're still trying to figure out what's causing this behavior exactly to see what we can do.

[avacaondata]

Okay, that's exactly my case, with spot instances... Therefore this isn't something we can change in any way to be able to load the dataset faster? I mean, what do you do internally at huggingface for being able to use spot instances with datasets efficiently?

[lhoestq]

There are no solutions yet unfortunately.
We're still trying to figure out a way to make the loading instantaneous on such disks, I'll keep you posted

[Tommigun1980]

@lhoestq This happens to me on my Mac (Tahoe) as well. If I load a 2TB dataset via datasets.load_from_disk(dataset_dir) (from an external SSD drive) it spends hours loading the files. Safe to say this is untenable.

I can also see in Activity Monitor that it seems to iterate ("read in") all the files as the 'Bytes Read' column shows the approximate size of the dataset once it finally finishes. RAM usage doesn't go up so it doesn't consume RAM, but it sure does iterate the entire dataset contents.

Is it intended that the entire dataset is iterated at load time? This makes it sound like it is not intended:

We load datasets using memory mapping. In general it makes it possible to load very big files instantaneously since it doesn't have to read the file (it just assigns virtual memory to the file on disk).

However; reading the sources it does seem like the HF dataset does iterate the entire dataset at load time:

Dataset.load_from_disk:

table_cls = InMemoryTable if keep_in_memory else MemoryMappedTable

arrow_table = concat_tables(
    thread_map(
        table_cls.from_file,
            ...

MemoryMappedTable.from_file:

@classmethod
def from_file(cls, filename: str, replays=None):
    table = _memory_mapped_arrow_table_from_file(filename)
    table = cls._apply_replays(table, replays)
    return cls(table, filename, replays)

def _memory_mapped_record_batch_reader_from_file(filename: str) -> pa.RecordBatchStreamReader:
    memory_mapped_stream = pa.memory_map(filename)
    return pa.ipc.open_stream(memory_mapped_stream)

def _memory_mapped_arrow_table_from_file(filename: str) -> pa.Table:
    opened_stream = _memory_mapped_record_batch_reader_from_file(filename)
    pa_table = opened_stream.read_all()
    return pa_table

So the entire stream does seem to be consumed in opened_stream.read_all(). So if the design is that the table is backed by mmap'd memory, but the file is read sequentially to build the table in virtual memory, then I can't see how the statement

[...]makes it possible to load very big files instantaneously since it doesn't have to read the file (it just assigns virtual memory to the file on disk).

applies here, as load time seems to scale roughly with dataset size and takes hours for larger datasets. Or am I just misreading the sources?

I am in a bit of a blocked state here because resuming training takes several hours every time I use a larger dataset, so any information on intended behavior and mitigation would be hugely appreciated. Thanks!