DOCSEARCHbyAlgolia
Submit Docs

MosaicML Streaming

repository·main·Indexed 23 days ago

https://github.com/mosaicml/streaming

A high-performance data loading library for training large models from cloud storage (S3, GCS, OCI). It provides fast, deterministic, and scalable data streaming that integrates with PyTorch, supporting map-style random access, dataset mixing via the Stream class, and local cache management. The library utilizes the Mosaic Data Shard (MDS) format and includes a Streaming Simulator for performance planning and debugging.

Tokens
32.4K
Snippets
96
Records
141
Agent score
80%

What's inside mosaicml-streaming

  1. Key features of StreamingDataset

    main

    StreamingDataset provides several features optimized for large-scale distributed training:

    • Elastic Determinism: Ensures samples are in the same order regardless of the number of GPUs, nodes, or CPU workers, allowing for reproducible debugging across different hardware configurations.
    • Instant Mid-Epoch Resumption: Enables resuming training in the middle of an epoch in seconds, reducing latency and costs.
    • High Throughput: Uses the MDS format to minimize sample retrieval latency.
    • Effective Shuffling: Implements specialized shuffling algorithms that maintain shuffle quality while reducing egress costs.
    • Random Access: Supports direct indexing via dataset[i] or NumPy-style indexing.
    • Flexible Data Mixing: Allows seamless, just-in-time shuffling and mixing of different data sources using specific batching and sampling methods.
    • Disk Usage Limits: Can dynamically delete least recently used (LRU) shards to stay under a specified local disk limit.
    • Parallelism-aware: Supports data, sequence, and tensor parallelism by ensuring correct sample replication across GPUs.

  2. Supported Dataset Formats in Streaming

    main

    To use StreamingDataset, raw data must be converted into one of the following supported serialized formats. The choice of format impacts the performance of cold random access during training.

    • MDS (Mosaic Data Shard): The most performant format, optimized for fast sample random-access. It stores data in a serialized tabular form. Use streaming.MDSWriter to create MDS datasets.
    • CSV/TSV: Plaintext tabular formats using comma or tab delimiters. Use streaming.CSVWriter, streaming.TSVWriter, or streaming.XSVWriter to create these.
    • JSONL: A format where each sample is a JSON dictionary terminated by a newline. Use streaming.JSONWriter to create JSONL datasets.

    These formats can encode and decode most Python objects, including images, text, video, and multimodal data.

  3. Optimize MDS serialization efficiency for ndarrays

    main

    When using MDSWriter to store ndarrays, you can optimize storage space by being more specific in the columns configuration. The efficiency hierarchy (from least to most efficient) is:

    1. Dynamic shape and dynamic dtype: Use ndarray.
    2. Dynamic shape and fixed dtype: Use ndarray:dtype.
    3. Fixed shape and fixed dtype: Use ndarray:dtype:shape (e.g., ndarray:float32:1,28,28).

  4. How shuffling works in Streaming

    main

    Streaming's shuffling process follows four main steps to balance randomness with download efficiency:

    1. Partitioning: StreamingDataset downloads shard metadata (index.json) and partitions sample IDs among nodes, devices, and workers. Shards and samples are initially not shuffled.
    2. Shard Shuffling: The order of the shards is shuffled.
    3. Canonical Node Splitting: Shards are split into num_canonical_nodes (buckets of samples). Some shards may be split between two canonical nodes.
    4. Intra-canonical Node Shuffling: Samples are shuffled within each canonical node using the specified shuffle_algo. These nodes are then assigned to physical training nodes.

  5. Configure data replication for TP and SP strategies

    main

    Parallelism strategies like Tensor Parallelism (TP) and Sequence Parallelism (SP) require multiple devices to receive the same data samples (replication) rather than unique samples.

    To implement this, set the replication argument when initializing StreamingDataset. This argument specifies how many consecutive devices should receive the same data.

  6. How StreamingDataset partitions samples for distributed training

    main

    StreamingDataset is designed for distributed model training by splitting samples across nodes (host CPU systems), ranks (typically GPUs), and workers (CPU processes handling fetching).

    1. Partitioning: Upon initialization, StreamingDataset downloads index.json files and partitions samples to minimize redundant downloads.
    2. Shuffling: You can enable shuffling by setting shuffle=True. To maintain performance and control download demand, StreamingDataset performs intra-node shuffling, which minimizes duplicate shard downloads between different nodes.
    3. Retrieval: Shards are downloaded from remote storage to local disk on-the-fly. Dataloader workers check if the required shard is on disk; if not, they download it before loading the sample for training.

  7. Use `assert` for invariants, not data validation

    main

    In Streaming, assert statements should only be used in test cases or for verifying invariants (to assist type checking).

    Do not use assert for data validation, because assertions can be disabled in Python using the -O flag. For data validation, explicitly raise exceptions:

    if parameter is None:
    raise ValueError("parameter must be specified and cannot be None")

  8. Iterate over datasets of any size without divisibility requirements

    main

    StreamingDataset does not require the total number of samples to be perfectly divisible by the number of training devices. Instead of dropping samples to satisfy divisibility, the dataset ensures that each device processes the same count by repeating a different selection of samples each epoch, ensuring no samples are lost.

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

  9. Achieve elastic determinism across varying GPU counts

    main

    Streaming supports elastically deterministic training and resumption, allowing you to change the number of GPUs used in a training job while maintaining the same loss curve and global batch size.

    To ensure determinism between different runs (e.g., moving from 32 GPUs to 8 GPUs), you must satisfy two requirements:

    1. Constant Global Batch Size: The total number of samples processed across all GPUs per step must remain the same. If you decrease the number of GPUs, you must increase the per-device batch_size proportionally.
    2. Consistent num_canonical_nodes: You must explicitly set the num_canonical_nodes parameter to the same value used in the original run. If not specified, Streaming defaults this value to the number of physical nodes used in the first run.

    Requirements for Global Batch Size: The global batch size must be divisible by all the numbers of GPUs you intend to use. For example, a global batch size of 18 allows deterministic training on 1, 2, 3, 6, 9, or 18 GPUs, but not on 7 GPUs.

  10. Difference between StreamingDataset, streams, and epoch_size

    main

    Core Concepts

    • StreamingDataset: The main dataset class. It is an IterableDataset that downloads shard files to provide a continuous flow of samples. It can combine multiple streams (data sources) into one dataset.
    • Streams: Individual data sources that the StreamingDataset consumes.

    Size and Length

    • epoch_size: The number of samples per epoch of training.
    • __len__(): Returns epoch_size divided by the number of devices (samples seen per device, per epoch).
    • size(): Returns the total number of unique samples in the underlying dataset (may differ from epoch_size due to upsampling/downsampling).

  11. Convert vision datasets to Mosaic Dataset Shard (MDS) format

    main

    To use the Streaming Dataset, you must first convert datasets from their native formats into the Mosaic Dataset Shard (MDS) format. Once converted, MDS files can be stored on local file systems (disk, NAS) or object stores (GCS, OCS, S3) and streamed efficiently for deep learning training.

    Conversion is typically performed using specialized scripts provided in the streaming/vision/convert/ directory. These scripts take an input root directory containing the native dataset and an output root directory where the MDS shard files will be saved.

  12. Calculate and configure the required cache limit for storing shards

    main

    When shards are downloaded, they are stored on the node's local disk. If the node's available disk space is less than the required cache_limit, performance will suffer.

    If cache_limit is not explicitly set, the required cache size per node ($L$) is calculated as: $L = \frac{S \cdot N}{P}$ (where $S$ is average shard size, $N$ is total shards, and $P$ is physical nodes).

    For optimal performance, you can set a lower cache_limit based on your shuffling strategy:

    • For shuffle=False or 'py1s'/'py2s' algorithms: $L = 2 \cdot S \cdot \lceil\frac{C}{P}\rceil$ (where $C$ is the number of canonical nodes).

    • For shuffle-block-based algorithms ('py1e' or 'py1br'): $L = k \cdot S \lceil \frac{B}{Q} \rceil \cdot \lceil\frac{C}{P}\rceil$

      • $k=1$ for 'py1e' (better cache performance).
      • $k=2$ for 'py1br'.
      • $B$ is the shuffle block size (number of samples).
      • $Q$ is the average number of samples per shard.