Support large size index.json (20GB +)
andreamad8 opened this issue · comments
🚀 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).
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.
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?