/streaming

A Data Streaming Library for Efficient Neural Network Training

Primary LanguagePythonApache License 2.0Apache-2.0


Fast, accurate streaming of training data from cloud storage

PyPi Version PyPi Package Version Unit test PyPi Downloads Documentation Chat @ Slack License


๐Ÿ‘‹ Welcome

We built StreamingDataset to make training on large datasets from cloud storage as fast, cheap, and scalable as possible.

Itโ€™s specially designed for multi-node, distributed training for large modelsโ€”maximizing correctness guarantees, performance, and ease of use. Now, you can efficiently train anywhere, independent of your training data location. Just stream in the data you need, when you need it. To learn more about why we built StreamingDataset, read our announcement blog.

StreamingDataset is compatible with any data type, including images, text, video, and multimodal data.

With support for major cloud storage providers (AWS, OCI, and GCS are supported today; Azure is coming soon), and designed as a drop-in replacement for your PyTorch IterableDataset class, StreamingDataset seamlessly integrates into your existing training workflows.

The flow of samples from shards in the cloud to devices in your cluster

๐Ÿš€ Getting Started

๐Ÿ’พ Installation

Streaming can be installed with pip:

pip install mosaicml-streaming

๐Ÿ Quick Start

1. Prepare Your Data

Convert your raw dataset into one of our supported streaming formats:

  • MDS (Mosaic Data Shard) format which can encode and decode any Python object
  • CSV / TSV
  • JSONL
import numpy as np
from PIL import Image
from streaming import MDSWriter

# Local or remote directory in which to store the compressed output files
data_dir = 'path-to-dataset'

# A dictionary mapping input fields to their data types
columns = {
    'image': 'jpeg',
    'class': 'int'
}

# Shard compression, if any
compression = 'zstd'

# Save the samples as shards using MDSWriter
with MDSWriter(out=data_dir, columns=columns, compression=compression) as out:
    for i in range(10000):
        sample = {
            'image': Image.fromarray(np.random.randint(0, 256, (32, 32, 3), np.uint8)),
            'class': np.random.randint(10),
        }
        out.write(sample)

2. Upload Your Data to Cloud Storage

Upload your streaming dataset to the cloud storage of your choice (AWS, OCI, or GCP). Below is one example of uploading a directory to an S3 bucket using the AWS CLI.

$ aws s3 cp --recursive path-to-dataset s3://my-bucket/path-to-dataset

3. Build a StreamingDataset and DataLoader

from torch.utils.data import DataLoader
from streaming import StreamingDataset

# Remote path where full dataset is persistently stored
remote = 's3://my-bucket/path-to-dataset'

# Local working dir where dataset is cached during operation
local = '/tmp/path-to-dataset'

# Create streaming dataset
dataset = StreamingDataset(local=local, remote=remote, shuffle=True)

# Let's see what is in sample #1337...
sample = dataset[1337]
img = sample['image']
cls = sample['class']

# Create PyTorch DataLoader
dataloader = DataLoader(dataset)

๐Ÿ“š What next?

Getting started guides, examples, API references, and other useful information can be found in our docs.

We have end-to-end tutorials for training a model on:

We also have starter code for the following popular datasets, which can be found in the streaming directory:

Dataset Task Read Write
LAION-400M Text and image Read Write
WebVid Text and video Read Write
C4 Text Read Write
EnWiki Text Read Write
Pile Text Read Write
ADE20K Image segmentation Read Write
CIFAR10 Image classification Read Write
COCO Image classification Read Write
ImageNet Image classification Read Write

To start training on these datasets:

  1. Convert raw data into .mds format using the corresponding script from the convert directory.

For example:

$ python -m streaming.multimodal.convert.webvid --in <CSV file> --out <MDS output directory>
  1. Import dataset class to start training the model.
from streaming.multimodal import StreamingInsideWebVid
dataset = StreamingInsideWebVid(local=local, remote=remote, shuffle=True)

๐Ÿ”‘ Key Features


True Determinism

A unique feature of our solution: samples are in the same order regardless of the number of GPUs, nodes, or CPU workers. This makes it easier to:

  • Reproduce and debug training runs and loss spikes
  • Load a checkpoint trained on 64 GPUs and debug on 8 GPUs with reproducibility

See the figure below โ€” training a model on 1, 8, 16, 32, or 64 GPUs yields the exact same loss curve (up to the limitations of floating point math!)

Plot of elastic determinism

Instant Mid-Epoch Resumption

It can be expensive โ€” and annoying โ€” to wait for your job to resume while your dataloader spins after a hardware failure or loss spike. Thanks to our deterministic sample ordering, StreamingDataset lets you resume training in seconds, not hours, in the middle of a long training run.

Minimizing resumption latency can save thousands of dollars in egress fees and idle GPU compute time compared to existing solutions.

High throughput

Our MDS format cuts extraneous work to the bone, resulting in ultra-low sample latency and higher throughput compared to alternatives for workloads bottlenecked by the dataloader.

Tool Throughput
StreamingDataset ~19000 img/sec
ImageFolder ~18000 img/sec
WebDataset ~16000 img/sec

Results shown are from ImageNet + ResNet-50 training, collected over 5 repetitions after the data is cached after the first epoch.

Equal Convergence

Model convergence from using StreamingDataset is just as good as using local disk, thanks to our shuffling algorithm.

Plot of equal convergence

Below are results from ImageNet + ResNet-50 training, collected over 5 repetitions.

Tool Top-1 Accuracy
StreamingDataset 76.51% +/- 0.09
ImageFolder 76.57% +/- 0.10
WebDataset 76.23% +/- 0.17

StreamingDataset shuffles across all samples assigned to a node, whereas alternative solutions only shuffle samples in a smaller pool (within a single process). Shuffling across a wider pool spreads out adjacent samples more. In addition, our shuffling algorithm minimizes dropped samples. We have found both of these shuffling features advantageous for model convergence.

Random access

Access the data you need when you need it.

Even if a sample isnโ€™t downloaded yet, you can access dataset[i] to get sample i. The download will kick off immediately and the result will be returned when itโ€™s done - similar to a map-style PyTorch dataset with samples numbered sequentially and accessible in any order.

dataset = StreamingDataset(...)
sample = dataset[19543]

No divisibility requirements

StreamingDataset will happily iterate over any number of samples. You do not have to forever delete samples so that the dataset is divisible over a baked-in number of devices. Instead, each epoch a different selection of samples are repeated (none dropped) so that each device processes the same count.

dataset = StreamingDataset(...)
dl = DataLoader(dataset, num_workers=...)

๐Ÿ† Project Showcase

Here are some projects and experiments that used StreamingDataset. Got something to add? Email community@mosaicml.com or join our Community Slack.

  • BioMedLM: a Domain Specific Large Language Model for BioMedicine by MosaicML and Stanford CRFM
  • Mosaic Diffusion Models: Training Stable Diffusion from Scratch Costs <$160k
  • Mosaic LLMs: GPT-3 quality for <$500k
  • Mosaic ResNet: Blazingly Fast Computer Vision Training with the Mosaic ResNet and Composer
  • Mosaic DeepLabv3: 5x Faster Image Segmentation Training with MosaicML Recipes
  • โ€ฆmore to come! Stay tuned!

๐Ÿ’ซ Contributors

We welcome any contributions, pull requests, or issues.

To start contributing, see our Contributing page.

P.S.: We're hiring!

If you like this project, give us a star โญ and check out our other projects:

  • Composer - a modern PyTorch library that makes scalable, efficient neural network training easy
  • MosaicML Examples - reference examples for training ML models quickly and to high accuracy - featuring starter code for GPT / Large Language Models, Stable Diffusion, BERT, ResNet-50, and DeepLabV3
  • MosaicML Cloud - our training platform built to minimize training costs for LLMs, Diffusion Models, and other large models - featuring multi-cloud orchestration, effortless multi-node scaling, and under-the-hood optimizations for speeding up training time

โœ๏ธ Citation

@misc{mosaicml2022streaming,
    author = {The Mosaic ML Team},
    title = {streaming},
    year = {2022},
    howpublished = {\\url{<https://github.com/mosaicml/streaming/>}},
}