Training Data Pipeline: Streaming Tokenization at Scale

BASE_URL = "https://huggingface.co/datasets/karpathy/fineweb-edu-100b-shuffle/resolve/main"
MAX_SHARD = 1822  # 1,823 files total (0-indexed)
DATA_DIR = os.path.join(base_dir, "base_data")
 
def list_parquet_files(data_dir=None):
    """List all parquet files in the data directory."""
    data_dir = DATA_DIR if data_dir is None else data_dir
    parquet_files = sorted([
        f for f in os.listdir(data_dir)
        if f.endswith('.parquet') and not f.endswith('.tmp')
    ])
    return [os.path.join(data_dir, f) for f in parquet_files]

def parquets_iter_batched(split, start=0, step=1):
    """Iterate through dataset in batches of row groups."""
    assert split in ["train", "val"]
    parquet_paths = list_parquet_files()
    
    # Last file = validation, rest = training
    parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]
    
    for filepath in parquet_paths:
        pf = pq.ParquetFile(filepath)
        for rg_idx in range(start, pf.num_row_groups, step):
            rg = pf.read_row_group(rg_idx)
            texts = rg.column('text').to_pylist()
            yield texts

pub fn train_from_iterator(
    &mut self,
    iterator: &PyAny,
    vocab_size: u32,
    buffer_size: usize,
    pattern: Option<String>,
) -> PyResult<()> {
    // Use GPT-4 style regex pattern
    let pattern_str = pattern.unwrap_or_else(|| GPT4_PATTERN.to_string());
    
    // Global chunk counts
    let mut counts: AHashMap<CompactString, i32> = AHashMap::new();
    let mut buf: Vec<String> = Vec::with_capacity(buffer_size);
    
    // Stream ingestion: refill under GIL, process without GIL (parallel)
    loop {
        let exhausted = refill(&mut buf)?;
        
        // Release GIL and process in parallel with rayon
        let local: AHashMap<CompactString, i32> = py.allow_threads(|| {
            buf.par_iter()
                .map(|s| {
                    let mut m: AHashMap<CompactString, i32> = AHashMap::new();
                    for mat in pattern.find_iter(s) {
                        let piece = mat.expect("regex match failed").as_str();
                        *m.entry(CompactString::from(piece)).or_default() += 1;
                    }
                    m
                })
                .reduce(|| AHashMap::new(), |mut a, b| {
                    for (k, v) in b {
                        *a.entry(k).or_default() += v;
                    }
                    a
                })
        });
        
        // Merge local into global
        for (k, v) in local {
            *counts.entry(k).or_default() += v;
        }
        
        if exhausted { break; }
    }
    
    // Train BPE on the collected statistics
    self.train_core_incremental(words, cvec, vocab_size);
}

class RustBPETokenizer:
    def __init__(self, enc, bos_token):
        self.enc = enc  # tiktoken.Encoding
        self.bos_token_id = self.encode_special(bos_token)
    
    def encode(self, text, prepend=None, append=None, num_threads=8):
        """Encode text using tiktoken's optimized C++ implementation."""
        if prepend is not None:
            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
        if append is not None:
            append_id = append if isinstance(append, int) else self.encode_special(append)
        
        if isinstance(text, str):
            ids = self.enc.encode_ordinary(text)
            if prepend is not None:
                ids.insert(0, prepend_id)
            if append is not None:
                ids.append(append_id)
        elif isinstance(text, list):
            # Batch encoding with multi-threading
            ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
            if prepend is not None:
                for ids_row in ids:
                    ids_row.insert(0, prepend_id)
            if append is not None:
                for ids_row in ids:
                    ids_row.append(append_id)
        
        return ids

SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""

def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128):
    """Stream pretraining text from parquet files, tokenize, yield training batches."""
    assert split in ["train", "val"]
    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
    
    needed_tokens = B * T + 1  # +1 for target at last position
    
    # Get tokenizer and BOS token
    tokenizer = get_tokenizer()
    bos_token = tokenizer.get_bos_token_id()
    
    # Token buffer streams tokens on the right, pops from the left
    token_buffer = deque()
    scratch = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=True)
    
    # Infinite iterator over document batches
    def document_batches():
        while True:
            # Distributed sharding: each rank processes every Nth row group
            for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size):
                # Further sub-batch for tokenizer
                for i in range(0, len(batch), tokenizer_batch_size):
                    yield batch[i:i+tokenizer_batch_size]
    
    batches = document_batches()
    
    batch_index = 0
    while True:
        # Accumulate enough tokens for one iteration
        while len(token_buffer) < needed_tokens:
            doc_batch = next(batches)
            token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
            for tokens in token_lists:
                token_buffer.extend(tokens)
            batch_index += 1
        
        # Move tokens from deque into scratch buffer
        for i in range(needed_tokens):
            scratch[i] = token_buffer.popleft()
        
        # Create inputs/targets as 1D tensors
        inputs_cpu = scratch[:-1].to(dtype=torch.int32)
        targets_cpu = scratch[1:]
        
        # Reshape to 2D and move to GPU async
        inputs = inputs_cpu.view(B, T).to(device="cuda", dtype=torch.int32, non_blocking=True)
        targets = targets_cpu.view(B, T).to(device="cuda", dtype=torch.int64, non_blocking=True)
        
        yield inputs, targets

# Each rank processes every Nth row group
for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size):
    ...

token_buffer = deque()  # Stream tokens on the right, pop from the left
 
while len(token_buffer) < needed_tokens:
    doc_batch = next(batches)
    token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
    for tokens in token_lists:
        token_buffer.extend(tokens)  # Concatenate all tokens

# Stage 1: Parquet row groups (~1024 documents)
for batch in parquets_iter_batched(...):
    
    # Stage 2: Tokenizer sub-batches (128 documents)
    for i in range(0, len(batch), tokenizer_batch_size):
        yield batch[i:i+tokenizer_batch_size]

# Pinned CPU memory for fast GPU transfer
scratch = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=True)
 
# ... fill scratch buffer ...
 
inputs_cpu = scratch[:-1].to(dtype=torch.int32)
targets_cpu = scratch[1:]
 
# Async GPU transfer (non-blocking)
inputs = inputs_cpu.view(B, T).to(device="cuda", dtype=torch.int32, non_blocking=True)
targets = targets_cpu.view(B, T).to(device="cuda", dtype=torch.int64, non_blocking=True)

def document_batches():
    while True:  # Infinite loop
        for batch in parquets_iter_batched(...):
            ...

train_loader = tokenizing_distributed_data_loader(device_batch_size, max_seq_len, split="train")
x, y = next(train_loader)  # Prefetch first batch
 
for step in range(num_iterations):
    loss = model(x, y)
    loss.backward()
    # ... optimizer steps ...
    x, y = next(train_loader)  # Prefetch next batch while GPU is busy

token_buffer = deque()  # Typical size: ~100K tokens × 8 bytes = 800 KB

needed_tokens = B * T + 1  # e.g., 32 × 2048 + 1 = 65,537
scratch = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=True)
# Size: 65,537 × 8 bytes = 524 KB

# Model
depth = 20
max_seq_len = 2048
device_batch_size = 32
total_batch_size = 524288  # tokens
 
# Data loading
train_loader = tokenizing_distributed_data_loader(
    device_batch_size, 
    max_seq_len, 
    split="train",
    tokenizer_threads=4,
    tokenizer_batch_size=128
)

train_loader = tokenizing_distributed_data_loader(device_batch_size, max_seq_len, split="train")
x, y = next(train_loader)  # Kick off load of the very first batch
 
for step in range(num_iterations + 1):
    # ... evaluation logic ...
    
    # Single training step
    torch.cuda.synchronize()
    t0 = time.time()
    
    for micro_step in range(grad_accum_steps):
        with autocast_ctx:
            loss = model(x, y)
        loss = loss / grad_accum_steps
        loss.backward()
        
        # Prefetch next batch while GPU is busy with forward/backward
        x, y = next(train_loader)
    
    # Gradient clipping
    if grad_clip > 0.0:
        torch.nn.utils.clip_grad_norm_(orig_model.parameters(), grad_clip)
    
    # Optimizer steps
    for opt in optimizers:
        opt.step()
    model.zero_grad(set_to_none=True)
    
    torch.cuda.synchronize()
    t1 = time.time()
    dt = t1 - t0

loss.backward()
x, y = next(train_loader)  # Prefetch DURING backward pass

tokens_per_fwdbwd = device_batch_size * max_seq_len
world_tokens_per_fwdbwd = tokens_per_fwdbwd * ddp_world_size
grad_accum_steps = total_batch_size // world_tokens_per_fwdbwd

def build_val_loader():
    return tokenizing_distributed_data_loader(device_batch_size, max_seq_len, split="val")
 
# During evaluation
val_loader = build_val_loader()
eval_steps = eval_tokens // (device_batch_size * max_seq_len * ddp_world_size)
 
with autocast_ctx:
    val_bpb = evaluate_bpb(model, val_loader, eval_steps, token_bytes)

def download_single_file(index):
    """Downloads a single file with exponential backoff retry."""
    filename = index_to_filename(index)
    filepath = os.path.join(DATA_DIR, filename)
    
    if os.path.exists(filepath):
        print(f"Skipping {filepath} (already exists)")
        return True
    
    url = f"{BASE_URL}/{filename}"
    max_attempts = 5
    
    for attempt in range(1, max_attempts + 1):
        try:
            response = requests.get(url, stream=True, timeout=30)
            response.raise_for_status()
            
            # Write to temporary file first
            temp_path = filepath + f".tmp"
            with open(temp_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        f.write(chunk)
            
            # Move temp file to final location (atomic)
            os.rename(temp_path, filepath)
            return True
            
        except (requests.RequestException, IOError) as e:
            print(f"Attempt {attempt}/{max_attempts} failed: {e}")
            # Exponential backoff
            if attempt < max_attempts:
                wait_time = 2 ** attempt
                time.sleep(wait_time)
    
    return False

# Download first 10 shards with 4 parallel workers
python -m nanochat.dataset -n 10 -w 4
 
# Download entire dataset (1,823 shards)
python -m nanochat.dataset -n -1 -w 8

import time
from nanochat.dataloader import tokenizing_distributed_data_loader
 
B, T = 32, 2048
loader = tokenizing_distributed_data_loader(B, T, "train", tokenizer_threads=4)
 
# Warmup
for _ in range(10):
    next(loader)
 
# Benchmark
t0 = time.time()
num_batches = 100
for _ in range(num_batches):
    x, y = next(loader)
t1 = time.time()
 
tokens_per_batch = B * T
total_tokens = tokens_per_batch * num_batches
throughput = total_tokens / (t1 - t0)
print(f"Throughput: {throughput/1e6:.2f}M tokens/sec")

import torch
import os
import psutil
 
process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / 1024**2  # MB
 
# Create loader and run 100 steps
loader = tokenizing_distributed_data_loader(32, 2048, "train")
for _ in range(100):
    x, y = next(loader)
    
mem_after = process.memory_info().rss / 1024**2  # MB
print(f"Memory increase: {mem_after - mem_before:.2f} MB")

# Precompute approach (NOT used)
for shard in parquet_files:
    tokens = tokenize(shard)
    save(tokens, f"tokens_{shard}.pt")

# Padding approach (NOT used)
for doc in documents:
    tokens = tokenize(doc)
    if len(tokens) < T:
        tokens += [PAD] * (T - len(tokens))
    yield tokens[:T]

# Shard assignment approach (NOT used)
shards_per_rank = len(parquet_files) // world_size
my_shards = parquet_files[rank * shards_per_rank : (rank+1) * shards_per_rank]