Lightning-AI/litData

LitData – PyTorch용 빠르고 클라우드 기반 데이터 파이프라인

무엇인가요 – LitData(라이트닝 AI 제작)는 거대한 학습 데이터셋을 빠르고 신뢰성 있게 로드할 수 있도록 도와주는 파이썬 라이브러리입니다. 이는 클라우드 스토리지(S3, GCS, Azure, 허깅페이스 허브 등)에서 직접 원시 파일을 스트리밍하거나, 한 번만 데이터를 컴팩트한 청크 기반 바이너리 형식으로 변환하여, 단순한 torch.utils.data.DataLoader보다 최대 20배 빠른 속도로 스트리밍할 수 있도록 합니다.

왜 중요한가요 – 현대 딥러닝 프로젝트에서는 I/O가 종종 버퍼넥이 됩니다. 데이터셋은 테라바이트 규모일 수 있고, 수백만 개의 작은 파일로 구성되거나, 원격 객체 스토리지에만 존재할 수 있습니다. LitData는 세 가지 주요 문제를 해결합니다:

준비 없이 스트리밍 – StreamingRawDataset는 완전히 비동기적이고 배치 다운로드를 지원하며 내장된 재시도 로직을 갖추고 있습니다. 로컬에 전체 데이터셋을 다운로드할 필요가 없습니다.
한 번의 최적화 – ld.optimize는 데이터셋을 LitData 고유의 청크 기반 바이너리 형식(또는 Parquet, MDS 등과도 호환)으로 재작성합니다. 생성된 청크는 로컬 및 버킷에 캐시되며, 재시작 가능하고 샤프링 인식 스트리밍이 가능해 최대 20배 빠른 성능을 제공합니다.
분산 대응 API – 라이브러리는 StreamingDataset/StreamingDataLoader(PyTorch Lightning, Fabric, Hugging Face와 호환)를 제공하며, map 프리미티브를 통해 이미지 리사이징, 임베딩 생성, 웹 스크래핑 등 임의의 전처리를 여러 머신에서 병렬로 실행할 수 있습니다.

핵심 개념 및 API

개념	대표 클래스/함수	기능
원시 스트리밍	`StreamingRawDataset`	로컬 또는 클라우드 경로에서 이미지, 오디오, 텍스트 등 어떤 파일도 원시 `bytes`로 읽습니다. `transform` 콜백을 제공하면 실시간으로 디코딩할 수 있습니다.
최적화된 스트리밍	`StreamingDataset` + `StreamingDataLoader`	`ld.optimize`로 사전 변환된 데이터를 로드합니다. 샤프링, drop-last, 멀티 GPU 재시작, 키 기반 검색을 처리합니다.
한 번의 변환	`ld.optimize(...)`	샘플을 생성하는 파이썬 함수를 받아, LitData의 청크 기반 바이너리 형식(청크 크기 설정 가능)에 쓰고, 선택적으로 랜덤 액세스를 위한 키 인덱스를 구축합니다.
병렬 전처리	`ld.map(...)`	사용자 제공 함수를 입력 리스트에 대해 병렬로 실행하고 결과를 목적지(로컬 또는 클라우드)에 씁니다. 이미지 리사이징, 임베딩 생성, 웹 스크래핑 등에 적합합니다.
허깅페이스 통합	`ld.optimize_hf(...)` / `StreamingDataset("hf://…")`	HF 허브에서 직접 데이터셋을 스트리밍하거나, 한 번만 LitData 청크로 변환하여 더 빠른 학습을 가능하게 합니다.

빠른 시작 (설치 및 기본 사용)

pip install litdata                # 핵심 패키지
pip install 'litdata[extras]'      # 더 빠른 asyncio를 위한 uvloop 추가

1️⃣ 원시 파일 스트리밍 (전처리 단계 없음)

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)

기능: 비동기 배치 다운로드, 자동 재시도, 로컬 index.json.zstd 캐시, 모든 클라우드 제공업체와 호환.

2️⃣ 한 번 최적화 후 최고 속도로 스트리밍

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")
# 폴더를 클라우드에 업로드, 예:
# 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"]은 Image 객체 리스트, batch["label"]은 정수 리스트
    train_step(batch)

결과: 동일한 원시 파일에 대해 기존 torch.utils.data.DataLoader보다 최대 20배 빠른 에포크 처리 시간.

LitData vs. 대안 사용 시기

LitData – 클라우드 독립적 스트리밍, 재시작 가능한 에포크, 대규모 샤프링, 한 번의 전처리로 여러 실험에 재사용하고 싶을 때 최적.
torchdata – 저수준 반복자 프리미티브 제공하지만, 내장된 저장 형식이나 청크 수준 캐시 없음. 단순한 파일 목록만 필요하고, 다른 모든 부분을 직접 구축하고 싶을 때 사용.
허깅페이스 datasets 스트리밍 – 빠른 프로토타이핑에 편리. LitData의 optimize_hf는 같은 데이터셋을 장기간 학습에 더 빠르게 만들 수 있습니다.

지원하는 데이터 모달리티

LitData는 쓰기/읽기 시 타입 정보를 유지하는 가벼운 래퍼를 제공합니다:

텍스트 / 토큰 – Text, Tokens
이미지 / JPEG / PIL – Image, Jpeg, Pil
오디오 / 비디오 – Audio, Video
3D 메시, Nifti 볼륨 – Mesh, Nifti
일반 파일 – File, Pdf
배열 및 텐서 – Numpy, Tensor
그래프 (PyG) – Graph
Parquet 테이블 – Parquet 이 래퍼들은 원시 바이트, NumPy 배열, PyTorch 텐서를 저장하고, 학습 시 동일한 파이썬 객체로 가져올 수 있게 합니다.

생태계 및 커뮤니티

Lightning Cloud – 원활한 통합; Lightning의 관리형 GPU 클러스터나 온프레미스 머신에서 LitData 파이프라인을 실행 가능.
Discord – 활발한 도움 채널 (https://discord.gg/VptPCZkGNa).
AI 에이전트 스킬 – Vercel 스타일의 '스킬' 제공으로 코드 생성 에이전트(Claude, Cursor 등)가 LitData API를 자동 완성할 수 있음.
사용 중인 곳 – Lightning Cloud에서 34만 명 이상의 개발자, 내부 연구 팀, 여러 공개 ML 프로젝트에서 사용 (리포지토리의 "Used by" 섹션에 링크 있음).

TL;DR

LitData는 PyTorch용 프로덕션 수준의 데이터 로딩 라이브러리로, 다음을 제공합니다:

어떤 클라우드 스토어에서든 비동기 배치 I/O로 원시 파일 스트리밍.
한 번만 데이터셋을 빠르고 청크 기반 바이너리 형식으로 변환하여 샤프링, 재시작, 키 기반 액세스 가능.
ld.map을 통해 머신 간 임의의 전처리 병렬화.
Drop-in 클래스인 StreamingDataset/StreamingDataLoader로 Lightning, 허깅페이스, PyTorch-Lightning 워크플로우에 통합.

학습 작업이 데이터 대기 시간에 많은 시간을 소비한다면, LitData는 그 대기 시간을 크게 줄이고 대규모 클라우드 기반 데이터셋의 엔지니어링을 단순화할 수 있습니다.