knighton
commented
1 year ago Thanks for reaching out!
With massive datasets, our serialization format choices are critical to ultimate observed performance of the system. If we really care about performance, we must own the format ourselves. Definitely tar is the general way to pack many blobs into a download-friendly form. But our use case, ML training, is much more specific: we also need extremely low latency cold random access on the granularity of individual samples, which all have the same fields.
Given those requirements, the high-level design: we have an index file and many shard files. Each shard file is a standalone dataset. The index file is a small JSON listing of shard metadata, which include the filename and how many samples the shard has, which concatenates the shards into a big dataset. A dataset is conceptualized as a table, of which each row is a sample. Columns have data types, which deserialize into int, str, PIL Image, etc.
For scalability, the shards are never loaded into memory; instead we read samples on the fly directly from disk. We store the byte offset where each sample is found at the head of the file as an int array. In the body of the file, samples are stored contiguously for fast random access with unbuffered i/o. In-between the two is a "neck" containing a tiny JSON dict of shard metadata (like column names and data types) which is skipped over. In the event the index is lost, it can be recovered using that data in the neck (this has happened).
The index (list of shard metadata) is kept in memory, including a dynamically generated O(1) sample ID range to shard lookup table. On StreamingDataset.__getitem__, we turn that global sample index into a shard and sample index within that shard, then seek directly to the sample (begin, end) in the file header, then seek to the sample data. This scheme of having shard metadata dicts in the index allows us to have other bells and whistles around compression and data integrity and so on, but the core thing is the very fast "two seek" random access.
The downside of the precomputed offsets is that once written, shards are set in stone. But there is nothing wrong with writing a new shard. A common idiom is to walk the samples of an existing dataset, doing some processing or filtering, and writing those samples as a new dataset. Recall that the index is in simple, hackable JSON; by modifying this listing one can subset, concatenate, etc. these datasets.
Why is that format better?
Performance! How does this script fare on your system? https://github.com/mosaicml/streaming/blob/main/scripts/serialization/compare.py
For features, see recent blog post https://www.mosaicml.com/blog/mosaicml-streamingdataset if you haven't.
If you don't like MDS, you can use other formats instead like JSONL or CSV, and StreamingDataset will merrily take your JSONL or CSV shards. The catch is that it will have to pair each JSONL/CSV file with an additional file that we generate containing precise sample byte offset info exactly as in the header/"neck" portion as described above for it to work. We haven't made much use of this instead of MDS as of yet, but the option is there.
What is that format exactly? Is there a spec?
Sorry, do you have some good examples of specs in mind that you can link us to? You are not the first to ask these questions and we need to do better than pointing people to code.
What are some ways to work with MDS files outside of this streamer library?
StreamingDataset can be used directly without a DataLoader, lazily downloading shards on sample access. It is also fine to not "stream" at all, just pointing it at local files.
If you don't like the intricate shards/samples to nodes/ranks/workers logic and want just a local map Dataset, see https://github.com/mosaicml/streaming/blob/main/streaming/base/local.py
Do you have any specific data processing use cases in mind that we aren't covering? We would be happy to gather any feedback you can share.
rom1504
jacobwjs
seralf
Hi, Thanks for the lib!
I think the 2 main alternatives in the pytorch world are webdataset and torch data. They both support tar files as shard format. The benefit of tar is that it's standard and supported by hundred of tools.
You decided to go instead with your own custom format for sharding. Do you have some pointers as to
I noticed that you used img2dataset for downloading laion400m, great! Maybe we could add support to WDS file type directly in img2dataset to avoid having to do a (costly) conversion after the fact.