diff --git a/README.md b/README.md index b1b32b2..0dd0146 100644 --- a/README.md +++ b/README.md @@ -165,7 +165,8 @@ data_loader: DistDataLoader = build_dataloader( seq_len=model.config.max_context, batch_size=32, base_url="karpathy/climbmix-400b-shuffle", # for starting point - max_shards=6542 # last shard id for the given dataset (if not provided, it will be computed by probing the server, which can take a while) + max_shards=6542, # last shard id for the given dataset (if not provided, it will be computed by probing the server, which can take a while) + packing_strategy="stream", # default ) ``` @@ -201,7 +202,13 @@ The data pipeline basically works with for any dataset available on internet tha ![sequence packing](./assets/sequence_packing.png) Sequence packing strategy. From [The Smol Training Playbook](https://huggingface.co/spaces/HuggingFaceTB/smol-training-playbook). -The `build_dataloader` function the data loader, which is accessible from [`gpt_lab.data.DistDataLoader`](./src/gpt_lab/data/loader.py). It employs a distributed streaming data pipeline over Parquet shards with on-the-fly tokenization and greedy document packing into fixed-length sequences. It creates cpu and gpu buffers to pre-load the data in contiguous memory, stream the data from local shards downloaded from the given dataset, and feed the model with the data with a packing strategy, to maximize the throughput of the training loop by avoiding the use of padding tokens. It also supports distributed training setups, and can be used with DDP or other distributed training frameworks. +The `build_dataloader` function, accessible through [`gpt_lab.data.DistDataLoader`](./src/gpt_lab/data/loader.py), exposes three explicit packing strategies: + +- `stream` (default) treats documents as one flat stream. It builds each batch from exactly `B*T+1` tokens and carries the final token into the next batch, so no source token or adjacent transition is silently lost. Rows are arbitrary views of that stream and need not start with BOS. +- `bos_aligned` is lossless and deterministic. Every row starts with BOS. A document suffix that crosses a row is retained and prioritized on the next row after a synthetic BOS. Full batches remain fully utilized, and finite sources stop before emitting an incomplete batch. `PackingStats` reports synthetic-BOS overhead. +- `bos_bestfit_crop` preserves the exact nanochat-style best-fit behavior. It keeps rows full and BOS-aligned by discarding suffixes when no complete document fits. This destructive crop rate can be substantial and is reported by `PackingStats`; its legacy resume position remains row-group approximate. `stream` and `bos_aligned` checkpoint the carry/continuation state exactly. + +These are context-layout policies, not document-isolation policies. A BOS token does **not** prevent causal attention from crossing document boundaries within the same row; GPT-Lab does not currently pass segment-aware attention masks to the model. The critical point regarding model training, is that we must make sure to have a good balance between loader time and model forward/backward time to avoid bottlenecks from the data loading process. Given that constraint, the implementated data loader is satisfying. @@ -986,4 +993,4 @@ END_SYSTEM_INSTRUCTION [nan]: https://aclanthology.org/2024.emnlp-main.40.pdf [nan]: https://aclanthology.org/2021.acl-long.243.pdf [nan]: https://www.academia.edu/download/62266271/Deep_Learning20200303-80130-1s42zvt.pdf -[unsupervised-multitask]: https://storage.prod.researchhub.com/uploads/papers/2020/06/01/language-models.pdf \ No newline at end of file +[unsupervised-multitask]: https://storage.prod.researchhub.com/uploads/papers/2020/06/01/language-models.pdf diff --git a/scripts/benchmark/dataloaders.py b/scripts/benchmark/dataloaders.py index ff4ebc6..59ada15 100644 --- a/scripts/benchmark/dataloaders.py +++ b/scripts/benchmark/dataloaders.py @@ -1,666 +1,948 @@ """ -benchmark_dataloaders.py (Sonnet 4.6 generated) -======================== -Compares PackedDataLoader (binary/pretokenized) vs the nanochat -tokenizing_distributed_data_loader_with_state_bos_bestfit loader. +dataloaders.py +==================================== +Compares GPT-Lab, custom PyTorch, and adapted nanochat dataloaders on the same +deterministic Parquet corpus. Implementations are compared only when they use +the same packing policy: flat stream or destructive BOS-aligned best-fit. + +A correctness gate runs before timing. Measurements cover loader work and +device transfer only; model execution, warmup, corpus preparation, and +pretokenization are excluded. Metrics reported ---------------- - - Throughput tokens/sec (inputs only, i.e. B*T per batch) - - Batch latency ms per batch (mean ± std) - - Packing efficiency fraction of non-padding tokens (always 1.0 for both, - but we measure tokens-per-row to verify) - - Crop rate fraction of tokens discarded due to cropping - - BOS alignment fraction of rows that begin with the BOS token id - - Buffer search time time spent inside the best-fit search loop (us) +- Throughput median tokens/sec across trials +- Batch latency p50/p95 and mean/std milliseconds per batch +- Source token utilization fraction of advanced source tokens preserved +- Target supervision utilization fraction of target positions representing + source-token transitions +- Source transition coverage fraction of advanced source transitions seen +- Crop rate fraction of source tokens destructively cropped +- BOS alignment fraction of rows beginning with the BOS token +- Buffer occupancy observed buffered tokens/documents +- Memory pressure peak host RSS and accelerator allocation + +Artifacts +--------- +Each run writes reproducibility metadata and results to JSON and CSV, plus an +HTML report and policy-specific plots unless disabled. Usage ----- - # Quick synthetic test (no real data needed) - uv run python -m scripts.benchmark_dataloaders --mode synthetic - - # Run for different buffer sizes + # Quick policy-matched benchmark + uv run python -m scripts.benchmark.dataloaders \ + --dataset-path path/to/parquets \ + --quick + + # Compare only flat-stream implementations + uv run python -m scripts.benchmark.dataloaders \ + --dataset-path path/to/parquets \ + --groups stream_packing \ + --implementations gpt_lab_stream,custom_pytorch,nanochat_stream + + # Compare destructive BOS-best-fit with different document-buffer sizes for buf in 100 500 1000 2000; do - uv run python -m scripts.benchmark_dataloaders --buffer_size $buf -done - - # Test against real data - uv run python -m scripts.benchmark_dataloaders --mode real \ - --bin path/to/data.bin \ - --idx path/to/data.idx \ - --parquet_dir path/to/parquets/ - - # Run only one loader - uv run python -m scripts.benchmark_dataloaders --mode synthetic --loader packed - uv run python -m scripts.benchmark_dataloaders --mode synthetic --loader nanochat + uv run python -m scripts.benchmark.dataloaders \ + --dataset-path path/to/parquets \ + --groups bos_aligned_best_fit \ + --implementations custom_pytorch,nanochat_best_fit \ + --best-fit-buffer-docs "$buf" + done + + # Benchmark pretokenized input only and skip optional visual artifacts + uv run python -m scripts.benchmark.dataloaders \ + --dataset-path path/to/parquets \ + --tokenization pretokenized \ + --no-plots \ + --no-html """ +from __future__ import annotations + import argparse -import time +import csv +import html +import json +import math +import platform import statistics import sys -from collections import defaultdict +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from itertools import cycle +from pathlib import Path +from typing import Any, Callable, Iterable, Iterator, Sequence +import pyarrow as pa +import pyarrow.parquet as pq +import psutil import torch -import numpy as np - -from gpt_lab.data.loader import build_dataloader -from gpt_lab.tokenizer import Tokenizer - -# ───────────────────────────────────────────── -# Synthetic data helpers -# ───────────────────────────────────────────── - -def make_synthetic_bin_idx(num_docs=10_000, min_len=32, max_len=1024, vocab=50_257, seed=42): - """ - Build in-memory fake .bin / .idx buffers that look like PretokenizedDataset files. - Returns (tokens_np, offsets_np) – same dtypes as the real files. - """ - rng = np.random.default_rng(seed) - lengths = rng.integers(min_len, max_len + 1, size=num_docs) - tokens = rng.integers(1, vocab, size=int(lengths.sum()), dtype=np.uint32) - offsets = np.concatenate([[0], lengths.cumsum()]).astype(np.uint64) - return tokens, offsets - - -class SyntheticPretokenizedDataset: - """Drop-in replacement for PretokenizedDataset that lives in RAM.""" - def __init__(self, tokens, offsets): - self.tokens = tokens - self.offsets = offsets - self.num_docs = len(offsets) - - def __len__(self): - return self.num_docs - - def get_doc(self, idx): - start = int(self.offsets[idx]) - end = int(self.offsets[idx + 1]) if idx + 1 < self.num_docs else len(self.tokens) - return torch.from_numpy(self.tokens[start:end].astype(np.int64)) +from gpt_lab.data.loader import ( + DistDataLoader, + PackingStats, + tokenizing_distributed_data_loader_with_state_bos_bestfit, +) +from gpt_lab.tokenizer import Tokenizer, TokenizerConfig +from gpt_lab.utils.default import DATA_DIR + + +STREAM = "flat_stream" +BEST_FIT = "bos_aligned_best_fit" +STREAM_GROUP = "stream_packing" +BEST_FIT_GROUP = "bos_aligned_best_fit" +MIB = 1024 ** 2 +PROCESS = psutil.Process() + + +@dataclass +class Counters(PackingStats): + search_ns: int = 0 + searches: int = 0 + + +@dataclass +class Result: + loader: str + throughput_tokens_s: float + mean_latency_ms: float + std_latency_ms: float + source_tokens_read: int + new_source_tokens_advanced: int + target_positions_emitted: int + destructively_cropped_tokens: int + skipped_adjacent_transitions: int + synthetic_bos_tokens_inserted: int + intentional_bos_boundaries: int + buffered_source_tokens_delta: int + source_token_utilization: float + target_supervision_utilization: float + source_transition_coverage: float + destructive_crop_rate: float + bos_row_alignment: float + + @property + def destructive_cropped_tokens(self): + return self.destructively_cropped_tokens + + +class IdentityTokenizer: + def __init__(self, bos_id: int, vocab_size: int): + self.bos_id, self.vocab_size = bos_id, vocab_size + + def get_bos_token_id(self): + return self.bos_id + + def __call__(self, tokens, **_): + return list(tokens) + + def encode(self, rows, **_): + return [list(map(int, row)) for row in rows] + + +def write_split(path: Path, documents: Sequence[Any], row_group_size: int) -> None: + path.mkdir(parents=True, exist_ok=True) + table = pa.table({"text": list(documents)}) + pq.write_table(table, path / "shard_00000.parquet", row_group_size=row_group_size) + pq.write_table(table.slice(0, min(len(table), row_group_size)), path / "shard_00001.parquet", row_group_size=row_group_size) + + +def check_accounting(*, source_tokens_read, new_source_tokens_advanced, + buffered_source_tokens_delta, target_positions_emitted, + skipped_adjacent_transitions, synthetic_bos_tokens_inserted): + """Check source conservation independently from next-token supervision.""" + assert source_tokens_read == new_source_tokens_advanced + buffered_source_tokens_delta + represented = new_source_tokens_advanced - skipped_adjacent_transitions + assert target_positions_emitted == represented + synthetic_bos_tokens_inserted + + +class TorchPackedDataset: + """Tiny reference implementation used by the comparison tests.""" + + def __init__(self, path, tokenizer, batch_size, seq_len, tokenized, packing, buffer_docs): + self.path, self.tokenizer = path, tokenizer + self.batch_size, self.seq_len = batch_size, seq_len + self.packing, self.buffer_docs = packing, buffer_docs + + def _stream(self, documents: Iterable[Sequence[int]]): + documents, pending, carry = iter(documents), [], [] + while True: + needed = self.batch_size * self.seq_len + 1 + read = 0 + buffered_before = len(pending) + len(carry) + while len(carry) + len(pending) < needed: + document = list(next(documents)) + pending.extend(document) + read += len(document) + take = needed - len(carry) + window = carry + pending[:take] + del pending[:take] + carry = window[-1:] + data = torch.tensor(window) + yield ( + data[:-1].view(self.batch_size, self.seq_len), + data[1:].view(self.batch_size, self.seq_len), + { + "source_tokens_read": read, + "new_source_tokens_advanced": needed - 1, + "destructive_cropped_tokens": 0, + "skipped_adjacent_transitions": 0, + "synthetic_bos_tokens_inserted": 0, + "intentional_bos_boundaries": 0, + "buffered_source_tokens_delta": len(pending) + 1 - buffered_before, + "actual_buffered_tokens": len(pending) + 1, + "actual_buffered_documents": None, + }, + ) + + def _bestfit(self, documents: Iterable[Sequence[int]]): + documents, buffer = iter(documents), [] + capacity = self.seq_len + 1 + while True: + rows = torch.empty((self.batch_size, capacity), dtype=torch.long) + read = advanced = cropped = 0 + for row_index in range(self.batch_size): + pos = 0 + while pos < capacity: + while len(buffer) < self.buffer_docs: + document = list(next(documents)) + buffer.append(document) + read += len(document) + remaining = capacity - pos + index = max( + (i for i, doc in enumerate(buffer) if len(doc) <= remaining), + key=lambda i: len(buffer[i]), + default=-1, + ) + if index < 0: + index = min(range(len(buffer)), key=lambda i: len(buffer[i])) + document = buffer.pop(index) + take = min(len(document), remaining) + rows[row_index, pos:pos + take] = torch.tensor(document[:take]) + pos += take + advanced += len(document) + cropped += len(document) - take + yield rows[:, :-1], rows[:, 1:], { + "source_tokens_read": read, + "new_source_tokens_advanced": advanced, + "destructive_cropped_tokens": cropped, + "skipped_adjacent_transitions": self.batch_size + cropped, + "synthetic_bos_tokens_inserted": 0, + "intentional_bos_boundaries": self.batch_size, + "buffered_source_tokens_delta": read - advanced, + "actual_buffered_tokens": sum(map(len, buffer)), + "actual_buffered_documents": len(buffer), + } + + +def _sync(device): + if device.type == "cuda": + torch.cuda.synchronize(device) + elif device.type == "mps": + torch.mps.synchronize() + + +def memory_snapshot(device): + """Return process RSS and live accelerator allocation in bytes.""" + device = torch.device(device) + accelerator = None + if device.type == "cuda": + accelerator = torch.cuda.memory_allocated(device) + elif device.type == "mps": + accelerator = torch.mps.current_allocated_memory() + return PROCESS.memory_info().rss, accelerator + + +def benchmark(loader_name: str, factory: Callable[[Counters], Iterator], + tokenization: str, tokenizer_name: str, bos_id: int, + eos_id, workload: str, model_target: str, model, torch_workers, + args, nano_search_us: float) -> Result: + del tokenization, tokenizer_name, eos_id, workload, model_target, torch_workers, nano_search_us + counters, latencies = Counters(), [] + loader = factory(counters) + + def accumulate(batch_stats): + if not isinstance(batch_stats, dict): + return + for key, value in batch_stats.items(): + if key == "buffered_source_tokens_delta": + counters.buffered_source_tokens += value + elif hasattr(counters, key): + setattr(counters, key, getattr(counters, key) + value) + + for _ in range(args.warmup_batches): + inputs, _, batch_stats = next(loader) + accumulate(batch_stats) + if model is not None: + model(inputs) + buffered = counters.buffered_source_tokens + counters.reset() + counters.buffered_source_tokens = buffered + + bos_rows = total_rows = 0 + for _ in range(args.batches): + _sync(args.device) + started = time.perf_counter() + inputs, _, batch_stats = next(loader) + accumulate(batch_stats) + if model is not None: + model(inputs) + _sync(args.device) + latencies.append(time.perf_counter() - started) + bos_rows += int((inputs[:, 0] == bos_id).sum()) + total_rows += inputs.shape[0] + + target_positions = args.batches * args.batch_size * args.seq_len + buffered_delta = counters.buffered_source_tokens - buffered + check_accounting( + source_tokens_read=counters.source_tokens_read, + new_source_tokens_advanced=counters.new_source_tokens_advanced, + buffered_source_tokens_delta=buffered_delta, + target_positions_emitted=target_positions, + skipped_adjacent_transitions=counters.skipped_adjacent_transitions, + synthetic_bos_tokens_inserted=counters.synthetic_bos_tokens_inserted, + ) + represented = counters.new_source_tokens_advanced - counters.skipped_adjacent_transitions + advanced = max(counters.new_source_tokens_advanced, 1) + return Result( + loader_name, target_positions / sum(latencies), statistics.mean(latencies) * 1000, + statistics.stdev(latencies) * 1000 if len(latencies) > 1 else 0.0, + counters.source_tokens_read, counters.new_source_tokens_advanced, target_positions, + counters.destructive_cropped_tokens, + counters.skipped_adjacent_transitions, counters.synthetic_bos_tokens_inserted, + counters.intentional_bos_boundaries, buffered_delta, + (counters.new_source_tokens_advanced - counters.destructive_cropped_tokens) / advanced, + represented / max(target_positions, 1), represented / advanced, + counters.destructive_cropped_tokens / advanced, + bos_rows / max(total_rows, 1), + ) + + +class SyntheticDocuments: + split, start_state = "train", None + + def __init__(self, documents): + self.documents = documents + + def __iter__(self): + for document in cycle(self.documents): + yield torch.tensor(document), None + + +@dataclass +class ComparisonResult: + comparison_group: str + implementation: str + implementation_id: str + packing_policy: str + provenance: str + tokenization: str + tokenizer: str + device: str + transfer_policy: str + batch_size: int + sequence_length: int + trials: int + batches_per_trial: int + total_measured_batches: int + throughput_median_tokens_s: float + throughput_p50_tokens_s: float + throughput_p95_tokens_s: float + throughput_mean_tokens_s: float + throughput_std_tokens_s: float + latency_median_ms: float + latency_p50_ms: float + latency_p95_ms: float + latency_mean_ms: float + latency_std_ms: float + host_rss_peak_mib: float + host_rss_peak_delta_mib: float + accelerator_peak_allocated_mib: float | None + accelerator_peak_delta_mib: float | None + destructive_crop_policy: str + destructively_cropped_tokens: int + destructive_crop_rate: float + bos_row_alignment: float + buffer_budget_tokens: int | None + buffer_budget_documents: int | None + actual_buffered_tokens_mean: float | None + actual_buffered_tokens_min: int | None + actual_buffered_tokens_max: int | None + actual_buffered_documents_mean: float | None + source_tokens_read: int + new_source_tokens_advanced: int + target_positions_emitted: int + skipped_adjacent_transitions: int + synthetic_bos_tokens_inserted: int + intentional_bos_boundaries: int + buffered_source_tokens_delta: int + source_token_utilization: float + target_supervision_utilization: float + source_transition_coverage: float + correctness_status: str + correctness_batches: int + pretokenization_seconds: float + + @property + def destructive_cropped_tokens(self): + return self.destructively_cropped_tokens + + +@dataclass +class LoaderSpec: + implementation: str + implementation_id: str + packing_policy: str + comparison_group: str + provenance: str + factory: Callable[[], Iterator] + + +class BenchmarkTokenizer: + def __init__(self, tokenizer): + self.tokenizer = tokenizer + @property + def vocab_size(self): + return int(self.tokenizer.vocab_size) -# ───────────────────────────────────────────── -# HuggingFace dataloader -# ───────────────────────────────────────────── + def get_bos_token_id(self): + return int(self.tokenizer.get_bos_token_id()) + def encode_one(self, text): + return list(self.tokenizer.encode(text, prepend_bos=True)) + def encode(self, texts, prepend=None, num_threads=1, **_): + try: + rows = self.tokenizer.encode(list(texts), num_threads=num_threads) + except TypeError: + rows = [self.tokenizer.encode(text) for text in texts] + return [([prepend] if prepend is not None else []) + list(row) for row in rows] -# ───────────────────────────────────────────── -# PyTorch-based dataloader -# from https://github.com/mddunlap924/PyTorch-LLM/blob/4bb378dcf6352c538b13f94ad5c325de5961f568/src/dataloading/preprocess.py -# ───────────────────────────────────────────── -from pathlib import Path -import pandas as pd - - -# Load Data -class LoadData: - """ - Load CSV Data Files - (Expand this class to other datasets suitable for your needs) - """ - - def __init__(self, base_dir: str): - """ - :param base_dir: Directory data files are stored - """ - self.base_dir = Path(base_dir) - - - def load(self, filename: str) -> pd.DataFrame: - """ - Pandas Read CSV filename - :param filename: Name of File to Load - :return: Data returned as a Pandas DataFrame - """ - return pd.read_csv(self.base_dir / filename, - low_memory=False) -from pathlib import Path -import pandas as pd -from torch.utils.data import Dataset -import torch -import numpy as np -from torch.utils.data import DataLoader +class CorpusDocuments: + split, start_state = "train", None + def __init__(self, raw, tokens, tokenizer, on_the_fly): + self.raw, self.tokens = raw, tokens + self.tokenizer, self.on_the_fly = tokenizer, on_the_fly -class CustomTextCollator: - """ - Data Collator used for a classification task. - - It uses a given tokenizer and label encoder to convert any text and labels to numbers that - can go straight into a GPT2 model. + def __iter__(self): + for index in cycle(range(len(self.raw))): + row = self.tokenizer.encode_one(self.raw[index]) if self.on_the_fly else self.tokens[index] + yield torch.tensor(row, dtype=torch.long), None - This class is built with reusability in mind: it can be used as is as long - as the `dataloader` outputs a batch in dictionary format that can be passed - straight into the model - `model(**batch)`. - Arguments: +class TransferIterator: + def __init__(self, iterator, device): + self.iterator, self.device = iter(iterator), torch.device(device) + self.use_cuda = torch.device(device).type == "cuda" - use_tokenizer (:obj:`transformers.tokenization_?`): - Transformer type tokenizer used to process raw text into numbers. + def __iter__(self): + return self - labels_ids (:obj:`dict`): - Dictionary to encode any labels names into numbers. Keys map to - labels names and Values map to number associated to those labels. + def __next__(self): + inputs, targets, stats = next(self.iterator) + if self.use_cuda: + inputs, targets = inputs.pin_memory(), targets.pin_memory() + return ( + inputs.to(self.device, non_blocking=self.use_cuda), + targets.to(self.device, non_blocking=self.use_cuda), + stats, + ) - max_sequence_len (:obj:`int`, `optional`) - Value to indicate the maximum desired sequence to truncate or pad text - sequences. If no value is passed it will used maximum sequence size - supported by the tokenizer and model. - """ +class StatsAdapter: + """Expose per-batch deltas from production aggregate PackingStats.""" - def __init__(self, tokenizer, tokenizer_cfg): + FIELDS = ( + "source_tokens_read", "new_source_tokens_advanced", + "destructive_cropped_tokens", "skipped_adjacent_transitions", + "synthetic_bos_tokens_inserted", "intentional_bos_boundaries", + ) - # Tokenizer to be used inside the class. - self.tokenizer = tokenizer + def __init__(self, iterator, stats): + self.iterator, self.stats = iter(iterator), stats - # Tokenizer configuration - self.tok_cfg = tokenizer_cfg - - # Check max sequence length. - self.max_sequence_len = tokenizer_cfg.max_length - return - - - def __call__(self, sequences): - """ - This function allows the class objects to be used as a function call. - Since the PyTorch DataLoader needs a collator function, this - class can be used as a function. - - Arguments: - - item (:obj:`list`): - List of texts and labels. - - Returns: - :obj:`Dict[str, object]`: Dictionary of inputs that feed into the model. - It holds the statement `model(**Returned Dictionary)`. - """ - - # Get all texts from sequences list. - texts = [sequence['text'] for sequence in sequences] - # Get all labels from sequences list. - labels = [sequence['label'] for sequence in sequences] - - # Call tokenizer on all texts to convert into tensors of numbers with - # appropriate padding. - # https://huggingface.co/docs/transformers/pad_truncation - inputs = self.tokenizer(text=texts, - return_tensors=self.tok_cfg.return_tensors, - padding=self.tok_cfg.padding, - truncation=self.tok_cfg.truncation, - max_length=self.max_sequence_len, - add_special_tokens=self.tok_cfg.add_special_tokens, - ) - # Update the inputs with the associated encoded labels as tensor. - inputs.update({'labels': torch.tensor(labels, dtype=torch.long)}) - return inputs - - -class TrainDataset(Dataset): - def __init__(self, - df: pd.DataFrame, - tok, - tok_cfg, - X_cols: list[str], - label: str, - encoder): - self.df = df - self.tokenizer = tok - self.tokenizer_cfg = tok_cfg - self.X_cols = X_cols - self.label = label - self.encoder = encoder - - - def __len__(self): - return len(self.df) - - - def __getitem__(self, idx): - # Extract all source fields into a list - text = [] - for col in self.X_cols: - if col == 'ZIP code': - feature = f'Zip code {self.df[col].iloc[idx]}' - elif col == 'Sub-issue': - feature = f'{self.df[col].iloc[idx]}' - elif col == 'Consumer complaint narrative': - feature = self.df[col].iloc[idx] - text.append(feature) - - # Combine the fields using special SEP token - text = '[SEP]'.join(text) - # Extract all source fields into a list - # text = self.df['Consumer complaint narrative'].iloc[idx] - - # Convert text labels into labels (e.g., if 18 classes then labels are 0-17) - label_text = self.df[self.label].iloc[idx] - label = self.encoder.transform([label_text])[0] - return {'text': text, 'label': label} - - -class TestDataset(Dataset): - def __init__(self, df, tokenizer, tokenizer_cfg): - self.tokenizer = tokenizer - self.tokenizer_cfg = tokenizer_cfg - self.texts = df['full_text'].values - - def __len__(self): - return len(self.texts) - - def __getitem__(self, item): - inputs = prepare_input(tokenizer=self.tokenizer, - cfg=self.tokenizer_cfg, - text=self.texts[item]) - input_ids = torch.tensor(inputs['input_ids'], dtype=torch.float) - return {'input_ids': input_ids} - - -def get_ds_dl(df, - cfg, - tokenizer, - encoder, - collator): - "Get the PyTorch Dataset (ds) and Dataloader (dl)" - # Dataset - ds = TrainDataset(df=df, - tok=tokenizer, - tok_cfg=cfg.tokenizer, - X_cols=cfg.data_info.source_fields, - label=cfg.data_info.target, - encoder=encoder) - - # Dataloader - dl = DataLoader(ds, - batch_size=cfg.batch_size, - collate_fn=collator, - shuffle=True, - num_workers=cfg.num_workers, - pin_memory=True, - ) - return ds, dl + def __iter__(self): + return self + def __next__(self): + before = {field: getattr(self.stats, field) for field in self.FIELDS} + buffered_before = self.stats.buffered_source_tokens + inputs, targets, _ = next(self.iterator) + values = {field: getattr(self.stats, field) - before[field] for field in self.FIELDS} + values["buffered_source_tokens_delta"] = self.stats.buffered_source_tokens - buffered_before + values["actual_buffered_tokens"] = self.stats.buffered_source_tokens + values["actual_buffered_documents"] = None + return inputs, targets, values -# ───────────────────────────────────────────── -# nanochat dataloader -# copied for easy comparision from -# https://github.com/karpathy/nanochat/blob/324e69c45d3606095adb6b409078647145165454/nanochat/dataloader.py -# ───────────────────────────────────────────── -""" -Distributed dataloaders for pretraining. - -BOS-aligned bestfit: - - Every row starts with BOS token - - Documents packed using best-fit algorithm to minimize cropping - - When no document fits remaining space, crops a document to fill exactly - - 100% utilization (no padding), ~35% tokens cropped at T=2048 - -Compared to the original tokenizing_distributed_data_loader: -BOS-aligned loses ~35% of tokens to cropping, but ensures that -there are fewer "confusing" tokens in the train/val batches as every token can -now attend back to the BOS token and sees the full context of the document. - -Fallback to the original if you have very limited data AND long documents: -https://github.com/karpathy/nanochat/blob/3c3a3d7/nanochat/dataloader.py#L78-L117 -""" +def nanochat_flat_stream(documents, batch_size, seq_len, device): + documents, token_buffer = iter(documents), [] + device, use_cuda = torch.device(device), torch.device(device).type == "cuda" + advance = batch_size * seq_len + while True: + buffered_before, source_read = len(token_buffer), 0 + while len(token_buffer) < advance + 1: + document = list(next(documents)) + token_buffer.extend(document) + source_read += len(document) + values = token_buffer[:advance + 1] + token_buffer = token_buffer[advance:] + scratch = torch.tensor(values, dtype=torch.long, pin_memory=use_cuda) + stats = { + "source_tokens_read": source_read, + "new_source_tokens_advanced": advance, + "destructive_cropped_tokens": 0, + "skipped_adjacent_transitions": 0, + "synthetic_bos_tokens_inserted": 0, + "intentional_bos_boundaries": 0, + "buffered_source_tokens_delta": len(token_buffer) - buffered_before, + "actual_buffered_tokens": len(token_buffer), + "actual_buffered_documents": None, + } + yield ( + scratch[:-1].view(batch_size, seq_len).to(device, non_blocking=use_cuda), + scratch[1:].view(batch_size, seq_len).to(device, non_blocking=use_cuda), + stats, + ) -import torch -import pyarrow.parquet as pq -from gpt_lab.utils.distributed import get_dist_info # replaced 'from nanochat.common import get_dist_info' -# L317-353 replaced 'from nanochat.dataset import list_parquet_files' -import os -from gpt_lab.utils.common import DATA_DIR -base_dir = DATA_DIR -def list_parquet_files(data_dir=None, warn_on_legacy=False): - """ Looks into a data dir and returns full paths to all parquet files. """ - data_dir = DATA_DIR if data_dir is None else data_dir - - # Legacy-supporting code due to the upgrade from FinewebEdu-100B to ClimbMix-400B - # This code will eventually be deleted. - if not os.path.exists(data_dir): - if warn_on_legacy: - print() - print("=" * 80) - print(" WARNING: DATASET UPGRADE REQUIRED") - print("=" * 80) - print() - print(f" Could not find: {data_dir}") - print() - print(" nanochat recently switched from FinewebEdu-100B to ClimbMix-400B.") - print(" Everyone who does `git pull` as of March 4, 2026 is expected to see this message.") - print(" To upgrade to the new ClimbMix-400B dataset, run these two commands:") - print() - print(" python -m nanochat.dataset -n 170 # download ~170 shards, enough for GPT-2, adjust as desired") - print(" python -m scripts.tok_train # re-train tokenizer on new ClimbMix data") - print() - print(" For now, falling back to your old FinewebEdu-100B dataset...") - print("=" * 80) - print() - # attempt a fallback to the legacy data directory - data_dir = os.path.join(base_dir, "base_data") - - parquet_files = sorted([ - f for f in os.listdir(data_dir) - if f.endswith('.parquet') and not f.endswith('.tmp') - ]) - parquet_paths = [os.path.join(data_dir, f) for f in parquet_files] - return parquet_paths - -def _document_batches(split, resume_state_dict, tokenizer_batch_size): - """ - Infinite iterator over document batches (list of text strings) from parquet files. - - Handles DDP sharding and approximate resume. Each yield is (text_batch, (pq_idx, rg_idx, epoch)) - where text_batch is a list of document strings, indices track position for resumption, - and epoch counts how many times we've cycled through the dataset (starts at 1). - """ - ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() - - warn_on_legacy = ddp_rank == 0 and split == "train" # rank 0 on train split will warn on legacy - parquet_paths = list_parquet_files(warn_on_legacy=warn_on_legacy) - assert len(parquet_paths) != 0, "No dataset parquet files found, did you run dataset.py?" - parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:] - - resume_pq_idx = resume_state_dict["pq_idx"] if resume_state_dict is not None else 0 - resume_rg_idx = resume_state_dict["rg_idx"] if resume_state_dict is not None else None - resume_epoch = resume_state_dict.get("epoch", 1) if resume_state_dict is not None else 1 - first_pass = True - pq_idx = resume_pq_idx - epoch = resume_epoch - - while True: # iterate infinitely (multi-epoch) - pq_idx = resume_pq_idx if first_pass else 0 - while pq_idx < len(parquet_paths): - filepath = parquet_paths[pq_idx] - pf = pq.ParquetFile(filepath) - # Start from resume point if resuming on same file, otherwise from DDP rank - if first_pass and (resume_rg_idx is not None) and (pq_idx == resume_pq_idx): - base_idx = resume_rg_idx // ddp_world_size - base_idx += 1 # advance by 1 so we don't repeat data after resuming - rg_idx = base_idx * ddp_world_size + ddp_rank - if rg_idx >= pf.num_row_groups: - pq_idx += 1 - continue - resume_rg_idx = None # only do this once +def percentile(values, fraction): + values = sorted(values) + index = (len(values) - 1) * fraction + low, high = math.floor(index), math.ceil(index) + return values[low] if low == high else values[low] + (values[high] - values[low]) * (index - low) + + +def read_documents(path, column, limit): + documents = [] + for shard in sorted(path.glob("*.parquet")): + parquet = pq.ParquetFile(shard) + if column not in parquet.schema_arrow.names: + raise ValueError(f"{column!r} is missing from {shard}") + for row_group in range(parquet.num_row_groups): + documents.extend(parquet.read_row_group(row_group, columns=[column]).column(0).to_pylist()) + if len(documents) >= limit: + return documents[:limit] + if not documents or not all(isinstance(value, str) for value in documents): + raise ValueError("The selected corpus must contain text documents") + return documents + + +def find_dataset(path): + path = path.expanduser().resolve() + if any(path.glob("*.parquet")): + return path + candidates = sorted({item.parent for item in path.rglob("*.parquet")}) if path.exists() else [] + if len(candidates) != 1: + raise ValueError(f"Expected one Parquet dataset under {path}; found {len(candidates)}") + return candidates[0] + + +def token_rows(raw, tokens, tokenizer, on_the_fly): + for index in cycle(range(len(raw))): + yield tokenizer.encode_one(raw[index]) if on_the_fly else list(tokens[index]) + + +def build_specs(raw, tokens, tokenizer, on_the_fly, raw_path, token_path, policy, selected, args): + group, specs = (STREAM_GROUP if policy == STREAM else BEST_FIT_GROUP), [] + rows = lambda: token_rows(raw, tokens, tokenizer, on_the_fly) + if policy == STREAM and "gpt_lab_stream" in selected: + def gpt_lab(): + stats = PackingStats() + loader = DistDataLoader( + CorpusDocuments(raw, tokens, tokenizer, on_the_fly), + args.batch_size, args.seq_len, device=args.device, + packing_strategy="stream", bos_token_id=tokenizer.get_bos_token_id(), + packing_stats=stats, + ) + return StatsAdapter(loader, stats) + specs.append(LoaderSpec("GPT-Lab DistDataLoader", "gpt_lab_stream", policy, group, "gpt_lab.data.loader.DistDataLoader", gpt_lab)) + if "custom_pytorch" in selected: + def custom(): + dataset = TorchPackedDataset(None, None, args.batch_size, args.seq_len, True, "stream" if policy == STREAM else "bestfit", args.best_fit_buffer_docs) + iterator = dataset._stream(rows()) if policy == STREAM else dataset._bestfit(rows()) + return TransferIterator(iterator, args.device) + specs.append(LoaderSpec("Custom PyTorch packer", "custom_pytorch", policy, group, "benchmark-local implementation; not PyTorch generally", custom)) + nano_id = "nanochat_stream" if policy == STREAM else "nanochat_best_fit" + if nano_id in selected: + if policy == STREAM: + factory = lambda: nanochat_flat_stream(rows(), args.batch_size, args.seq_len, args.device) + name = "nanochat flat stream (adapted)" + else: + def factory(): + stats = PackingStats() + path = raw_path if on_the_fly else token_path + active_tokenizer = tokenizer if on_the_fly else IdentityTokenizer(tokenizer.get_bos_token_id(), tokenizer.vocab_size) + loader = tokenizing_distributed_data_loader_with_state_bos_bestfit( + active_tokenizer, B=args.batch_size, T=args.seq_len, split="train", + tokenizer_threads=1, tokenizer_batch_size=1, device=args.device, + buffer_size=args.best_fit_buffer_docs, base_path=path, packing_stats=stats, + ) + return StatsAdapter(loader, stats) + name = "nanochat BOS best-fit (adapted)" + specs.append(LoaderSpec(name, nano_id, policy, group, "vendored/adapted from nanochat 3c3a3d7; no upstream runtime import", factory)) + return specs + + +def correctness_gate(specs, documents, bos, args): + reference_dataset = TorchPackedDataset(None, None, args.batch_size, args.seq_len, True, "stream" if specs[0].packing_policy == STREAM else "bestfit", args.best_fit_buffer_docs) + reference = reference_dataset._stream(cycle(documents)) if specs[0].packing_policy == STREAM else reference_dataset._bestfit(cycle(documents)) + expected = [(x.clone(), y.clone(), stats) for x, y, stats in (next(reference) for _ in range(args.correctness_batches))] + alignments = {} + for spec in specs: + loader, previous, bos_rows, total_rows = spec.factory(), None, 0, 0 + for batch_index, (expected_inputs, expected_targets, expected_stats) in enumerate(expected): + inputs, targets, stats = next(loader) + _sync(args.device) + inputs, targets = inputs.detach().cpu().clone(), targets.detach().cpu().clone() + label = f"{spec.implementation} correctness batch {batch_index}" + if inputs.shape != (args.batch_size, args.seq_len) or targets.shape != inputs.shape: + raise RuntimeError(f"{label}: invalid shapes") + if not torch.equal(inputs, expected_inputs) or not torch.equal(targets, expected_targets): + raise RuntimeError(f"{label}: differs from independent {spec.packing_policy} reference") + if spec.packing_policy == STREAM: + if not torch.equal(inputs.flatten()[1:], targets.flatten()[:-1]): + raise RuntimeError(f"{label}: invalid flat shift") + if previous is not None and inputs.flatten()[0] != previous.flatten()[-1]: + raise RuntimeError(f"{label}: invalid one-token carry") + if stats["destructive_cropped_tokens"]: + raise RuntimeError(f"{label}: stream destructively cropped tokens") + if stats["skipped_adjacent_transitions"]: + raise RuntimeError(f"{label}: stream skipped adjacent transitions") else: - rg_idx = ddp_rank - while rg_idx < pf.num_row_groups: - rg = pf.read_row_group(rg_idx) - batch = rg.column('text').to_pylist() - for i in range(0, len(batch), tokenizer_batch_size): - yield batch[i:i+tokenizer_batch_size], (pq_idx, rg_idx, epoch) - rg_idx += ddp_world_size - pq_idx += 1 - first_pass = False - epoch += 1 - - -def tokenizing_distributed_data_loader_with_state_bos_bestfit( - tokenizer, B, T, split, - tokenizer_threads=4, tokenizer_batch_size=128, - device="cuda", resume_state_dict=None, - buffer_size=1000 -): - """ - BOS-aligned dataloader with Best-Fit Cropping. - - Reduces token waste compared to simple greedy cropping by searching a buffer - for documents that fit well, while maintaining 100% utilization (no padding). - - Algorithm for each row: - 1. From buffered docs, pick the LARGEST doc that fits entirely - 2. Repeat until no doc fits - 3. When nothing fits, crop a doc to fill remaining space exactly - - Key properties: - - Every row starts with BOS - - 100% utilization (no padding, every token is trained on) - - Approximately 35% of all tokens are discarded due to cropping - """ - assert split in ["train", "val"], "split must be 'train' or 'val'" - - row_capacity = T + 1 - batches = _document_batches(split, resume_state_dict, tokenizer_batch_size) - bos_token = tokenizer.get_bos_token_id() - doc_buffer = [] - pq_idx, rg_idx, epoch = 0, 0, 1 - - def refill_buffer(): - nonlocal pq_idx, rg_idx, epoch - doc_batch, (pq_idx, rg_idx, epoch) = next(batches) - token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads) - for tokens in token_lists: - doc_buffer.append(tokens) - - # Pre-allocate buffers once: layout is [inputs (B*T) | targets (B*T)] - # This gives us contiguous views and a single HtoD transfer - use_cuda = device == "cuda" - row_buffer = torch.empty((B, row_capacity), dtype=torch.long) # for building rows without creating Python lists - cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=use_cuda) # staging area (CPU) - gpu_buffer = torch.empty(2 * B * T, dtype=torch.long, device=device) # on-device buffer - cpu_inputs = cpu_buffer[:B * T].view(B, T) # a few views into these buffers just for convenience - cpu_targets = cpu_buffer[B * T:].view(B, T) - inputs = gpu_buffer[:B * T].view(B, T) - targets = gpu_buffer[B * T:].view(B, T) - - while True: - for row_idx in range(B): - pos = 0 - while pos < row_capacity: - # Ensure buffer has documents - while len(doc_buffer) < buffer_size: - refill_buffer() - - remaining = row_capacity - pos - - # Find largest doc that fits entirely - best_idx = -1 - best_len = 0 - for i, doc in enumerate(doc_buffer): - doc_len = len(doc) - if doc_len <= remaining and doc_len > best_len: - best_idx = i - best_len = doc_len - - if best_idx >= 0: - doc = doc_buffer.pop(best_idx) - doc_len = len(doc) - row_buffer[row_idx, pos:pos + doc_len] = torch.tensor(doc, dtype=torch.long) - pos += doc_len - else: - # No doc fits - crop shortest in buffer to fill remaining and minimize waste - shortest_idx = min(range(len(doc_buffer)), key=lambda i: len(doc_buffer[i])) - doc = doc_buffer.pop(shortest_idx) - row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining], dtype=torch.long) - pos += remaining - - # Copy to pinned CPU buffer, then single HtoD transfer - cpu_inputs.copy_(row_buffer[:, :-1]) - cpu_targets.copy_(row_buffer[:, 1:]) - - state_dict = {"pq_idx": pq_idx, "rg_idx": rg_idx, "epoch": epoch} - - # Single HtoD copy into persistent GPU buffer and yield - gpu_buffer.copy_(cpu_buffer, non_blocking=use_cuda) - yield inputs, targets, state_dict - -def tokenizing_distributed_data_loader_bos_bestfit(*args, **kwargs): - """Helper that omits state_dict from yields.""" - for inputs, targets, state_dict in tokenizing_distributed_data_loader_with_state_bos_bestfit(*args, **kwargs): - yield inputs, targets - -# ───────────────────────────────────────────── -# Benchmark runner -# ───────────────────────────────────────────── - -def benchmark_loader(name, loader, B, T, num_batches, bos_id, device): - latencies = [] - last_stats = {} - - # warm-up - for _ in range(2): + if not torch.equal(inputs[:, 1:], targets[:, :-1]) or not bool(torch.all(inputs[:, 0] == bos)): + raise RuntimeError(f"{label}: invalid BOS alignment or row shift") + if stats["destructive_cropped_tokens"] != expected_stats["destructive_cropped_tokens"]: + raise RuntimeError(f"{label}: crop accounting differs from reference") + try: + check_accounting( + source_tokens_read=stats["source_tokens_read"], + new_source_tokens_advanced=stats["new_source_tokens_advanced"], + buffered_source_tokens_delta=stats["buffered_source_tokens_delta"], + target_positions_emitted=args.batch_size * args.seq_len, + skipped_adjacent_transitions=stats["skipped_adjacent_transitions"], + synthetic_bos_tokens_inserted=stats["synthetic_bos_tokens_inserted"], + ) + except AssertionError as error: + raise RuntimeError(f"{label}: aggregate accounting failed") from error + bos_rows += int((inputs[:, 0] == bos).sum()); total_rows += inputs.shape[0] + previous = targets + alignments[spec.implementation_id] = bos_rows / total_rows + return alignments + + +def benchmark_trial(spec, args): + _sync(args.device) + host_baseline, accelerator_baseline = memory_snapshot(args.device) + if args.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(args.device) + host_peak, accelerator_peak = host_baseline, accelerator_baseline + + def sample_memory(): + nonlocal host_peak, accelerator_peak + host, accelerator = memory_snapshot(args.device) + host_peak = max(host_peak, host) + if accelerator is not None: + accelerator_peak = max(accelerator_peak or 0, accelerator) + + loader = spec.factory() + sample_memory() + for _ in range(args.warmup_batches): next(loader) - - bos_rows = 0 - total_rows = 0 - - for _ in range(num_batches): - t0 = time.perf_counter() - inputs, targets, stats = next(loader) - t1 = time.perf_counter() - latencies.append((t1 - t0) * 1000) - last_stats = stats - - # check BOS alignment - bos_rows += int((inputs[:, 0] == bos_id).sum().item()) - total_rows += B - - return { - "name": name, - "mean_latency_ms": statistics.mean(latencies), - "std_latency_ms": statistics.stdev(latencies) if len(latencies) > 1 else 0.0, - "throughput_tok_s": (B * T * 1000) / statistics.mean(latencies), - "crop_rate": last_stats.get("cropped_tokens", 0) / max(last_stats.get("total_tokens", 1), 1), - "mean_search_us": statistics.mean(last_stats["search_times_us"]) if last_stats.get("search_times_us") else 0.0, - "bos_alignment": bos_rows / max(total_rows, 1), + _sync(args.device) + sample_memory() + _sync(args.device) + latencies, stats = [], [] + for _ in range(args.batches): + _sync(args.device); started = time.perf_counter() + inputs, targets, batch_stats = next(loader) + _sync(args.device); latencies.append((time.perf_counter() - started) * 1000) + sample_memory() + stats.append(batch_stats) + del inputs, targets + if args.device.type == "cuda": + accelerator_peak = max(accelerator_peak or 0, torch.cuda.max_memory_allocated(args.device)) + tokens = args.batches * args.batch_size * args.seq_len + for batch_stats in stats: + check_accounting( + source_tokens_read=batch_stats["source_tokens_read"], + new_source_tokens_advanced=batch_stats["new_source_tokens_advanced"], + buffered_source_tokens_delta=batch_stats["buffered_source_tokens_delta"], + target_positions_emitted=args.batch_size * args.seq_len, + skipped_adjacent_transitions=batch_stats["skipped_adjacent_transitions"], + synthetic_bos_tokens_inserted=batch_stats["synthetic_bos_tokens_inserted"], + ) + memory = { + "host_rss_peak_mib": host_peak / MIB, + "host_rss_peak_delta_mib": max(0, host_peak - host_baseline) / MIB, + "accelerator_peak_allocated_mib": None if accelerator_peak is None else accelerator_peak / MIB, + "accelerator_peak_delta_mib": None if accelerator_peak is None else max(0, accelerator_peak - (accelerator_baseline or 0)) / MIB, } - - -def print_results(results: list[dict]): - keys = [ - ("mean_latency_ms", "Latency mean (ms)", ".2f"), - ("std_latency_ms", "Latency std (ms)", ".2f"), - ("throughput_tok_s", "Throughput (tok/s)", ",.0f"), - ("crop_rate", "Crop rate", ".2%"), - ("mean_search_us", "Search time mean (μs)", ".3f"), - ("bos_alignment", "BOS alignment", ".2%"), - ] - - col_w = 32 - name_w = 42 - - header = f"{'Metric':<{col_w}}" + "".join(f"{r['name']:<{name_w}}" for r in results) - print() - print("=" * (col_w + name_w * len(results))) - print(header) - print("-" * (col_w + name_w * len(results))) - - for key, label, fmt in keys: - row = f"{label:<{col_w}}" - for r in results: - val = r.get(key) - if val is None: - row += f"{'N/A':<{name_w}}" - else: - row += f"{format(val, fmt):<{name_w}}" - print(row) - - print("=" * (col_w + name_w * len(results))) - print() - -# ───────────────────────────────────────────── -# Main -# ───────────────────────────────────────────── - -def main(): - parser = argparse.ArgumentParser(description="Benchmark PackedDataLoader vs nanochat best-fit") - parser.add_argument("--mode", choices=["synthetic", "real"], default="synthetic") - parser.add_argument("--loader", choices=["glab", "glab-on-the-fly", "glab-tokenized", "nanochat", "all"], default="all") - parser.add_argument("--batch_size", type=int, default=8) - parser.add_argument("--seq_len", type=int, default=2048) - parser.add_argument("--num_batches", type=int, default=20, - help="Number of batches to time (after 2 warm-up batches)") - parser.add_argument("--buffer_size", type=int, default=1000, - help="Document buffer size for both loaders") - parser.add_argument("--num_docs", type=int, default=20_000, - help="Number of synthetic documents") - parser.add_argument("--device", default="cpu", - help="torch device, e.g. cpu or cuda") - parser.add_argument("--bos_id", type=int, default=1, - help="BOS token id (for nanochat loader)") - # real-data paths - parser.add_argument("--bin", default=None, help="Path to pretokenized .bin file") - parser.add_argument("--idx", default=None, help="Path to pretokenized .idx file") - args = parser.parse_args() - - print(f"\n{'─'*60}") - print(f" Dataloader benchmark") - print(f" mode={args.mode} B={args.batch_size} T={args.seq_len}") - print(f" batches={args.num_batches} buffer={args.buffer_size} device={args.device}") - print(f"{'─'*60}\n") - - if args.mode == "synthetic": - print(f"Generating {args.num_docs:,} synthetic documents …") - tokens, offsets = make_synthetic_bin_idx(num_docs=args.num_docs) - dataset = SyntheticPretokenizedDataset(tokens, offsets) - print(f" vocab=50257 len range=[32, 1024] total tokens={len(tokens):,}\n") - else: - if args.bin is None or args.idx is None: - sys.exit("--mode real requires --bin and --idx arguments.") - # import real dataset class + return tokens / (sum(latencies) / 1000), latencies, stats, memory + + +def aggregate(spec, trials, alignment, mode, tokenizer_name, pretokenization_seconds, args): + throughputs = [trial[0] for trial in trials] + latencies = [latency for _, values, _, _ in trials for latency in values] + stats = [value for _, _, values, _ in trials for value in values] + memories = [memory for _, _, _, memory in trials] + token_samples = [row["actual_buffered_tokens"] for row in stats if row.get("actual_buffered_tokens") is not None] + document_samples = [row["actual_buffered_documents"] for row in stats if row.get("actual_buffered_documents") is not None] + cropped = sum(row["destructive_cropped_tokens"] for row in stats) + advanced = sum(row["new_source_tokens_advanced"] for row in stats) + source_read = sum(row["source_tokens_read"] for row in stats) + target_positions = len(stats) * args.batch_size * args.seq_len + skipped = sum(row["skipped_adjacent_transitions"] for row in stats) + synthetic = sum(row["synthetic_bos_tokens_inserted"] for row in stats) + intentional = sum(row["intentional_bos_boundaries"] for row in stats) + buffered_delta = sum(row["buffered_source_tokens_delta"] for row in stats) + represented = advanced - skipped + check_accounting( + source_tokens_read=source_read, + new_source_tokens_advanced=advanced, + buffered_source_tokens_delta=buffered_delta, + target_positions_emitted=target_positions, + skipped_adjacent_transitions=skipped, + synthetic_bos_tokens_inserted=synthetic, + ) + return ComparisonResult( + spec.comparison_group, spec.implementation, spec.implementation_id, + spec.packing_policy, spec.provenance, mode, tokenizer_name, str(args.device), + "pinned CPU + non-blocking CUDA copy; regular copy otherwise", + args.batch_size, args.seq_len, args.trials, args.batches, args.trials * args.batches, + statistics.median(throughputs), percentile(throughputs, .5), percentile(throughputs, .95), + statistics.mean(throughputs), statistics.stdev(throughputs) if len(throughputs) > 1 else 0.0, + statistics.median(latencies), percentile(latencies, .5), percentile(latencies, .95), + statistics.mean(latencies), statistics.stdev(latencies) if len(latencies) > 1 else 0.0, + max(row["host_rss_peak_mib"] for row in memories), + max(row["host_rss_peak_delta_mib"] for row in memories), + max((row["accelerator_peak_allocated_mib"] for row in memories if row["accelerator_peak_allocated_mib"] is not None), default=None), + max((row["accelerator_peak_delta_mib"] for row in memories if row["accelerator_peak_delta_mib"] is not None), default=None), + "none" if spec.packing_policy == STREAM else "discard selected document remainder", + cropped, cropped / max(advanced, 1), alignment, + args.batch_size * args.seq_len + 1 if spec.packing_policy == STREAM else None, + args.best_fit_buffer_docs if spec.packing_policy == BEST_FIT else None, + statistics.mean(token_samples) if token_samples else None, + min(token_samples) if token_samples else None, max(token_samples) if token_samples else None, + statistics.mean(document_samples) if document_samples else None, + source_read, advanced, target_positions, + skipped, synthetic, intentional, buffered_delta, + (advanced - cropped) / max(advanced, 1), + represented / max(target_positions, 1), represented / max(advanced, 1), + "passed", args.correctness_batches, pretokenization_seconds, + ) + + +def write_outputs(results, metadata, output, plots, html_report): + rows = [asdict(result) for result in results] + (output / "results.json").write_text(json.dumps({"metadata": metadata, "results": rows}, indent=2)) + with (output / "results.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])); writer.writeheader(); writer.writerows(rows) + images = [] + if plots: try: - from gpt_lab.data.loader import ... # adjust import path as needed + import matplotlib.pyplot as plt + for policy in (STREAM, BEST_FIT): + selected = [row for row in results if row.packing_policy == policy] + if not selected: continue + labels = [f"{row.implementation}\n{row.tokenization}" for row in selected] + figure, axes = plt.subplots(1, 5, figsize=(max(22, len(selected) * 3), 5)) + axes[0].bar(labels, [row.throughput_median_tokens_s for row in selected]) + axes[1].bar(labels, [row.latency_p50_ms for row in selected], yerr=[row.latency_p95_ms - row.latency_p50_ms for row in selected], capsize=4) + axes[0].set_ylabel("tokens/s (trial median)"); axes[1].set_ylabel("latency p50, whisker to p95 (ms)") + x, width = list(range(len(selected))), 0.25 + for offset, field_name, label in ( + (-1, "source_token_utilization", "source preserved"), + (0, "source_transition_coverage", "source transitions covered"), + (1, "target_supervision_utilization", "targets supervising source"), + ): + axes[2].bar( + [value + offset * width for value in x], + [getattr(row, field_name) for row in selected], + width, + label=label, + ) + axes[2].set_xticks(x, labels); axes[2].set_ylim(0, 1.05) + axes[2].set_ylabel("utilization fraction"); axes[2].legend(fontsize=8) + axes[3].bar(labels, [row.destructive_crop_rate for row in selected]) + axes[3].set_ylim(0, 1.05) + axes[3].set_ylabel("destructive crop rate") + memory_width = 0.35 + axes[4].bar( + [value - memory_width / 2 for value in x], + [row.host_rss_peak_delta_mib for row in selected], + memory_width, + label="host RSS", + ) + if any(row.accelerator_peak_delta_mib is not None for row in selected): + axes[4].bar( + [value + memory_width / 2 for value in x], + [row.accelerator_peak_delta_mib or 0 for row in selected], + memory_width, + label="accelerator", + ) + axes[4].set_xticks(x, labels) + axes[4].set_ylabel("peak above baseline (MiB)") + axes[4].legend(fontsize=8) + for axis in axes: axis.tick_params(axis="x", labelrotation=20); axis.grid(axis="y", alpha=.25) + figure.suptitle(f"Matched policy: {policy}"); figure.tight_layout() + path = output / f"{policy}.png"; figure.savefig(path, dpi=150); plt.close(figure); images.append(path) except ImportError: - sys.exit("Could not import PretokenizedDataset. " - "Make sure dataloader.py is on sys.path.") - dataset = PretokenizedDataset(args.bin, args.idx) - print(f"Loaded dataset: {len(dataset):,} docs\n") - - results = [] - - if args.loader in ("glab", "both"): - print("Running DataLoader with tokenized data…") - t0 = time.perf_counter() - glab_loader = build_dataloader( - dataset, batch_size=args.batch_size, seq_len=args.seq_len, - buffer_size=args.buffer_size, device=args.device, - ) - init_time = time.perf_counter() - t0 - print(f" initialization time: {init_time:.2f}s") - r = benchmark_loader("DataLoader (O log N) - tokenized data", glab_loader, args.batch_size, args.seq_len, - args.num_batches, args.buffer_size, args.bos_id, args.device) - r["initialization_time_s"] = init_time - results.append(r) - print(" done.\n") - - if args.loader in ("glab", "both") : - print("Running DataLoader with data tokenization on-flight…") - t0 = time.perf_counter() - tokenizer = Tokenizer.from_pretrained("gpt2") # or your custom tokenizer - glab_loader = build_dataloader( - dataset, batch_size=args.batch_size, seq_len=args.seq_len, - buffer_size=args.buffer_size, device=args.device, tokenizer=tokenizer + print("matplotlib unavailable; skipping plots", file=sys.stderr) + if html_report: + columns = [ + "comparison_group", "packing_policy", "implementation", "tokenization", + "throughput_median_tokens_s", "latency_p50_ms", "latency_p95_ms", + "host_rss_peak_mib", "host_rss_peak_delta_mib", + "accelerator_peak_allocated_mib", "accelerator_peak_delta_mib", + "source_tokens_read", "new_source_tokens_advanced", "target_positions_emitted", + "destructively_cropped_tokens", "skipped_adjacent_transitions", + "synthetic_bos_tokens_inserted", "intentional_bos_boundaries", + "source_token_utilization", + "target_supervision_utilization", "source_transition_coverage", + "destructive_crop_rate", "bos_row_alignment", "actual_buffered_tokens_mean", + ] + heading = "".join(f"{column}" for column in columns) + body = "".join("" + "".join(f"{html.escape(str(asdict(row)[column]))}" for column in columns) + "" for row in results) + pictures = "".join(f'{path.stem}' for path in images) + (output / "report.html").write_text(f"

Policy-matched dataloaders

{html.escape(json.dumps(metadata, indent=2))}
{heading}{body}
{pictures}") + + +def parse_args(): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--dataset-path", type=Path, default=DATA_DIR); parser.add_argument("--column", default="text") + parser.add_argument("--output-dir", type=Path, default=Path("benchmark-results") / datetime.now().strftime("dataloaders-%Y%m%d-%H%M%S")) + parser.add_argument("--tokenizers", default="gpt2"); parser.add_argument("--tokenization", choices=("on-the-fly", "pretokenized", "both"), default="both") + parser.add_argument("--groups", default=f"{STREAM_GROUP},{BEST_FIT_GROUP}") + parser.add_argument("--implementations", default="gpt_lab_stream,custom_pytorch,nanochat_stream,nanochat_best_fit") + parser.add_argument("--device", type=torch.device, default=torch.device("cuda" if torch.cuda.is_available() else "cpu")) + parser.add_argument("--batch-size", type=int, default=4); parser.add_argument("--seq-len", type=int, default=128) + parser.add_argument("--batches", type=int, default=300); parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--warmup-batches", type=int, default=20); parser.add_argument("--correctness-batches", type=int, default=3) + parser.add_argument("--max-docs", type=int, default=4096); parser.add_argument("--row-group-size", type=int, default=256) + parser.add_argument("--best-fit-buffer-docs", type=int, default=1000) + parser.add_argument("--no-plots", action="store_true"); parser.add_argument("--html", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--quick", action="store_true") + args = parser.parse_args() + if args.quick: + args.max_docs, args.batches, args.trials, args.warmup_batches = min(args.max_docs, 256), 10, 2, 3 + args.best_fit_buffer_docs = min(args.best_fit_buffer_docs, 64) + args.output_dir = args.output_dir.expanduser().resolve() + if args.device.type == "cuda" and not torch.cuda.is_available(): parser.error("CUDA unavailable") + if args.device.type == "mps" and not torch.backends.mps.is_available(): parser.error("MPS unavailable") + if min(args.batches, args.trials, args.warmup_batches, args.correctness_batches) < 1: parser.error("timing counts must be positive") + return args + + +def load_tokenizer(name): + try: + return Tokenizer.from_pretrained(name) + except UnboundLocalError: + return Tokenizer.from_config(TokenizerConfig(name=name, source="tiktoken")) + + +def main() -> None: + args = parse_args(); args.output_dir.mkdir(parents=True, exist_ok=True) + dataset = find_dataset(args.dataset_path); raw = read_documents(dataset, args.column, args.max_docs) + selected, groups = set(filter(None, args.implementations.split(","))), set(filter(None, args.groups.split(","))) + known = {"gpt_lab_stream", "custom_pytorch", "nanochat_stream", "nanochat_best_fit"} + if selected - known: raise ValueError(f"Unknown implementations: {sorted(selected - known)}") + if groups - {STREAM_GROUP, BEST_FIT_GROUP}: raise ValueError(f"Unknown groups: {sorted(groups)}") + modes = ("on-the-fly", "pretokenized") if args.tokenization == "both" else (args.tokenization,) + results, execution_orders, pretokenization = [], [], {} + for tokenizer_name in filter(None, args.tokenizers.split(",")): + tokenizer = BenchmarkTokenizer(load_tokenizer(tokenizer_name)); bos = tokenizer.get_bos_token_id() + started = time.perf_counter(); tokens = tokenizer.encode(raw, prepend=bos); elapsed = time.perf_counter() - started + pretokenization[tokenizer_name] = {"seconds": elapsed, "documents": len(raw), "tokens": sum(map(len, tokens))} + corpus = args.output_dir / "corpus" / tokenizer_name.replace("/", "_") + raw_path, token_path = corpus / "raw", corpus / "pretokenized" + write_split(raw_path, raw, args.row_group_size); write_split(token_path, tokens, args.row_group_size) + for mode in modes: + for policy, group in ((STREAM, STREAM_GROUP), (BEST_FIT, BEST_FIT_GROUP)): + if group not in groups: continue + specs = build_specs(raw, tokens, tokenizer, mode == "on-the-fly", raw_path, token_path, policy, selected, args) + if len(specs) < 2: continue + print(f"Correctness gate: {group} / {mode} / {tokenizer_name}") + alignments = correctness_gate(specs, tokens, bos, args) + trials = {spec.implementation_id: [] for spec in specs} + for trial in range(args.trials): + offset = trial % len(specs); order = [*specs[offset:], *specs[:offset]] + execution_orders.append({"tokenizer": tokenizer_name, "tokenization": mode, "packing_policy": policy, "trial": trial, "order": [spec.implementation_id for spec in order]}) + for spec in order: + print(f"Trial {trial + 1}/{args.trials}: {spec.implementation} / {policy} / {mode}") + trials[spec.implementation_id].append(benchmark_trial(spec, args)) + results.extend(aggregate(spec, trials[spec.implementation_id], alignments[spec.implementation_id], mode, tokenizer_name, elapsed, args) for spec in specs) + if not results: raise RuntimeError("No selected matched group contains at least two implementations") + metadata = { + "created_at": datetime.now().astimezone().isoformat(), "command": sys.argv, + "python": sys.version, "platform": platform.platform(), "torch": torch.__version__, "device": str(args.device), + "corpus": {"path": str(dataset), "documents": len(raw), "order": "sorted shards/row groups/rows, repeated exactly"}, + "config": {"batch_size": args.batch_size, "sequence_length": args.seq_len, "trials": args.trials, "batches_per_trial": args.batches, "warmup_batches": args.warmup_batches, "correctness_batches": args.correctness_batches, "best_fit_buffer_documents": args.best_fit_buffer_docs, "tokenization": args.tokenization, "tokenizers": args.tokenizers}, + "pretokenization": pretokenization, + "matched_groups": {STREAM_GROUP: ["GPT-Lab", "Custom PyTorch", "nanochat adapted"], BEST_FIT_GROUP: ["Custom PyTorch", "nanochat adapted"]}, + "gpt_lab_best_fit": "Unavailable: GPT-Lab bos_aligned retains suffixes and is not destructive best-fit, so it is excluded.", + "invalid_previous_comparisons": ["stream versus destructive BOS-best-fit", "B*(T+1) row stream versus flat B*T+1 carry stream", "loader timing obscured by model execution"], + "timing": {"scope": "loader-only next() and device transfer", "excluded": ["construction", "correctness", "warmup", "pretokenization", "model execution"], "accelerator_sync": "before/after each measured batch", "rotated_orders": execution_orders}, + "memory": { + "scope": "loader construction, warmup, and measured batches", + "host": "process RSS sampled before construction and after construction/warmup/batches", + "accelerator": "CUDA allocator peak; live allocation samples on MPS; unavailable on CPU", + "aggregation": "maximum peak and peak-above-baseline across trials", + }, + "metric_definitions": { + "source_tokens_read": "Tokenized source tokens entering the packer during measured batches, including original document BOS tokens; excludes warmup and synthetic BOS.", + "new_source_tokens_advanced": "Unique source-stream tokens permanently advanced during measured batches, including discarded tokens and excluding a retained carry token.", + "target_positions_emitted": "Language-model target tensor positions produced; exactly trials * batches * B * T.", + "host_rss_peak_mib": "Maximum process resident memory observed during a trial, in MiB.", + "host_rss_peak_delta_mib": "Maximum process RSS above the pre-construction trial baseline, in MiB.", + "accelerator_peak_allocated_mib": "Maximum live tensor memory allocated on the selected accelerator, in MiB; null on CPU.", + "accelerator_peak_delta_mib": "Maximum accelerator allocation above the pre-construction trial baseline, in MiB; null on CPU.", + "destructively_cropped_tokens": ( + "Original tokenized source tokens permanently discarded by the packing algorithm " + "and never emitted or retained for a future batch. Excludes the shifted-batch carry, " + "synthetic BOS, temporarily buffered tails, and separately reported skipped " + "adjacent transitions." + ), + "skipped_adjacent_transitions": "Adjacent source-token pairs not emitted as (input, target), including crop loss and intentional row segmentation at BOS.", + "synthetic_bos_tokens_inserted": "Continuation BOS tokens absent from the source; never counted as source read or advanced.", + "intentional_bos_boundaries": "Subset of skipped transitions deliberately segmented because a row starts at an original source BOS.", + "buffered_source_tokens_delta": "Final minus initial buffered source tokens over measured batches, including retained carry tokens.", + "bos_row_alignment": "Fraction of rows whose first input token is BOS.", + "source_token_utilization": "(advanced - destructively cropped) / advanced; source preservation.", + "target_supervision_utilization": "Represented adjacent source transitions / emitted target positions; excludes synthetic supervision.", + "source_transition_coverage": "Represented adjacent source transitions / advanced source tokens.", + "destructive_crop_rate": "Destructively cropped source tokens / advanced source tokens.", + }, + "accounting_invariants": { + "source_balance": "source_tokens_read = new_source_tokens_advanced + buffered_source_tokens_delta", + "supervision_balance": "target_positions_emitted = new_source_tokens_advanced - skipped_adjacent_transitions + synthetic_bos_tokens_inserted", + "target_capacity": "target_positions_emitted = trials * batches * B * T", + "fifo": "Flat B*T+1 stream retains one carry and has zero destructive crop and zero skipped adjacent transitions.", + }, + "instrumentation": "Packing paths emit aggregate integer deltas only. Exact discarded suffix provenance is enabled only by debug tests and is outside throughput timing.", + "interpretation": "Implementation-level, policy-matched results; not a global framework ranking.", + } + write_outputs(results, metadata, args.output_dir, not args.no_plots, args.html) + for result in results: + print( + f"{result.packing_policy:24} {result.implementation:34} {result.tokenization:12} " + f"{result.throughput_median_tokens_s:12,.0f} tok/s " + f"source={result.source_token_utilization:.2%} " + f"coverage={result.source_transition_coverage:.2%} " + f"supervision={result.target_supervision_utilization:.2%} " + f"crop={result.destructive_crop_rate:.2%} " + f"skips={result.skipped_adjacent_transitions:,} " + f"synth_BOS={result.synthetic_bos_tokens_inserted:,} " + f"BOS_rows={result.bos_row_alignment:.2%} " + f"host_peak={result.host_rss_peak_mib:.1f}MiB " + f"host_delta={result.host_rss_peak_delta_mib:.1f}MiB " + f"accelerator_delta={'n/a' if result.accelerator_peak_delta_mib is None else f'{result.accelerator_peak_delta_mib:.1f}MiB'}" ) - r = benchmark_loader("DataLoader (O log N) - tokenization on-flight", glab_loader, args.batch_size, args.seq_len, - args.num_batches, args.buffer_size, args.bos_id, args.device) - init_time = time.perf_counter() - t0 - r["initialization_time_s"] = init_time - results.append(r) - print(" done.\n") - - if args.loader in ("nanochat", "both"): - print("Running nanochat BOS best-fit …") - r = benchmark_loader("nanochat BOS best-fit (O N)", dataset, args.batch_size, args.seq_len, - args.num_batches, args.buffer_size, - args.bos_id, args.device) - results.append(r) - print(" done.\n") - - print_results(results) - - if len(results) == 2: - ratio = results[0]["mean_latency_ms"] / results[1]["mean_latency_ms"] - faster = results[0]["name"] if ratio < 1 else results[1]["name"] - slower = results[1]["name"] if ratio < 1 else results[0]["name"] - print(f" {faster} is {abs(1 - ratio):.1%} {'faster' if ratio < 1 else 'slower'} " - f"than {slower} on average.\n") + print(f"Artifacts: {args.output_dir}") if __name__ == "__main__": diff --git a/scripts/train_base.py b/scripts/train_base.py index fff0fe7..8e9bfdd 100644 --- a/scripts/train_base.py +++ b/scripts/train_base.py @@ -91,8 +91,12 @@ def get_common_arguments(prs: ArgumentParser): # Temporary prs.add_argument("--device-batch-size", type=int, default=32, help="(default: 32) Batch size for each device during training. Batch size define further effective batch size as device_batch_size * max_seq_len * n_acc_steps.") - # For tests - prs.add_argument("--use-nanochat-dataloader", action="store_true", help="(default: False) Whether to use the nanochat dataloader instead of the default dataloader.") + prs.add_argument( + "--packing-strategy", + choices=("stream", "bos_aligned", "bos_bestfit_crop"), + default="stream", + help="Token packing policy. bos_bestfit_crop may discard document suffixes.", + ) return prs @@ -313,7 +317,7 @@ def get_common_arguments(prs: ArgumentParser): max_shards=None, # TODO: configured based on configs/data.yaml batch_size=base_training_config.get("device_batch_size", args.device_batch_size), dist_info=dist_info, - use_nanochat=args.use_nanochat_dataloader, + packing_strategy=args.packing_strategy, ) # TODO: add option to configure buffer size train_loader = build_dataloader(split="train", resume_state=resume_state, **loader_common_kwargs) @@ -405,4 +409,4 @@ def get_common_arguments(prs: ArgumentParser): # ------------------------------------------------------------------------------ board.close() # wandb run finish - cleanup_dist_groups() \ No newline at end of file + cleanup_dist_groups() diff --git a/src/gpt_lab/data/__init__.py b/src/gpt_lab/data/__init__.py index acda9e4..ed077b0 100644 --- a/src/gpt_lab/data/__init__.py +++ b/src/gpt_lab/data/__init__.py @@ -1,3 +1,4 @@ -from .loader import DistDataLoader, build_dataloader, DataLoaderState +from .loader import DistDataLoader, PackingStats, build_dataloader, DataLoaderState +from gpt_lab.utils.types import PackingStrategy -__all__ = ["DistDataLoader", "build_dataloader", "DataLoaderState"] \ No newline at end of file +__all__ = ["DistDataLoader", "PackingStats", "PackingStrategy", "build_dataloader", "DataLoaderState"] diff --git a/src/gpt_lab/data/loader.py b/src/gpt_lab/data/loader.py index a35207d..fe58df6 100644 --- a/src/gpt_lab/data/loader.py +++ b/src/gpt_lab/data/loader.py @@ -1,14 +1,57 @@ from collections import deque +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Callable, Iterator, Optional, Tuple, Union, Literal +from typing import Callable, Iterator, Optional, Tuple, Union import torch from gpt_lab.utils.default import DATA_DIR from gpt_lab.utils.distributed import get_dist_info from gpt_lab.utils.schemas import DataLoaderState +from gpt_lab.utils.types import PackingStrategy from gpt_lab.data.sharder import ShardManager + +@dataclass +class PackingStats: + """Optional aggregate packing counters; omitted in normal training. + + ``source_tokens_read`` counts tokenized source tokens entering the packer, + including document BOS tokens. ``new_source_tokens_advanced`` counts unique + source-stream tokens permanently advanced, excluding a retained FIFO carry. + Cropped tokens are source tokens advanced without being emitted. Skipped + transitions include destructive suffix loss and deliberate row segmentation + at an original BOS. Synthetic BOS tokens are never source tokens. + """ + + destructive_cropped_tokens: int = 0 + source_tokens_read: int = 0 + new_source_tokens_advanced: int = 0 + skipped_adjacent_transitions: int = 0 + synthetic_bos_tokens_inserted: int = 0 + intentional_bos_boundaries: int = 0 + buffered_source_tokens: int = 0 + source_tokens_emitted: int = 0 + source_bos_tokens_emitted: int = 0 + rows: int = 0 + batches: int = 0 + debug_discarded_suffixes: Optional[list[tuple[int, ...]]] = None + + def reset(self) -> None: + self.destructive_cropped_tokens = 0 + self.source_tokens_read = 0 + self.new_source_tokens_advanced = 0 + self.skipped_adjacent_transitions = 0 + self.synthetic_bos_tokens_inserted = 0 + self.intentional_bos_boundaries = 0 + self.buffered_source_tokens = 0 + self.source_tokens_emitted = 0 + self.source_bos_tokens_emitted = 0 + self.rows = 0 + self.batches = 0 + if self.debug_discarded_suffixes is not None: + self.debug_discarded_suffixes.clear() + # --------------------------------------------------------------------------- # Dataset # --------------------------------------------------------------------------- @@ -57,7 +100,9 @@ def __init__( self.tokenizer_threads = tokenizer_threads def __iter__(self) -> Iterator[Tuple[torch.Tensor, DataLoaderState]]: - for texts, state in self.sm.iterate(start_state=self.start_state): + # One document per source state makes checkpoint resume exact: no + # untracked documents remain suspended inside this generator. + for texts, state in self.sm.iterate(start_state=self.start_state, batch_size=1): for txt in texts: tokens = ( self.tokenizer(txt, prepend_bos=True, threads=self.tokenizer_threads) @@ -80,29 +125,13 @@ class DistDataLoader: """ Packs tokenized documents into fixed-shape (B, T) input/target tensors. - Design notes - ------------ - * Buffer is a deque of (Tensor, state) pairs — O(1) popleft and head - update, versus O(n) for list.pop(0). - - * Buffer is sized in **tokens**, not in documents. The old heuristic - `buffer_size = B * T * 16` was meant to be a token budget but was - interpreted as a document count, which could queue hundreds of - thousands of documents and consume gigabytes of memory before the - first batch. - - * Documents are separated by an explicit BOS token so the model always - sees clean boundaries: …doc_N_last [BOS] doc_N+1_first… - Without this, cross-document targets are silently trained on, - which is especially harmful for value-embedding architectures where - BOS is the per-document anchor. - - * Partial-document tail is kept as a tensor *view* (tokens[take:]), - not a Python list slice, so no allocation occurs when a document - spans multiple batches. - - * Output tensors (self.inputs, self.targets) are pre-allocated once - and reused every step. + ``stream`` (the default) is a corrected flat B*T+1 stream with one-token + carry between batches. ``bos_aligned`` starts every row with BOS and keeps + every continuation suffix, inserting a synthetic BOS before it. + + BOS alignment does not isolate documents in attention. Documents packed in + the same row can attend to earlier documents unless the model is separately + given a segment-aware attention mask. """ def __init__( @@ -112,22 +141,57 @@ def __init__( seq_len: int, device: str = "cuda", buffer_size: Optional[int] = None, - eos_token_id: Optional[int] = None, + packing_strategy: PackingStrategy = "stream", + bos_token_id: Optional[int] = None, + resume_state: Optional[DataLoaderState] = None, + packing_stats: Optional[PackingStats] = None, ): + if packing_strategy not in ("stream", "bos_aligned"): + raise ValueError(f"DistDataLoader does not implement {packing_strategy!r}") + if batch_size < 1 or seq_len < 1: + raise ValueError("batch_size and seq_len must be positive") + if packing_strategy == "bos_aligned" and bos_token_id is None: + raise ValueError("bos_aligned packing requires bos_token_id") self.dataset = dataset self.B = batch_size self.T = seq_len self.device = torch.device(device) - self.eos_token_id = eos_token_id + self.packing_strategy = packing_strategy + self.bos_token_id = bos_token_id + self.packing_stats = packing_stats or PackingStats() + self.last_batch_stats = PackingStats() - # Token budget for the pre-fetch buffer. - # Default: ~16 full batches so the GPU is never starved. + # Retained for API compatibility. On-demand reads make a checkpoint + # exact without serializing a multi-batch prefetch queue. self.token_buffer_size = buffer_size or (batch_size * seq_len * 16) self.iterator = iter(dataset) - self.buffer: deque[Tuple[torch.Tensor, DataLoaderState]] = deque() - self._buffered_tokens: int = 0 + self.buffer: deque[Tuple[torch.Tensor, Optional[DataLoaderState], bool]] = deque() self.last_state: Optional[DataLoaderState] = None + self._carry_token: Optional[int] = None + self._continuation_pending = False + self._finished = False + + state = resume_state or getattr(dataset, "start_state", None) + if state is not None and state.packing_strategy is not None: + if state.packing_strategy != packing_strategy: + raise ValueError( + f"Cannot resume {state.packing_strategy!r} state with {packing_strategy!r} packing" + ) + self.last_state = state + if state.pending_tokens: + self.buffer.append(( + torch.tensor(state.pending_tokens, dtype=torch.long), + state, + state.pending_is_document_start, + )) + self._carry_token = state.carry_token + self._continuation_pending = state.continuation_pending + saved = { + key: value for key, value in state.packing_stats.items() + if key in PackingStats.__dataclass_fields__ + } + self.packing_stats = packing_stats or PackingStats(**saved) total = batch_size * (seq_len + 1) @@ -137,7 +201,6 @@ def __init__( pin_memory=(self.device.type == "cuda"), ) self.gpu = torch.empty(total, dtype=torch.long, device=self.device) - # Single contiguous allocation; inputs and targets are non-overlapping # views into it. Avoids two separate allocations per forward pass. _out = torch.empty( @@ -146,88 +209,149 @@ def __init__( self.inputs = _out[:batch_size * seq_len].view(batch_size, seq_len) self.targets = _out[batch_size * seq_len:].view(batch_size, seq_len) - def _refill(self) -> None: - """Pull documents from the dataset until the token budget is met.""" - while self._buffered_tokens < self.token_buffer_size: - try: - tokens, state = next(self.iterator) - except StopIteration: - if getattr(self.dataset, "split", None) == "val": - # Validation set is finite — wrap around silently. - self.iterator = iter(self.dataset) - continue - # Training iterator (ShardManager.iterate) is infinite; - # reaching here means something went wrong upstream. - break - - self.buffer.append((tokens, state)) - self._buffered_tokens += len(tokens) + def _pull_document(self) -> bool: + if self.buffer: + return True + try: + tokens, state = next(self.iterator) + except StopIteration: + return False + if not isinstance(tokens, torch.Tensor): + tokens = torch.tensor(tokens, dtype=torch.long) + self.last_state = state + self.buffer.append((tokens, state, True)) + self.packing_stats.source_tokens_read += len(tokens) + self.packing_stats.buffered_source_tokens += len(tokens) + self.last_batch_stats.source_tokens_read += len(tokens) + return True + + def _copy_source(self, destination: torch.Tensor, limit: int) -> int: + if not self._pull_document(): + return 0 + tokens, state, is_start = self.buffer[0] + take = min(len(tokens), limit) + destination[:take].copy_(tokens[:take]) + for stats in (self.packing_stats, self.last_batch_stats): + stats.source_tokens_emitted += take + stats.new_source_tokens_advanced += take + if is_start and take: + stats.source_bos_tokens_emitted += 1 + self.packing_stats.buffered_source_tokens -= take + if take == len(tokens): + self.buffer.popleft() + else: + self.buffer[0] = (tokens[take:], state, False) + return take def __iter__(self): return self def __next__(self) -> Tuple[torch.Tensor, torch.Tensor, DataLoaderState]: - self._refill() - B, T = self.B, self.T - total = B * (T + 1) + if self._finished: + raise StopIteration + self.last_batch_stats = PackingStats() + if self.packing_strategy == "stream": + self._next_stream() + else: + self._next_bos_aligned() + for stats in (self.packing_stats, self.last_batch_stats): + stats.batches += 1 + stats.rows += self.B + return self.inputs, self.targets, self._state() + + def _next_stream(self) -> None: + total = self.B * self.T + 1 pos = 0 - + had_carry = self._carry_token is not None + if had_carry: + self.packing_stats.buffered_source_tokens -= 1 + self.cpu[0] = self._carry_token + pos = 1 while pos < total: - if not self.buffer: - # Mid-batch top-up (rare: only if a single doc > token_buffer_size). - self._refill() - if not self.buffer: - raise RuntimeError( - "DistDataLoader buffer is empty mid-batch. " - "The training ShardManager iterator should be infinite — " - "check shard availability and ShardManager.iterate()." - ) - - # O(1) peek at deque head — no pop yet. - tokens, state = self.buffer[0] - self.last_state = state - - remaining = total - pos - take = min(len(tokens), remaining) - - # Zero-copy write into the pinned CPU buffer. - # tokens[:take] is a tensor view (no new allocation). - self.cpu[pos : pos + take].copy_(tokens[:take]) + take = self._copy_source(self.cpu[pos:total], total - pos) + if not take: + break pos += take - self._buffered_tokens -= take - - if take < len(tokens): - # Document spans into the next batch. - # Update the head in place with a view of the remainder — - # deque[0] access and assignment are both O(1). - self.buffer[0] = (tokens[take:], state) - else: - # Document fully consumed — O(1) removal. - self.buffer.popleft() - - # Async host-to-device transfer. + if pos < total: + self._finished = True + raise StopIteration + if not had_carry: + # The first source token seeds the flat window; only the following + # B*T tokens advance the stream and contribute target positions. + for stats in (self.packing_stats, self.last_batch_stats): + stats.new_source_tokens_advanced -= 1 + self.gpu[:total].copy_(self.cpu[:total], non_blocking=(self.device.type == "cuda")) + self.inputs.copy_(self.gpu[:total - 1].view(self.B, self.T)) + self.targets.copy_(self.gpu[1:total].view(self.B, self.T)) + self._carry_token = int(self.cpu[total - 1]) + self.packing_stats.buffered_source_tokens += 1 + + def _next_bos_aligned(self) -> None: + B, capacity = self.B, self.T + 1 + rows = self.cpu[:B * capacity].view(B, capacity) + for row_index in range(B): + pos = 0 + if self._continuation_pending: + rows[row_index, 0] = self.bos_token_id + for stats in (self.packing_stats, self.last_batch_stats): + stats.synthetic_bos_tokens_inserted += 1 + stats.skipped_adjacent_transitions += 1 + self._continuation_pending = False + pos = 1 + while pos < capacity: + if not self._pull_document(): + self._finished = True + raise StopIteration + tokens, _, is_start = self.buffer[0] + if is_start and (not len(tokens) or int(tokens[0]) != self.bos_token_id): + raise ValueError("bos_aligned requires every source document to begin with BOS") + if pos == 0 and is_start: + for stats in (self.packing_stats, self.last_batch_stats): + stats.intentional_bos_boundaries += 1 + stats.skipped_adjacent_transitions += 1 + take = self._copy_source(rows[row_index, pos:capacity], capacity - pos) + pos += take + if self.buffer and pos == capacity: + self._continuation_pending = True self.gpu.copy_(self.cpu, non_blocking=(self.device.type == "cuda")) - - data = self.gpu.view(B, T + 1) + data = self.gpu.view(B, capacity) self.inputs.copy_(data[:, :-1]) self.targets.copy_(data[:, 1:]) - return self.inputs, self.targets, self.last_state + def _state(self) -> DataLoaderState: + base = self.last_state or DataLoaderState() + pending_tokens: list[int] = [] + pending_is_start = False + if self.buffer: + tokens, _, pending_is_start = self.buffer[0] + pending_tokens = tokens.tolist() + values = base.model_dump(exclude={ + "packing_strategy", "pending_tokens", "pending_is_document_start", + "carry_token", "continuation_pending", "packing_stats", + }) + return DataLoaderState( + **values, + packing_strategy=self.packing_strategy, + pending_tokens=pending_tokens, + pending_is_document_start=pending_is_start, + carry_token=self._carry_token, + continuation_pending=self._continuation_pending, + packing_stats=asdict(self.packing_stats), + ) # just to compare """ Distributed dataloaders for pretraining. -BOS-aligned bestfit: +BOS-aligned destructive best-fit: - Every row starts with BOS token - Documents packed using best-fit algorithm to minimize cropping - When no document fits remaining space, crops a document to fill exactly - 100% utilization (no padding), ~35% tokens cropped at T=2048 -Compared to the original tokenizing_distributed_data_loader: -BOS-aligned loses ~35% of tokens to cropping, but ensures that -there are fewer "confusing" tokens in the train/val batches as every token can -now attend back to the BOS token and sees the full context of the document. +This legacy strategy may lose substantial source content to cropping. Starting +a row with BOS does not isolate documents packed later in that row: causal +attention still crosses document boundaries without segment-aware masks. Fallback to the original if you have very limited data AND long documents: https://github.com/karpathy/nanochat/blob/3c3a3d7/nanochat/dataloader.py#L78-L117 @@ -295,6 +419,7 @@ def tokenizing_distributed_data_loader_with_state_bos_bestfit( tokenizer_threads=4, tokenizer_batch_size=128, device="cuda", resume_state_dict=None, buffer_size=1000, base_path=None, + packing_stats: Optional[PackingStats] = None, ): """ BOS-aligned dataloader with Best-Fit Cropping. @@ -326,10 +451,13 @@ def refill_buffer(): token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads) for tokens in token_lists: doc_buffer.append(tokens) + if packing_stats is not None: + packing_stats.source_tokens_read += len(tokens) + packing_stats.buffered_source_tokens += len(tokens) # Pre-allocate buffers once: layout is [inputs (B*T) | targets (B*T)] # This gives us contiguous views and a single HtoD transfer - use_cuda = device == "cuda" + use_cuda = torch.device(device).type == "cuda" row_buffer = torch.empty((B, row_capacity), dtype=torch.long) # for building rows without creating Python lists cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=use_cuda) # staging area (CPU) gpu_buffer = torch.empty(2 * B * T, dtype=torch.long, device=device) # on-device buffer @@ -340,6 +468,11 @@ def refill_buffer(): while True: for row_idx in range(B): + if packing_stats is not None: + # The transition into this row's original BOS is deliberately + # segmented, rather than lost accidentally or made synthetic. + packing_stats.intentional_bos_boundaries += 1 + packing_stats.skipped_adjacent_transitions += 1 pos = 0 while pos < row_capacity: # Ensure buffer has documents @@ -366,8 +499,19 @@ def refill_buffer(): # No doc fits - crop shortest in buffer to fill remaining and minimize waste shortest_idx = min(range(len(doc_buffer)), key=lambda i: len(doc_buffer[i])) doc = doc_buffer.pop(shortest_idx) - row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining], dtype=torch.long) - pos += remaining + take = min(len(doc), remaining) + row_buffer[row_idx, pos:pos + take] = torch.tensor(doc[:take], dtype=torch.long) + if packing_stats is not None: + discarded = len(doc) - take + packing_stats.destructive_cropped_tokens += discarded + packing_stats.skipped_adjacent_transitions += discarded + if packing_stats.debug_discarded_suffixes is not None and discarded: + packing_stats.debug_discarded_suffixes.append(tuple(doc[take:])) + pos += take + + if packing_stats is not None: + packing_stats.new_source_tokens_advanced += len(doc) + packing_stats.buffered_source_tokens -= len(doc) # Copy to pinned CPU buffer, then single HtoD transfer cpu_inputs.copy_(row_buffer[:, :-1]) @@ -404,11 +548,18 @@ def build_dataloader( resume_state: Optional[DataLoaderState] = None, dist_info: Optional[dict] = None, tokenizer_threads: int = 4, - use_nanochat: bool = False, + packing_strategy: PackingStrategy = "stream", + bos_token_id: Optional[int] = None, + use_nanochat: Optional[bool] = None, + packing_stats: Optional[PackingStats] = None, ) -> DistDataLoader: if dist_info is None: dist_info = get_dist_info() if use_nanochat: + if packing_strategy != "stream": + raise ValueError("Use packing_strategy instead of combining it with use_nanochat") + packing_strategy = "bos_bestfit_crop" + if packing_strategy == "bos_bestfit_crop": # This is the original dataloader from nanochat, from which I derived the pipeline for gpt-lab. # So it is based on the same underlying data loading and on-the-fly tokenization if resume_state is not None: @@ -428,10 +579,13 @@ def build_dataloader( tokenizer_batch_size=128, device=dist_info["DEVICE"], resume_state_dict=resume_state_dict, - buffer_size=buffer_size or (batch_size * seq_len * 16), - base_path=datadir / name + buffer_size=buffer_size or 1000, + base_path=(Path(datadir) if datadir is not None else DATA_DIR) / name, + packing_stats=packing_stats, ) - else: + elif packing_strategy in ("stream", "bos_aligned"): + if bos_token_id is None and tokenizer is not None and hasattr(tokenizer, "get_bos_token_id"): + bos_token_id = tokenizer.get_bos_token_id() ds = ShardedDataset( name=name, tokenizer=tokenizer, @@ -440,9 +594,10 @@ def build_dataloader( base_url=base_url, shard_limit=shard_limit, max_shards=max_shards, - cachedir=cachedir, + cachedir=cachedir or DATA_DIR, start_state=resume_state, dist_info=dist_info, + tokenizer_threads=tokenizer_threads, ) dataloader = lambda: DistDataLoader( ds, @@ -450,8 +605,14 @@ def build_dataloader( seq_len=seq_len, device=dist_info["DEVICE"], buffer_size=buffer_size, + packing_strategy=packing_strategy, + bos_token_id=bos_token_id, + resume_state=resume_state, + packing_stats=packing_stats, ) + else: + raise ValueError(f"Unknown packing strategy: {packing_strategy!r}") if split == "val": return dataloader else: - return dataloader() \ No newline at end of file + return dataloader() diff --git a/src/gpt_lab/data/sharder.py b/src/gpt_lab/data/sharder.py index 866931d..f61f280 100644 --- a/src/gpt_lab/data/sharder.py +++ b/src/gpt_lab/data/sharder.py @@ -261,15 +261,7 @@ def iterate( pf = pq.ParquetFile(shard_path) - if is_resuming: - base = state.row_group_idx // self.world_size - state.row_group_idx = (base + 1) * self.world_size + self.ddp_rank - if state.row_group_idx >= pf.num_row_groups: - state.shard_idx += 1 # go to resuming shard id - state.row_group_idx = self.ddp_rank # start at the first row group for the next shard - state.offset_in_row_group = 0 - continue - else: + if not is_resuming: state.row_group_idx = self.ddp_rank while state.row_group_idx < pf.num_row_groups: @@ -280,11 +272,14 @@ def iterate( continue if is_resuming: is_resuming = False # only do this once - yield batch[i:i+batch_size], DataLoaderState( + chunk = batch[i:i+batch_size] + yield chunk, DataLoaderState( shard_idx=state.shard_idx, global_shard_idx=state.global_shard_idx, # for debbugging - we keep track of original shard idx row_group_idx=state.row_group_idx, - offset_in_row_group=i, + # State always names the next unread document, so a + # checkpoint resumes without repeating this chunk. + offset_in_row_group=i + len(chunk), epoch=state.epoch, ) state.offset_in_row_group = 0 @@ -295,4 +290,4 @@ def iterate( state.shard_idx = 0 state.row_group_idx = self.ddp_rank state.offset_in_row_group = 0 - state.epoch += 1 \ No newline at end of file + state.epoch += 1 diff --git a/src/gpt_lab/utils/schemas.py b/src/gpt_lab/utils/schemas.py index d4a4372..e2d1583 100644 --- a/src/gpt_lab/utils/schemas.py +++ b/src/gpt_lab/utils/schemas.py @@ -33,6 +33,7 @@ LossTypes, NormalizationTypes, PositionalEncodingTypes, + PackingStrategy, TfTypes, TokenizerSources, TpModes, @@ -196,6 +197,7 @@ class DataLoaderConfig(BaseModel): buffer_size: int = 10000 device: str = "cuda" use_pin_memory: bool = False + packing_strategy: PackingStrategy = "stream" class BaseConfig(BaseModel): data_dir: Union[str, Path] = DATA_DIR @@ -616,7 +618,12 @@ class DataLoaderState(BaseModel): row_group_idx: int = 0 offset_in_row_group: int = 0 epoch: int = 1 - # Add more fields as needed to track the state of the current iteration over a data shard + packing_strategy: Optional[PackingStrategy] = None + pending_tokens: List[int] = Field(default_factory=list) + pending_is_document_start: bool = False + carry_token: Optional[int] = None + continuation_pending: bool = False + packing_stats: Dict[str, Any] = Field(default_factory=dict) class TrainerState(BaseModel): step: int = 0 @@ -723,4 +730,4 @@ def get_config_from_huggingface(model_name: str) -> TransformerConfig: tokenizer=tokenizer.encode, pad_id=pad_id, vocab_size=vocab_size - ) \ No newline at end of file + ) diff --git a/src/gpt_lab/utils/types.py b/src/gpt_lab/utils/types.py index 2877b7b..e7963bc 100644 --- a/src/gpt_lab/utils/types.py +++ b/src/gpt_lab/utils/types.py @@ -27,6 +27,7 @@ LossTypes = Literal["cross_entropy", "kl_divergence"] LossReductionTypes = Literal["none", "mean", "sum"] TpModes = Literal["row", "column"] +PackingStrategy = Literal["stream", "bos_aligned", "bos_bestfit_crop"] # str_to_torch_dtype = { # "bool": torch.bool, @@ -48,4 +49,4 @@ # str_to_torch_dtype["uint32"] = torch.uint32 # str_to_torch_dtype["uint64"] = torch.uint64 -# Dtypes = Literal[tuple(str_to_torch_dtype.keys())] \ No newline at end of file +# Dtypes = Literal[tuple(str_to_torch_dtype.keys())] diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..4a53966 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,239 @@ +from argparse import Namespace +from pathlib import Path + +import pytest +import torch + +from gpt_lab.data.loader import ( + DistDataLoader, + PackingStats, + tokenizing_distributed_data_loader_with_state_bos_bestfit, +) +from gpt_lab.utils.schemas import DataLoaderState +from scripts.benchmark.dataloaders import ( + Counters, + IdentityTokenizer, + LoaderSpec, + TorchPackedDataset, + aggregate, + benchmark, + benchmark_trial, + check_accounting, + write_split, +) + + +class DeterministicDataset: + def __init__(self, documents, split="train"): + self.documents = documents + self.split = split + + def __iter__(self): + state = DataLoaderState() + for document in self.documents: + yield torch.tensor(document, dtype=torch.long), state + + +def loader_batches(kind, documents, batch_size=2, seq_len=3): + if kind == "gpt_lab": + loader = DistDataLoader( + DeterministicDataset(documents), + batch_size=batch_size, + seq_len=seq_len, + device="cpu", + buffer_size=batch_size * seq_len, + ) + return loader + + dataset = TorchPackedDataset( + Path("."), None, batch_size, seq_len, True, "stream", 1 + ) + return dataset._stream(iter(documents)) + + +def take_batches(loader, count): + batches = [] + for _ in range(count): + inputs, targets = next(loader)[:2] + batches.append((inputs.clone(), targets.clone())) + return batches + + +def flatten_batches(batches): + windows = [torch.cat((inputs.flatten(), targets.flatten()[-1:])) for inputs, targets in batches] + return torch.cat((windows[0], *[window[1:] for window in windows[1:]])) + + +@pytest.mark.parametrize("kind", ["gpt_lab", "pytorch"]) +def test_flat_stream_preserves_every_transition_across_rows_and_batches(kind): + loader = loader_batches(kind, [list(range(30))]) + batches = take_batches(loader, 3) + + for inputs, targets in batches: + assert inputs.shape == targets.shape == (2, 3) + flat = torch.cat((inputs.flatten(), targets.flatten()[-1:])) + assert torch.equal(targets.flatten(), flat[1:]) + assert targets[0, -1] == inputs[1, 0] + + for (_, targets), (inputs, _) in zip(batches, batches[1:]): + assert inputs.flatten()[0] == targets.flatten()[-1] + + observed = flatten_batches(batches) + assert torch.equal(observed, torch.arange(19)) + + +@pytest.mark.parametrize("kind", ["gpt_lab", "pytorch"]) +def test_flat_stream_keeps_document_tails_and_bos_tokens(kind): + bos = 99 + documents = [ + [bos, *range(0, 10)], + [bos, *range(10, 20)], + [bos, *range(20, 30)], + ] + expected = torch.tensor([token for document in documents for token in document]) + loader = loader_batches(kind, documents) + batches = take_batches(loader, 5) + observed = flatten_batches(batches) + + assert torch.equal(observed, expected[: len(observed)]) + assert observed.tolist().count(bos) == expected[: len(observed)].tolist().count(bos) + + +@pytest.mark.parametrize("kind", ["gpt_lab", "pytorch"]) +def test_fifo_accounting_has_no_crop_or_skipped_transitions(kind): + documents = [list(range(30))] + if kind == "gpt_lab": + stats = PackingStats() + loader = DistDataLoader( + DeterministicDataset(documents), 2, 3, device="cpu", buffer_size=6, + packing_stats=stats, + ) + buffered_start = stats.buffered_source_tokens + batches = take_batches(loader, 2) + buffered_delta = stats.buffered_source_tokens - buffered_start + values = vars(stats) + else: + dataset = TorchPackedDataset(Path("."), None, 2, 3, True, "stream", 1) + loader = dataset._stream(iter(documents)) + batch_stats = [] + batches = [] + for _ in range(2): + inputs, targets, current = next(loader) + batches.append((inputs, targets)) + batch_stats.append(current) + values = { + key: sum(row[key] for row in batch_stats) + for key in ( + "source_tokens_read", "new_source_tokens_advanced", + "destructive_cropped_tokens", "skipped_adjacent_transitions", + "synthetic_bos_tokens_inserted", + ) + } + buffered_delta = sum(row["buffered_source_tokens_delta"] for row in batch_stats) + + targets = sum(target.numel() for _, target in batches) + assert values["destructive_cropped_tokens"] == 0 + assert values["skipped_adjacent_transitions"] == 0 + assert targets == 2 * 2 * 3 + check_accounting( + source_tokens_read=values["source_tokens_read"], + new_source_tokens_advanced=values["new_source_tokens_advanced"], + buffered_source_tokens_delta=buffered_delta, + target_positions_emitted=targets, + skipped_adjacent_transitions=values["skipped_adjacent_transitions"], + synthetic_bos_tokens_inserted=values["synthetic_bos_tokens_inserted"], + ) + + +def test_bos_bestfit_reports_exact_discarded_suffix(tmp_path): + bos = 99 + write_split(tmp_path, [[bos, 1, 2, 3, 4, 5]], row_group_size=1) + stats = PackingStats(debug_discarded_suffixes=[]) + loader = tokenizing_distributed_data_loader_with_state_bos_bestfit( + IdentityTokenizer(bos, 128), B=1, T=3, split="train", device="cpu", + tokenizer_batch_size=1, buffer_size=1, base_path=tmp_path, + packing_stats=stats, + ) + + inputs, targets, _ = next(loader) + + assert inputs.tolist() == [[bos, 1, 2]] + assert targets.tolist() == [[1, 2, 3]] + assert stats.destructive_cropped_tokens == 2 + assert stats.debug_discarded_suffixes == [(4, 5)] + assert stats.skipped_adjacent_transitions == 3 # two crop skips + one BOS row boundary + assert stats.intentional_bos_boundaries == 1 + + +def test_synthetic_bos_is_not_counted_as_source(): + # Four source tokens advance; one real adjacency is segmented and replaced + # by one synthetic-BOS target. The synthetic token never enters source read. + check_accounting( + source_tokens_read=4, + new_source_tokens_advanced=4, + buffered_source_tokens_delta=0, + target_positions_emitted=4, + skipped_adjacent_transitions=1, + synthetic_bos_tokens_inserted=1, + ) + + +def test_benchmark_excludes_warmup_counters(): + def factory(counters: Counters): + def batches(): + while True: + counters.source_tokens_read += 4 + counters.new_source_tokens_advanced += 4 + tokens = torch.tensor([[99, 1, 2, 3]]) + yield tokens, tokens, {} + return batches() + + args = Namespace( + warmup_batches=1, batches=2, device=torch.device("cpu"), + batch_size=1, seq_len=4, torch_packing="stream", + ) + result = benchmark( + "fake", factory, "pretokenized", "identity", 99, None, + "none", "none", None, None, args, 0.0, + ) + + assert result.source_tokens_read == 8 + assert result.new_source_tokens_advanced == 8 + assert result.target_positions_emitted == 8 + + +def test_benchmark_trial_reports_memory_pressure(): + def factory(): + def batches(): + while True: + tokens = torch.tensor([[99, 1, 2, 3]]) + yield tokens, tokens, { + "source_tokens_read": 4, + "new_source_tokens_advanced": 4, + "destructive_cropped_tokens": 0, + "buffered_source_tokens_delta": 0, + "skipped_adjacent_transitions": 0, + "synthetic_bos_tokens_inserted": 0, + "intentional_bos_boundaries": 0, + "actual_buffered_tokens": 0, + "actual_buffered_documents": 0, + } + return batches() + + spec = LoaderSpec("fake", "fake", "flat_stream", "stream_packing", "test", factory) + args = Namespace( + warmup_batches=1, batches=2, device=torch.device("cpu"), + batch_size=1, seq_len=4, trials=1, best_fit_buffer_docs=1, + correctness_batches=1, + ) + + trial = benchmark_trial(spec, args) + _, _, _, memory = trial + result = aggregate(spec, [trial], 0.5, "pretokenized", "identity", 0.0, args) + + assert memory["host_rss_peak_mib"] > 0 + assert memory["host_rss_peak_delta_mib"] >= 0 + assert memory["accelerator_peak_allocated_mib"] is None + assert memory["accelerator_peak_delta_mib"] is None + assert result.host_rss_peak_mib == memory["host_rss_peak_mib"] + assert result.accelerator_peak_allocated_mib is None diff --git a/tests/test_dataloader_destructive_crop.py b/tests/test_dataloader_destructive_crop.py new file mode 100644 index 0000000..732bd7b --- /dev/null +++ b/tests/test_dataloader_destructive_crop.py @@ -0,0 +1,191 @@ +from itertools import cycle +from types import SimpleNamespace + +import torch + +import gpt_lab.data.loader as loader_module +from gpt_lab.data.loader import DistDataLoader, PackingStats +from scripts.benchmark.dataloaders import ( + Counters, + IdentityTokenizer, + TorchPackedDataset, + benchmark, +) + + +class IdentifiableDocuments: + split = "train" + + def __init__(self, documents): + self.documents = documents + + def __iter__(self): + for document in cycle(self.documents): + yield torch.tensor(document), None + + +def packed_tokens(inputs, targets): + return [*inputs.flatten().tolist(), targets.flatten().tolist()[-1]] + + +def test_gpt_lab_fifo_retains_long_document_tail(): + stats = PackingStats() + loader = DistDataLoader( + IdentifiableDocuments(([1, 2, 3, 4, 5, 6, 7, 8], [20, 21, 22])), + batch_size=1, + seq_len=4, + device="cpu", + buffer_size=1, + packing_stats=stats, + ) + + first_inputs, first_targets, _ = next(loader) + first_inputs, first_targets = first_inputs.clone(), first_targets.clone() + second_inputs, second_targets, _ = next(loader) + + assert packed_tokens(first_inputs, first_targets) == [1, 2, 3, 4, 5] + assert packed_tokens(second_inputs, second_targets) == [5, 6, 7, 8, 20] + assert stats.destructive_cropped_tokens == 0 + + +def test_pytorch_stream_retains_long_document_tail(): + dataset = TorchPackedDataset( + path=None, + tokenizer=IdentityTokenizer(0, 100), + batch_size=1, + seq_len=4, + tokenized=True, + packing="stream", + buffer_docs=1, + ) + stream = dataset._stream(iter(([1, 2, 3, 4, 5, 6, 7, 8], [20, 21, 22]))) + + first_inputs, first_targets, first_stats = next(stream) + second_inputs, second_targets, second_stats = next(stream) + + assert packed_tokens(first_inputs, first_targets) == [1, 2, 3, 4, 5] + assert packed_tokens(second_inputs, second_targets) == [5, 6, 7, 8, 20] + assert first_stats["destructive_cropped_tokens"] == 0 + assert second_stats["destructive_cropped_tokens"] == 0 + + +def test_nanochat_bestfit_counts_awkward_suffix_and_long_document(monkeypatch): + document_batches = iter( + ( + ([[0, 31, 32, 33], [0, 11, 12], [0, 21, 22, 23, 24, 25, 26, 27]], (0, 0, 1)), + ([[0, 41, 42, 43, 44, 45, 46, 47], [0, 51, 52, 53, 54, 55, 56, 57, 58]], (0, 1, 1)), + ) + ) + monkeypatch.setattr(loader_module, "_document_batches", lambda *args, **kwargs: document_batches) + stats = PackingStats() + loader = loader_module.tokenizing_distributed_data_loader_with_state_bos_bestfit( + IdentityTokenizer(0, 100), + B=1, + T=4, + split="train", + device="cpu", + buffer_size=3, + packing_stats=stats, + ) + + inputs, targets, _ = next(loader) + + assert packed_tokens(inputs, targets) == [0, 31, 32, 33, 0] + assert stats.destructive_cropped_tokens == 2 + + long_batches = cycle( + [([[0, 61, 62, 63, 64, 65, 66, 67], [0, 71, 72, 73, 74, 75, 76, 77, 78]], (0, 0, 1))] + ) + monkeypatch.setattr(loader_module, "_document_batches", lambda *args, **kwargs: long_batches) + long_stats = PackingStats() + long_loader = loader_module.tokenizing_distributed_data_loader_with_state_bos_bestfit( + IdentityTokenizer(0, 100), + B=1, + T=4, + split="train", + device="cpu", + buffer_size=2, + packing_stats=long_stats, + ) + + long_inputs, long_targets, _ = next(long_loader) + + assert packed_tokens(long_inputs, long_targets) == [0, 61, 62, 63, 64] + assert long_stats.destructive_cropped_tokens == 3 + + +def test_pytorch_bestfit_counts_awkward_suffix_and_long_document(): + dataset = TorchPackedDataset( + path=None, + tokenizer=IdentityTokenizer(0, 100), + batch_size=1, + seq_len=4, + tokenized=True, + packing="bestfit", + buffer_docs=3, + ) + awkward = dataset._bestfit( + cycle(([0, 31, 32, 33], [0, 11, 12], [0, 21, 22, 23, 24, 25, 26, 27])) + ) + + inputs, targets, stats = next(awkward) + + assert packed_tokens(inputs, targets) == [0, 31, 32, 33, 0] + assert stats["destructive_cropped_tokens"] == 2 + + long_dataset = TorchPackedDataset( + path=None, + tokenizer=IdentityTokenizer(0, 100), + batch_size=1, + seq_len=4, + tokenized=True, + packing="bestfit", + buffer_docs=2, + ) + long = long_dataset._bestfit( + iter(([0, 61, 62, 63, 64, 65, 66, 67], [0, 71, 72, 73, 74, 75, 76, 77, 78])) + ) + + long_inputs, long_targets, long_stats = next(long) + + assert packed_tokens(long_inputs, long_targets) == [0, 61, 62, 63, 64] + assert long_stats["destructive_cropped_tokens"] == 3 + + +def test_benchmark_resets_destructive_crop_counter_after_warmup(): + def factory(counters: Counters): + def batches(): + for discarded in (7, 3): + counters.destructive_cropped_tokens += discarded + counters.source_tokens_read += 5 + discarded + counters.new_source_tokens_advanced += 5 + discarded + counters.skipped_adjacent_transitions += 1 + discarded + data = torch.tensor([[0, 1, 2, 3, 4]]) + yield data[:, :-1], data[:, 1:], None + + return batches() + + args = SimpleNamespace( + warmup_batches=1, + batches=1, + batch_size=1, + seq_len=4, + device=torch.device("cpu"), + ) + + result = benchmark( + "nanochat", + factory, + "pretokenized", + "identity", + 0, + None, + "none", + "none", + None, + None, + args, + 0.0, + ) + + assert result.destructively_cropped_tokens == 3 diff --git a/tests/test_packing_strategies.py b/tests/test_packing_strategies.py new file mode 100644 index 0000000..b33ad1f --- /dev/null +++ b/tests/test_packing_strategies.py @@ -0,0 +1,177 @@ +import inspect + +import pytest +import pyarrow as pa +import pyarrow.parquet as pq +import torch + +from gpt_lab.data.loader import DistDataLoader, build_dataloader +from gpt_lab.utils.schemas import DataLoaderState + + +BOS = 99 + + +class ResumableDocuments: + split = "train" + + def __init__(self, documents, start_state=None): + self.documents = documents + self.start_state = start_state + + def __iter__(self): + start = self.start_state.offset_in_row_group if self.start_state else 0 + for index, document in enumerate(self.documents[start:], start): + yield torch.tensor(document), DataLoaderState(offset_in_row_group=index + 1) + + +class IdentityTokenizer: + def get_bos_token_id(self): + return BOS + + def __call__(self, tokens, **_): + return tokens + + +def make_loader(documents, *, strategy="bos_aligned", batch_size=2, seq_len=4, state=None): + dataset = ResumableDocuments(documents, state) + return DistDataLoader( + dataset, + batch_size=batch_size, + seq_len=seq_len, + device="cpu", + packing_strategy=strategy, + bos_token_id=BOS, + resume_state=state, + ) + + +def clone_batch(batch): + inputs, targets, state = batch + return inputs.clone(), targets.clone(), state + + +def source_rows(batch): + inputs, targets, _ = batch + rows = [] + for row_inputs, row_targets in zip(inputs, targets): + rows.append([int(row_inputs[0]), *row_targets.tolist()]) + return rows + + +def drain(loader): + batches = [] + while True: + try: + batches.append(clone_batch(next(loader))) + except StopIteration: + return batches + + +def test_bos_aligned_rows_start_with_bos_and_batch_size_gt_one(): + loader = make_loader([[BOS, *range(1, 18)], [BOS, 20], [BOS, 21, 22]], batch_size=3) + for inputs, targets, _ in drain(loader): + assert inputs.shape == targets.shape == (3, 4) + assert torch.all(inputs[:, 0] == BOS) + + +def test_bos_aligned_emits_each_source_content_once_and_counts_synthetic_bos(): + documents = [ + [BOS, 1, 2], # shorter than row capacity + [BOS, 3, 4, 5, 6], # equal to row capacity + [BOS, 7, 8, 9, 10, 11, 12, 13, 14], # longer than row capacity + ] + loader = make_loader(documents) + rows = [row for batch in drain(loader) for row in source_rows(batch)] + observed_content = [token for row in rows for token in row if token != BOS] + expected_content = [token for document in documents for token in document[1:]] + + assert observed_content == expected_content + assert loader.packing_stats.source_tokens_read == sum(map(len, documents)) + assert loader.packing_stats.source_tokens_emitted == sum(map(len, documents)) + assert loader.packing_stats.source_bos_tokens_emitted == len(documents) + assert loader.packing_stats.synthetic_bos_tokens_inserted > 0 + assert loader.packing_stats.destructive_cropped_tokens == 0 + + +def test_bos_aligned_packs_multiple_short_documents_in_one_row(): + loader = make_loader([[BOS, 1], [BOS, 2], [BOS, 3]], batch_size=1, seq_len=5) + batch = clone_batch(next(loader)) + assert source_rows(batch)[0] == [BOS, 1, BOS, 2, BOS, 3] + + +def test_bos_aligned_long_document_spans_rows_and_batches_without_loss(): + document = [BOS, *range(1, 25)] + loader = make_loader([document], batch_size=2, seq_len=4) + batches = drain(loader) + rows = [row for batch in batches for row in source_rows(batch)] + + assert len(batches) > 1 + assert [token for row in rows for token in row if token != BOS] == document[1:] + assert loader.packing_stats.synthetic_bos_tokens_inserted == len(rows) - 1 + + +def test_validation_iteration_is_deterministic(): + documents = [[BOS, 1, 2], [BOS, 3], [BOS, *range(4, 14)]] + first = drain(make_loader(documents)) + second = drain(make_loader(documents)) + assert len(first) == len(second) + for (x1, y1, _), (x2, y2, _) in zip(first, second): + assert torch.equal(x1, x2) + assert torch.equal(y1, y2) + + +@pytest.mark.parametrize("strategy", ["stream", "bos_aligned"]) +def test_checkpoint_resume_continues_at_exact_next_batch(strategy): + documents = [[BOS, *range(1, 10)], [BOS, *range(10, 30)]] + loader = make_loader(documents, strategy=strategy, batch_size=2, seq_len=4) + _, _, state = clone_batch(next(loader)) + expected_inputs, expected_targets, _ = clone_batch(next(loader)) + + resumed = make_loader(documents, strategy=strategy, batch_size=2, seq_len=4, state=state) + actual_inputs, actual_targets, _ = clone_batch(next(resumed)) + assert torch.equal(actual_inputs, expected_inputs) + assert torch.equal(actual_targets, expected_targets) + + +def test_default_strategy_remains_stream(): + assert inspect.signature(build_dataloader).parameters["packing_strategy"].default == "stream" + loader = make_loader([[BOS, *range(20)]], strategy="stream") + assert loader.packing_strategy == "stream" + + +@pytest.mark.parametrize("strategy", ["stream", "bos_aligned"]) +def test_incomplete_batches_are_not_emitted(strategy): + loader = make_loader([[BOS, 1]], strategy=strategy, batch_size=1, seq_len=4) + with pytest.raises(StopIteration): + next(loader) + + +def test_build_dataloader_validation_and_shard_resume(tmp_path): + dataset = tmp_path / "packing" + dataset.mkdir() + documents = [[BOS, *range(i, i + 7)] for i in range(0, 42, 7)] + table = pa.table({"text": documents}) + pq.write_table(table, dataset / "shard_00000.parquet", row_group_size=3) + pq.write_table(table, dataset / "shard_00001.parquet", row_group_size=3) + common = dict( + name="packing", + tokenizer=IdentityTokenizer(), + cachedir=tmp_path, + batch_size=2, + seq_len=4, + max_shards=1, + packing_strategy="bos_aligned", + dist_info={"RANK": 0, "WORLD_SIZE": 1, "DEVICE": "cpu"}, + ) + + loader = build_dataloader(split="train", **common) + _, _, state = clone_batch(next(loader)) + expected = clone_batch(next(loader))[:2] + resumed = build_dataloader(split="train", resume_state=state, **common) + actual = clone_batch(next(resumed))[:2] + assert all(torch.equal(left, right) for left, right in zip(expected, actual)) + + validation = build_dataloader(split="val", **common) + first, second = clone_batch(next(validation()))[:2], clone_batch(next(validation()))[:2] + assert all(torch.equal(left, right) for left, right in zip(first, second))