Lightning-AI/litData

LitData – Fast, cloud‑native data pipelines for PyTorch

What it is – LitData (by Lightning AI) is a Python library that makes loading huge training datasets fast and reliable. It does this by either streaming raw files directly from cloud storage (S3, GCS, Azure, Hugging Face Hub, etc.) or by converting the data once into a compact, chunked binary format that can be streamed at up to 20 × the speed of a naïve torch.utils.data.DataLoader.

Why it matters – In modern deep‑learning projects the bottleneck is often I/O: datasets may be terabytes in size, consist of millions of tiny files, or live only in remote object stores. LitData solves three pain points:

Zero‑prep streaming – StreamingRawDataset pulls files as raw bytes with fully asynchronous, batched downloads and built‑in retry logic. No need to download the whole dataset locally.
One‑time optimization – ld.optimize rewrites a dataset into LitData’s own chunked binary format (or can work with Parquet, MDS, etc.). The resulting chunks are cached locally and on the bucket, enabling resumable, shuffling‑aware streaming that can be up to 20 × faster.
Distributed‑ready API – The library ships StreamingDataset/StreamingDataLoader (compatible with PyTorch Lightning, Fabric, and Hugging Face) and a map primitive that lets you run arbitrary preprocessing (e.g., resizing images, creating embeddings, web‑scraping) across many machines.

Core concepts & API

Concept	Typical class / function	What it does
Raw streaming	`StreamingRawDataset`	Reads any file (image, audio, text, etc.) from local or cloud paths as raw `bytes`. You can supply a `transform` callable to decode on the fly.
Optimized streaming	`StreamingDataset` + `StreamingDataLoader`	Loads data that has been pre‑converted with `ld.optimize`. Handles shuffling, drop‑last, multi‑GPU resume, and key‑based look‑ups.
One‑time conversion	`ld.optimize(...)`	Takes a Python function that yields samples, writes them into LitData’s chunked binary format (configurable chunk size), and optionally builds a key index for random access.
Parallel preprocessing	`ld.map(...)`	Executes a user‑provided function over a list of inputs in parallel, writing the results to a destination (local or cloud). Great for image resizing, embedding generation, web scraping, etc.
Hugging Face integration	`ld.optimize_hf(...)` / `StreamingDataset("hf://…")`	Streams datasets directly from the HF Hub, or converts them once into LitData chunks for faster training.

Quick start (install & basic usage)

pip install litdata                # core package
pip install 'litdata[extras]'      # adds optional uvloop for faster asyncio

1️⃣ Stream raw files (no preprocessing step)

from litdata import StreamingRawDataset
from torch.utils.data import DataLoader
from PIL import Image
import io

ds = StreamingRawDataset(
    "s3://my-bucket/raw-images/",
    transform=lambda b: Image.open(io.BytesIO(b)).convert("RGB"),
)
loader = DataLoader(ds, batch_size=32, num_workers=8)
for batch in loader:
    train_step(batch)

Features: async batched downloads, automatic retries, local index.json.zstd cache, works with any cloud provider.

2️⃣ Optimize once, then stream at max speed

import litdata as ld, numpy as np

def make_sample(i):
    img = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8)
    return {"index": i, "image": ld.Image(array=img, quality=95, format="jpeg"), "label": np.random.randint(10)}

ld.optimize(fn=make_sample, inputs=list(range(1000)), output_dir="fast_data", chunk_bytes="64MB")
# upload the folder to cloud, e.g.:
# aws s3 cp --recursive fast_data s3://my-bucket/fast_data

ds = ld.StreamingDataset('s3://my-bucket/fast_data', shuffle=True, seed=42)
loader = ld.StreamingDataLoader(ds, batch_size=64)
for batch in loader:
    # batch["image"] is a list of Image objects, batch["label"] a list of ints
    train_step(batch)

Result: up to 20 × faster epoch times compared with a vanilla torch.utils.data.DataLoader on the same raw files.

When to use LitData vs. alternatives

LitData – best when you need cloud‑agnostic streaming, resume‑able epochs, large‑scale shuffling, or you want to pre‑process once and reuse the optimized format across many experiments.
torchdata – provides low‑level iterator primitives but no built‑in storage format or chunk‑level caching. Use it if you only need simple file listing and want to build every other piece yourself.
Hugging Face datasets streaming – convenient for quick prototypes; LitData’s optimize_hf can make those same datasets much faster for long‑running training.

Supported data modalities

LitData ships lightweight wrappers that preserve type information when writing/reading:

Text / Tokens – Text, Tokens
Images / Jpeg / Pil – Image, Jpeg, Pil
Audio / Video – Audio, Video
3‑D meshes, Nifti volumes – Mesh, Nifti
Generic files – File, Pdf
Arrays & tensors – Numpy, Tensor
Graphs (PyG) – Graph
Parquet tables – Parquet These wrappers let you store raw bytes, NumPy arrays, or PyTorch tensors and retrieve them with the same Python objects during training.

Ecosystem & community

Lightning Cloud – seamless integration; you can run LitData pipelines on Lightning’s managed GPU clusters or on‑premise machines.
Discord – active help channel (https://discord.gg/VptPCZkGNa).
AI‑agent skills – a Vercel‑style skill is provided so code‑generation agents (Claude, Cursor, etc.) can autocomplete the LitData API.
Used by – over 340 k developers on Lightning Cloud, internal research teams, and several public ML projects (links listed in the repo’s “Used by” section).

TL;DR

LitData is a production‑grade data‑loading library for PyTorch that lets you:

Stream raw files directly from any cloud store with async, batched I/O.
Convert datasets once into a fast, chunked binary format that can be shuffled, resumed, and accessed by key.
Parallelize arbitrary preprocessing across machines via ld.map.
Plug into Lightning, Hugging Face, and PyTorch‑Lightning workflows with drop‑in StreamingDataset/StreamingDataLoader classes.

If your training jobs spend a lot of time waiting for data, LitData can cut that wait dramatically and simplify the engineering around large‑scale, cloud‑based datasets.