Storage I/O Engineering: The Art of Letting Your Dataset Slow Down Training (and Prevention)

# === Storage I/O bench — DataLoader'ın gerçek throughput'unu ölç ===
import time, torch
from torch.utils.data import DataLoader, IterableDataset
from datasets import load_dataset
 
dataset = load_dataset("malhajar/alpaca-gpt4-tr", split="train")
# Tokenize ve format
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
 
def tokenize(ex):
    return tok(ex["text"], truncation=True, max_length=2048, padding=False)
 
dataset = dataset.map(tokenize, batched=True, num_proc=8)
dataset.set_format("torch", columns=["input_ids", "attention_mask"])
 
# Bench farklı config'ler
def bench(loader, n=200):
    it = iter(loader)
    next(it)  # warmup
    t0 = time.perf_counter()
    n_tokens = 0
    for i in range(n):
        batch = next(it)
        n_tokens += batch["input_ids"].numel()
    elapsed = time.perf_counter() - t0
    return n_tokens / elapsed   # tokens/sec
 
for nw, pf, pin in [(0, 2, False), (2, 2, True), (4, 2, True), (4, 4, True), (8, 4, True)]:
    loader = DataLoader(dataset, batch_size=8,
                        num_workers=nw, prefetch_factor=pf,
                        pin_memory=pin, persistent_workers=(nw > 0))
    tps = bench(loader)
    print(f"num_workers={nw:2d}  prefetch={pf}  pin={pin}  →  {tps/1e6:5.2f} M tokens/s")
 
# Typical RTX 4090 + Ryzen 7950X output:
#  num_workers=0  prefetch=2  pin=False  →   0.42 M tokens/s    ← CPU-bound
#  num_workers=2  prefetch=2  pin=True   →   2.10 M tokens/s
#  num_workers=4  prefetch=2  pin=True   →   3.85 M tokens/s
#  num_workers=4  prefetch=4  pin=True   →   4.20 M tokens/s    ← sweet spot
#  num_workers=8  prefetch=4  pin=True   →   3.90 M tokens/s    ← context switch overhead