2nd/10_Wiki/Topics/AI_and_ML/Just-in-time-Data-Loading.md

---
id: wiki-2026-0508-just-in-time-data-loading
title: Just-in-time Data Loading
category: 10_Wiki/Topics
status: verified
canonical_id: self
aliases: [JIT Data Loading, Lazy Loading, Streaming Datasets, On-demand Loading]
duplicate_of: none
source_trust_level: A
confidence_score: 0.9
verification_status: applied
tags: [data-loading, pytorch, huggingface, mmap, streaming, ml-infra]
raw_sources: []
last_reinforced: 2026-05-10
github_commit: pending
tech_stack:
  language: python
  framework: pytorch-huggingface
---

# Just-in-time Data Loading

## 매 한 줄
> **"매 모든 데이터를 메모리에 올리는 시대는 끝났다 — 필요할 때 필요한 만큼"**. JIT/lazy data loading 은 데이터 전체를 RAM 으로 미리 적재하지 않고 학습/추론 시점에 부분만 fetch 하는 패턴으로, PyTorch DataLoader streaming, HuggingFace datasets streaming, mmap, Arrow IPC 등이 TB-PB 규모 학습의 표준 도구가 되었다.

## 매 핵심

### 매 핵심 기법
- **mmap**: 파일을 가상 메모리에 매핑 — OS page cache 활용.
- **Arrow / Parquet**: 컬럼 기반, zero-copy slice.
- **Streaming dataset**: HF `streaming=True`, IterableDataset.
- **Sharding + shuffling**: 큰 shuffle buffer 없이 shard-level shuffle.
- **Async prefetch**: DataLoader workers + pinned memory.
- **WebDataset**: tar shard streaming.
- **MosaicML Streaming (StreamingDataset)**: cloud-native, fast resume.
- **Ray Data**: 분산 lazy.
- **NVIDIA DALI**: GPU 측 디코드/증강.

### 매 병목 지점
- IO bandwidth (네트워크/디스크).
- Decompression CPU.
- Decoding (image/audio).
- Augmentation CPU.
- Host → Device 전송 (PCIe).

### 매 응용
1. ImageNet/LAION TB scale 학습.
2. LLM pretraining (수 PB tokens).
3. 동영상/오디오 모델 학습.
4. Inference batch (대규모 evaluation).

## 💻 패턴

### 1. PyTorch IterableDataset (스트리밍 라인)
```python
import torch
from torch.utils.data import IterableDataset, DataLoader

class JsonlStream(IterableDataset):
    def __init__(self, path): self.path = path
    def __iter__(self):
        worker = torch.utils.data.get_worker_info()
        with open(self.path) as f:
            for i, line in enumerate(f):
                if worker is None or i % worker.num_workers == worker.id:
                    yield json.loads(line)

loader = DataLoader(JsonlStream("data.jsonl"),
                    batch_size=32, num_workers=4, pin_memory=True, prefetch_factor=4)
```

### 2. HuggingFace `datasets` streaming
```python
from datasets import load_dataset
ds = load_dataset("c4", "en", split="train", streaming=True)
ds = ds.shuffle(buffer_size=10_000, seed=0).take(1_000_000)
for ex in ds:
    yield tokenize(ex["text"])
# 디스크 다운로드 없음, network 에서 chunk-by-chunk.
```

### 3. mmap 으로 큰 numpy 배열 로딩
```python
import numpy as np
arr = np.memmap("embeddings.f16", dtype=np.float16, mode="r",
                shape=(50_000_000, 1024))
# slice 만 OS page-in
batch = np.array(arr[indices])  # copy out
```

### 4. Parquet zero-copy iterator (pyarrow)
```python
import pyarrow.parquet as pq
pf = pq.ParquetFile("shard.parquet")
for batch in pf.iter_batches(batch_size=8192, columns=["text", "label"]):
    yield batch.to_pydict()
```

### 5. WebDataset tar shards
```python
import webdataset as wds
url = "pipe:aws s3 cp s3://bucket/shard-{000000..001023}.tar -"
ds = (wds.WebDataset(url, shardshuffle=True)
        .shuffle(1000)
        .decode("pil")
        .to_tuple("jpg", "cls")
        .batched(64))
loader = wds.WebLoader(ds, num_workers=8)
```

### 6. MosaicML StreamingDataset (resume-safe)
```python
from streaming import StreamingDataset
ds = StreamingDataset(
    remote="s3://my-bucket/shards", local="/tmp/cache",
    shuffle=True, batch_size=64, predownload=2_000,
)
loader = torch.utils.data.DataLoader(ds, batch_size=64, num_workers=8)
# epoch 중단/재시작 시 정확히 같은 sample sequence 보장.
```

### 7. NVIDIA DALI GPU 디코드
```python
from nvidia.dali import pipeline_def, fn, types
@pipeline_def(batch_size=128, num_threads=4, device_id=0)
def pipe():
    jpegs, labels = fn.readers.file(file_root="/data")
    images = fn.decoders.image(jpegs, device="mixed")
    images = fn.resize(images, size=224)
    return images, labels
p = pipe(); p.build()
```

### 8. Ray Data (분산 lazy)
```python
import ray
ds = (ray.data.read_parquet("s3://bucket/")
        .map(tokenize)
        .iter_torch_batches(batch_size=1024, prefetch_batches=4))
for batch in ds:
    train_step(batch)
```

### 9. 동기/비동기 prefetch (CUDA streams)
```python
import torch
class Prefetcher:
    def __init__(self, loader, device):
        self.loader, self.device = iter(loader), device
        self.stream = torch.cuda.Stream()
        self._next()
    def _next(self):
        try: self.batch = next(self.loader)
        except StopIteration: self.batch = None; return
        with torch.cuda.stream(self.stream):
            self.batch = {k: v.to(self.device, non_blocking=True) for k,v in self.batch.items()}
    def get(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        b = self.batch; self._next(); return b
```

### 10. profile-driven 튜닝 체크리스트
```python
# 1) torch.profiler 로 dataloader vs compute 시간 측정
# 2) num_workers: CPU core 의 1-2x 부터 → GPU util 95%+ 까지 증가
# 3) prefetch_factor: 2 → 4 → 8
# 4) pin_memory=True, non_blocking=True
# 5) 디코드가 CPU bound → DALI/torchcodec 로 GPU 이동
# 6) 네트워크 bound → shard 크기 64-256MB, 동시 connection 다중화
```

## 매 결정 기준
| 상황 | Approach |
|---|---|
| 데이터 < RAM | 일반 in-memory Dataset 도 OK |
| 데이터 > RAM, 단일 노드 | mmap 또는 Parquet iter |
| TB+ cloud, resume 중요 | MosaicML Streaming 또는 WebDataset |
| LLM pretraining (PB) | tokenized shards + StreamingDataset + global shuffle |
| 이미지 / 비디오 디코드 bottleneck | DALI 또는 torchcodec |
| 분산 + lazy transform | Ray Data |

**기본값**: > 100 GB 데이터는 streaming + sharded shuffle 이 기본. < 그 이하는 mmap + 일반 Dataset.

## 🔗 Graph
- Adjacent: [[Parquet]]

## 🤖 LLM 활용
**언제**: DataLoader 코드 변환 (in-memory → streaming), shard schema 설계, prefetch 튜닝 체크리스트 생성.
**언제 X**: 실제 GPU util 측정 / profiler 트레이스 — 직접 실행해야 truth.

## ❌ 안티패턴
- **streaming + 작은 shuffle buffer**: in-shard 순서 잔존 → 학습 편향.
- **num_workers=0 + 큰 데이터**: 메인 스레드 IO block — GPU 0% util.
- **shard 너무 작음 (< 16MB)**: object-store 호출 폭주.
- **shard 너무 큼 (> 1GB)**: resume / shuffle 비효율.
- **pin_memory + CPU dataset**: 의미 없음. GPU 학습일 때만.
- **augmentation on CPU 만**: GPU starvation — DALI 검토.

## 🧪 검증 / 중복
- Verified (PyTorch DataLoader docs, HF datasets streaming guide, MosaicML Streaming docs, NVIDIA DALI docs 2026).
- 신뢰도 A.

## 🕓 Changelog
| 날짜 | 변경 |
|---|---|
| 2026-05-08 | Phase 1 |
| 2026-05-10 | Manual cleanup — streaming/mmap/DALI/MosaicML 패턴 |