search

mosaicml / streaming

A Data Streaming Library for Efficient Neural Network Training
https://streaming.docs.mosaicml.com
Apache License 2.0
1.09k stars 136 forks source link

Support large size index.json (20GB +) #662

Open andreamad8 opened 5 months ago

andreamad8 commented 5 months ago

🚀 Feature Request

Large index.json are slow to load. Currently, I am trying to increase shard size, so stream.py#L473 will be faster (hopefully).

Motivation

These two steps are very slow for large index.json files.

https://github.com/mosaicml/streaming/blob/main/streaming/base/stream.py#L461

and 

https://github.com/mosaicml/streaming/blob/main/streaming/base/stream.py#L473

especially with large scale dataset (e.g, Billion same).

ASchneidman commented 5 months ago

Some more context, we have a dataset with ~1.2 billion samples at like 1MB/sample. The index.json file of the merged dataset will be in the tens of GBs, which makes the dataset prohibitively slow to initialize.

snarayan21 commented 4 months ago

Hey, we have seen index.json load times be slow. I think that this is because we download the index file on every single rank, rather than downloading it on just one rank and then broadcasting its contents to other ranks. Downloading a file that's a few GB from cloud storage just on one rank should be relatively fast. This would be a good enhancement but isn't high priority for us right now -- if it's not too much of a hassle, mind submitting a PR?