LitData

Use CombinedStreamingDataset to mix multiple datasets into a single stream. This is useful for creating data mixtures (e.g., combining SlimPajama and StarCoder).

Mixing Modes

• iterate_over_all=True (default): Iterates until all datasets are exhausted. Do not pass weights (LitData derives them from dataset lengths; passing both raises ValueError).
• iterate_over_all=False: Stops when any dataset is exhausted. You must pass explicit weights for your mixture.

Batching Methods (batching_method)

• stratified (default): Each batch contains a mix of samples from multiple datasets according to the provided weights.
• per_stream: Each batch comes from only one randomly selected dataset. Use this if datasets have different shapes or dtypes.

Other Options

• seed: Random seed (default 42).
• force_override_state_dict=True: Allows local constructor arguments to override a loaded checkpoint.

from litdata import StreamingDataset, CombinedStreamingDataset, StreamingDataLoader, TokensLoader
import os

train_datasets = [
    StreamingDataset(
        input_dir="s3://tinyllama-template/slimpajama/train/",
        item_loader=TokensLoader(block_size=2048 + 1),
        shuffle=True,
        drop_last=True,
    ),
    StreamingDataset(
        input_dir="s3://tinyllama-template/starcoder/",
        item_loader=TokensLoader(block_size=2048 + 1),
        shuffle=True,
        drop_last=True,
    ),
]

weights = (0.693584, 0.306416)
combined_dataset = CombinedStreamingDataset(
    datasets=train_datasets,
    seed=42,
    weights=weights,
    iterate_over_all=False,
)

train_dataloader = StreamingDataLoader(combined_dataset, batch_size=8, pin_memory=True, num_workers=os.cpu_count())

While CombinedStreamingDataset picks one sample from one dataset per iteration, ParallelStreamingDataset fetches a sample from every wrapped dataset at each iteration. This is ideal for generating new data on-the-fly from multiple sources.

Using a Transform

Provide a transform function to combine the samples. To ensure reproducible and resumable transformations, use the internal random number generators provided in the rngs dictionary (keys: "random", "numpy", and "torch").

from litdata import StreamingDataset, ParallelStreamingDataset, StreamingDataLoader
from typing import Tuple, Any, Dict

def transform(samples: Tuple[Any, ...], rngs: Dict[str, Any]):
    sample_1, sample_2 = samples
    rng = rngs["random"]
    return rng.random() * sample_1 + rng.random() * sample_2

parallel_dataset = ParallelStreamingDataset(
    [dset_1, dset_2],
    transform=transform
)

dataloader = StreamingDataLoader(parallel_dataset)

LitData uses an adaptive concurrency model to maximize throughput (samples/s) while avoiding network or prefix congestion. Instead of requiring users to manually tune max_concurrent_downloads for every combination of num_workers and file size, LitData implements a size-aware budget based on Little's Law.

Concurrency Logic

1. Automatic Mode (max_concurrent_downloads=None): LitData calculates an aggregate_budget based on a target bandwidth (100 MiB/s) and the median file size. This budget is then distributed across workers:

   • If num_workers <= 1: The budget is capped at 128.
   • If num_workers > 1: The budget is calculated as max(8, aggregate_budget // num_workers).
   • The budget is clamped between 32 and 512.

2. Explicit Mode (max_concurrent_downloads=int): If you provide an integer, LitData uses exactly that many permits and disables the adaptive clamping logic.

Key Concepts

• Division of Labor: Clients (like Botocore) handle rate-limiting/retries for 503/SlowDown errors. LitData manages concurrency and prefetch depth to prevent these errors from occurring in the first place.
• Bandwidth vs. Latency Bounding: For small files, the budget is latency-bounded. For large files, the budget is bandwidth-bounded to prevent too many multi-GB files from being in flight simultaneously.

# Automatic adaptive concurrency (default)
# LitData will calculate permits based on median file size and num_workers
max_concurrent_downloads = None

# Explicit concurrency control
# LitData will use exactly 64 permits total
max_concurrent_downloads = 64

Shuffling in LitData is deterministic and designed for distributed training. It works by assigning chunks to ranks/workers and then permuting items within those chunks. The permutation is based on the seed, the epoch, and chunk metadata.

Key Parameters:

• seed: Controls the permutation. Keep this stable when resuming training to ensure the same data order.
• drop_last: When True, ensures all ranks/workers have the same length (default is True under DDP).
• shuffle: Controls whether chunks and items are permuted.

Resuming: You can resume training using loader.state_dict() and load_state_dict(). If you want to ignore checkpointed shuffle settings during a resume, set force_override_state_dict=True on the dataset.

from litdata import StreamingDataset, StreamingDataLoader

train = StreamingDataset(
    "s3://my-bucket/train",
    shuffle=True,
    drop_last=True,  # keep every rank/worker at the same length (default True under DDP)
    seed=42,         # default is 42; keep stable when resuming
)
loader = StreamingDataLoader(train, batch_size=64, num_workers=8)

# shuffle=/drop_last= on the loader override the dataset
loader = StreamingDataLoader(train, batch_size=64, shuffle=True, drop_last=True)

When working in Lightning Studios, use /teamspace/... paths to enable direct bucket I/O. This is significantly faster than reading through the FUSE mount because LitData resolves these paths to the underlying object store URL (e.g., s3:// or gs://) and uses direct connections.

Path Mapping in Studios:

Important Notes:

• Credentials: Outside of Studio, /teamspace/... paths will not resolve; use standard cloud URIs (s3://, etc.) with storage_options.
• Optimization: Using optimize with a /teamspace/... output directory uploads chunks directly to the bucket.
• Distributed Jobs: When using num_nodes with optimize or map, LitData launches a Studio job. Local /this_studio outputs are sent to job artifacts.

from litdata import StreamingDataset, StreamingRawDataset, optimize

# Stream optimized data from an attached S3 connection (direct bucket download)
dataset = StreamingDataset("/teamspace/s3_connections/my-data-1/fast_data", shuffle=True, drop_last=True)

# Stream raw files from a connection
raw = StreamingRawDataset("/teamspace/s3_connections/my-bucket-1/raw")

# Optimize *into* a connection — chunks upload straight to the bucket
def should_keep(data):
    if data % 2 == 0:
        yield data

if __name__ == "__main__":
    optimize(
        fn=should_keep,
        inputs=list(range(1000)),
        output_dir="/teamspace/s3_connections/my-data-1/output",
        chunk_bytes="64MB",
        num_workers=1,
    )

To handle interruptions during long training runs, StreamingDataLoader is stateful. You can save its state using state_dict() and restore it using load_state_dict(). This allows you to pick up exactly where you left off in the dataset iteration.

Note: This is highly recommended for large-scale pre-training to recover from network or CUDA errors.

import os
import torch
from litdata import StreamingDataset, StreamingDataLoader

dataset = StreamingDataset("s3://my-bucket/my-data", shuffle=True)
dataloader = StreamingDataLoader(dataset, num_workers=os.cpu_count(), batch_size=64)

# Restore the dataLoader state if it exists
if os.path.isfile("dataloader_state.pt"):
    state_dict = torch.load("dataloader_state.pt")
    dataloader.load_state_dict(state_dict)

# Iterate over the data
for batch_idx, batch in enumerate(dataloader):
    # Store the state every 1000 batches
    if batch_idx % 1000 == 0:
        torch.save(dataloader.state_dict(), "dataloader_state.pt")

To achieve up to 20x faster training, follow a three-step workflow to convert raw data into an optimized chunked binary format.

Step 1: Optimize the data

Use ld.optimize to transform raw data into efficient chunks.

Important constraints for the mapping function (fn):

• Stable Keys/Types: Keys and types must remain stable across all samples. List lengths and types must be fixed.
• Chunking Configuration: You must provide exactly one of chunk_bytes or chunk_size.
• Data Format: For images, prefer returning compressed formats like JPEG to avoid the 10x+ size increase of uncompressed PIL RAW data.

import io
import numpy as np
from PIL import Image
import litdata as ld

def my_mapping_fn(index):
    # Example: creating a dummy JPEG image
    img = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="JPEG", quality=95)
    buf.seek(0)
    jpeg_image = Image.open(buf)
    
    return {"index": index, "image": jpeg_image, "class": np.random.randint(10)}

ld.optimize(
    fn=my_mapping_fn,
    inputs=list(range(1000)),
    output_dir="fast_data",
    num_workers=4,
    chunk_bytes="64MB"  # Use either chunk_bytes or chunk_size
)

Step 2: Upload to Cloud

Upload the output_dir to your cloud storage (e.g., S3):

aws s3 cp --recursive fast_data s3://my-bucket/fast_data

Step 3: Stream during training

Use ld.StreamingDataset and ld.StreamingDataLoader to load the optimized data.

import litdata as ld

dataset = ld.StreamingDataset(
    's3://my-bucket/fast_data',
    shuffle=True,
    drop_last=True,  # Recommended for multi-GPU to ensure consistent length
    seed=42,
)

def collate_fn(batch):
    return {
        "image": [sample["image"] for sample in batch],
        "class": [sample["class"] for sample in batch],
    }

dataloader = ld.StreamingDataLoader(dataset, batch_size=64, collate_fn=collate_fn)
for sample in dataloader:
    img, cls = sample["image"], sample["class"]

import litdata as ld

dataset = ld.StreamingDataset(
    's3://my-bucket/fast_data',
    shuffle=True,
    drop_last=True,
    seed=42,
)

dataloader = ld.StreamingDataLoader(dataset, batch_size=64)

When calling ld.optimize(), setting keep_data_ordered=False enables a shared queue. This helps balance the load across multiple workers, which is particularly beneficial if some workers are slower than others or if you want to improve fault tolerance against OOM errors.

Note: This option impacts the optimization time, not the subsequent streaming speed.

import numpy as np
from PIL import Image
import litdata as ld

def random_images(index):
    fake_images = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
    fake_labels = np.random.randint(10)
    return {"index": index, "image": fake_images, "class": fake_labels}

if __name__ == "__main__":
    ld.optimize(
        fn=random_images,
        inputs=list(range(1000)),
        output_dir="fast_data",
        num_workers=4,
        chunk_bytes="64MB",
        keep_data_ordered=False,  # Enables shared queue
    )