mosaicml/streaming: shard the dataset, stream the shards, keep the epoch deterministic

The problem StreamingDataset is aimed at

Training on a dataset that lives in a bucket is not the same problem as training on a dataset on a local disk. The naive fix is to copy the whole thing down before the job starts, which costs wall-clock time proportional to dataset size and multiplies storage by the number of nodes. The other naive fix is a plain PyTorch IterableDataset that opens objects as it goes, which is easy until you run on more than one process and discover that each worker is making its own requests and each epoch sees a different order of samples. The README frames the goal directly: make training on large datasets from cloud storage "as fast, cheap, and scalable as possible," with an explicit emphasis on multi-node, distributed training for large models and on correctness guarantees. That last phrase is the interesting one. The project is not only about throughput. It is about making the sample ordering a function of the dataset and the epoch number rather than of how many workers happened to be running. The intended user is someone training a model that does not fit alongside its data, on a cluster, with the dataset sitting in S3, GCS, Azure Blob, OCI, Databricks storage, or an S3-compatible store such as Cloudflare R2, CoreWeave, or Backblaze B2.

Shards, a local cache, and a sample order that does not depend on worker count

The mechanism has two halves. The first is a conversion step: MDSWriter takes your raw samples and writes them into shard files in MDS (Mosaic Data Shard) format, or into CSV, TSV, or JSONL. The README's example declares a columns dictionary mapping each field to a type, for instance 'image' to 'jpeg' and 'class' to 'int', plus an optional compression setting such as 'zstd', and then writes samples one at a time inside a with block. MDS is described as able to encode and decode any Python object, which is what makes it the general-purpose option among the four formats. The second half is StreamingDataset itself, constructed with two paths: remote, the persistent location of the full dataset, and local, a working directory where the dataset is cached during operation. The README's example uses remote='s3://my-bucket/path-to-dataset' and local='/tmp/path-to-dataset'. Samples are then indexed directly, as in dataset[1337], and the object is passed straight into a torch.utils.data.DataLoader. The cache is the part that makes the multi-node story work. Because each rank reads shards through a local directory rather than pulling every sample from the bucket, repeated access within an epoch and across epochs does not translate into repeated object store requests, and shuffle=True is resolved against the shard layout rather than against whatever order the network happened to deliver. The README calls StreamingDataset a drop-in replacement for IterableDataset, and the API surface shown is small enough that the claim is plausible: two paths, a shuffle flag, and indexing. What you are buying with that small surface is a deterministic mapping from epoch and rank to samples.

Getting it running: install, convert, upload, construct

The install is a single pip command: pip install mosaicml-streaming. Conversion is a Python script using MDSWriter, with out set to a local or remote directory, columns as a field-to-type mapping, and compression set to something like 'zstd'. Uploading is left to whatever tool you already use for your object store; the README shows aws s3 cp --recursive path-to-dataset s3://my-bucket/path-to-dataset as one example. Consumption is four lines: import StreamingDataset from streaming, construct it with local, remote, and shuffle=True, index a sample to check it, and hand it to DataLoader. There is no separate configuration file, no daemon, and no service to run. That is a real advantage for a library in this space, because the failure modes are confined to the conversion script and the two paths. It also means the operational surface is exactly the two directories. If local points at a filesystem that fills up, or at a path that is not shared correctly between ranks, the failure will look like a training problem rather than a storage problem. The README does not spell out a cache eviction policy or a size limit, so the size of local is a parameter you have to reason about yourself based on the dataset and the number of concurrent jobs.

Where this design gets in the way

The conversion step is the honest cost. Your data has to be rewritten into MDS, CSV, TSV, or JSONL before training, which means a pipeline stage that did not exist before and a second copy of the data during the transition. For a dataset that changes frequently, or one assembled on the fly from a database query, that stage is friction rather than a one-time expense. The second constraint is the cache directory. StreamingDataset is not a stateless reader; it needs a local path with room for shards, and the README does not describe a bounded cache, so on a node with a small ephemeral disk the local directory is a resource you have to budget. Third, the formats are not equivalent. MDS is the one described as able to encode and decode any Python object; CSV, TSV, and JSONL carry their own type limitations, so a dataset with nested or binary fields will push you toward MDS whether or not that is what you wanted. Finally, the drop-in claim has a boundary. A drop-in replacement for IterableDataset is still an IterableDataset, so anything that depends on a map-style dataset with a stable __len__ and random access semantics will need checking, even though the README's example does index a sample directly. None of these are defects in the design; they are the terms of the trade.

What you give up compared to a plain local dataset

The obvious alternative is not another streaming library. It is downloading the data once and using a normal map-style PyTorch Dataset over local files. That approach has no conversion format, no cache directory, no object store credentials in the training process, and no dependency on network behaviour during a run. Its costs are the ones StreamingDataset exists to remove: staging time proportional to dataset size, storage multiplied across nodes, and a pipeline that has to be re-run whenever the data changes. The difference in approach is where the work happens. A local dataset pays at the start of the job and then reads at disk speed. StreamingDataset pays a conversion cost once and a cache-management cost continuously, in exchange for a training process that does not care how big the dataset is relative to the node. For a dataset in the tens of gigabytes on a single machine, the local approach is simpler and there is no reason to take on the shard format. The calculus flips when the dataset is large enough that staging it is the dominant cost of a run, or when the number of nodes is large enough that per-node copies are wasteful. The README's own framing, multi-node distributed training for large models, is a fair statement of where the line sits.

Maintenance, versioning, and the licence

The project is not archived and has a regular release cadence, with v0.11.0 in January 2025, v0.12.0 in April 2025, and v0.13.0 in July 2025, and the repository shows activity after that. The version numbers are still in the 0.x range, which is worth noting for anyone pinning a dependency: minor releases in a 0.x line can carry breaking changes by convention, so an upgrade from v0.12.0 to v0.13.0 deserves a look at the release notes rather than an unpinned install. The licence is Apache-2.0, which permits commercial use and modification and includes an explicit patent grant; it also carries notice and attribution obligations for redistributed copies. That is a summary of the licence identifier as stated in the repository, not legal advice, and anyone embedding the library in a distributed product should read the LICENSE file rather than this paragraph. The practical maintenance cost is low in one sense, because there is no service to operate, and higher in another, because the conversion scripts and the cache layout are yours to maintain and they are coupled to the library's shard format. If the MDS format changes across a release, the data you converted under an older version is the thing you have to think about.

Who should adopt it, and what to check before you do

The fit is narrow and specific. You are training a model on a dataset that is too large to stage on each node, you are already running across multiple machines, and your data sits in one of the supported object stores. In that situation the shard-plus-cache design solves a problem that is otherwise solved by waiting. You should not adopt it if the dataset fits on one node's disk, if the data is regenerated per run rather than persisted, or if adding a format conversion stage to your pipeline is not something you can do. The README's quick start is short enough that the decision is really about the conversion step, not about the API. Before committing, verify that your object store is on the supported list, that the local cache path has room for the shards your job will touch, and that every field in your samples maps cleanly onto a column type MDSWriter accepts. The three end-to-end tutorials linked from the README, CIFAR-10, FaceSynthetics, and SyntheticNLP, are the fastest way to see whether that mapping holds for data shaped like yours.

Editorial conclusion

Adopt StreamingDataset if your training set is too large to stage on every node and you already train across multiple machines, because the shard plus cache design is the part that makes node counts stop changing what the model sees. Do not adopt it for a dataset that fits comfortably on one node's disk, or if you cannot add a data conversion step to your pipeline. Before committing, verify three things against your own setup: that your object store is one of the supported backends (AWS, OCI, GCS, Azure, Databricks, or an S3-compatible store), that the node running the job has a local path with enough room for the cache, and that your sample types map onto the column types MDSWriter accepts, since MDS can encode and decode any Python object but the writer still needs a declared type per field.

Official sources

1. License: Apache-2.0
2. mosaicml/streaming on GitHub
3. Project website
4. README
5. Releases