Dataset uses excessive memory when loading files

[avishaiElmakies]

Describe the bug

Hi
I am having an issue when loading a dataset.
I have about 200 json files each about 1GB (total about 215GB). each row has a few features which are a list of ints.
I am trying to load the dataset using load_dataset.
The dataset is about 1.5M samples
I use num_proc=32 and a node with 378GB of memory.
About a third of the way there I get an OOM.
I also saw an old bug with a similar issue, which says to set writer_batch_size. I tried to lower it to 10, but it still crashed.
I also tried to lower the num_proc to 16 and even 8, but still the same issue.

Steps to reproduce the bug

dataset = load_dataset("json", data_dir=data_config.train_path, num_proc=data_config.num_proc, writer_batch_size=50)["train"]

Expected behavior

Loading a dataset with more than 100GB to spare should not cause an OOM error.
maybe i am missing something but I would love some help.

Environment info

• datasets version: 3.5.0
• Platform: Linux-6.6.20-aufs-1-x86_64-with-glibc2.36
• Python version: 3.11.2
• huggingface_hub version: 0.29.1
• PyArrow version: 19.0.1
• Pandas version: 2.2.3
• fsspec version: 2024.9.0

[avishaiElmakies]

small update: I converted the jsons to parquet and it now works well with 32 proc and the same node.
I still think this needs to be understood, since json is a very popular and easy-to-use format.

[lhoestq]

Hi ! The JSON loader loads full files in memory, unless they are JSON Lines. In this case it iterates on the JSON Lines in a memory efficient manner.

I know there is an ijson package that works similarly but for general JSON files, maybe it can help and remove the need to load full JSON files in memory

[avishaiElmakies]

Hi, i understand that json files are probably loaded into memory to read them but aren't they released when we write all the file content into arrow or something?

[lhoestq]

Yes correct, the JSON data is only in memory during the conversion to Arrow. Then, the data is memory mapped from you disk

[avishaiElmakies]

so the json files are all loaded into memory before converting to arrow? or do they convert 1 json at a time and then they are realeased?
I don't understand how 200GB worth of jsons fill a 378GB node's memory.

[lhoestq]

Each process converts one JSON file at at time, So the total memory usage is num_proc * json_file_size * overhead, where overhead can be around 2 or 3 for the conversion.

So it's indeed surprising that you run out of memory. Is the dataset available somewhere ? or a subset maybe ?

[avishaiElmakies]

This is a tokenized dataset I created for training a speech-language model with a few features (so it is not private but not easily available). I can send/upload a shard or two and you can copy them however many times you want so you can debug. this should give you something comparable to what I have, but will be easier than creating it yourself. so if you want that, let me know :)

[lhoestq]

Maybe you can measure the memory usage when loading 1 file with num_proc=1 ? This should already be helpful.

Memory usage for tokenized data can be bigger than just text, for example the tokens type can be inferred as int64 and the lists offsets are int32

[avishaiElmakies]

OK, I will try to do this in the near future. I am a little swamped at the moment. do you have a preferred tool?

also My data is just list of ints, there is no offsets

[jxma20]

so the json files are all loaded into memory before converting to arrow? or do they convert 1 json at a time and then they are realeased? I don't understand how 200GB worth of jsons fill a 378GB node's memory.

Hello! Is your query solved? I have the same confusion and would like to ask you for advice

[avishaiElmakies]

no, the issue is still present. I converted the json files to parquet, but json seems to have a problem.

Unfortunately i didn't have the time to try and profile the memory usage for 1 file. So if you want to do that, it will be great!

[jxma20]

My dataset is about image descriptions, stored as a 20MB JSON file on disk. However, I need to use the map function to preprocess the images, and after computation, the preprocessed dataset amounts to 70GB. My server has 122GB of RAM, but it still runs out of memory (OOM). This issue is very similar to yours.

After some research during this period, I found that the map function does not perform disk mapping in memory while working. Using the command find /DataB/mjx -type f -mmin -10, I discovered that no temporary cache files were modified or created during program execution, meaning the data was continuously loaded into memory. After several attempts, I found that adding the parameter cache_file_name="your/path" to the map function can enable memory-disk mapping. This is a strange setting, but after adding this parameter, the memory usage dropped to only 7GB, indicating that once the writer_batch_size worth of data is read into the disk cache, the corresponding data in memory is released. However, I don't think this is the intended behavior by the author, as memory-disk caching should have been enabled without needing this additional parameter.

Finally, here is my map function call. I hope it helps you.
train_data = train_data.map(process_fun, cache_file_name='./cache_file', remove_columns=['image_name', 'question_type', 'concern', 'question', 'candidate_answers', 'answer'])