134 lines
5.3 KiB
Python
134 lines
5.3 KiB
Python
"""
|
||
Dataset classes for LLM training.
|
||
|
||
TextDataset: Sliding window (stride 1) over a memory-mapped uint16 binary file.
|
||
PackedDataset: Non-overlapping windows (stride = seq_len) over the same file format.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from typing import Tuple, Union
|
||
|
||
import numpy as np
|
||
import torch
|
||
from torch.utils.data import Dataset
|
||
|
||
|
||
class TextDataset(Dataset):
|
||
"""
|
||
Sliding-window dataset over a memory-mapped numpy uint16 binary token file.
|
||
|
||
Each sample is a (input_ids, targets) pair of length seq_len, where
|
||
targets is input_ids shifted by one position. Windows overlap by
|
||
(seq_len - 1) tokens, i.e. stride = 1.
|
||
|
||
Args:
|
||
data_path: Path to the .bin file produced by data/prepare.py.
|
||
seq_len: Number of tokens per sample (context length).
|
||
"""
|
||
|
||
def __init__(self, data_path: Union[str, Path], seq_len: int) -> None:
|
||
super().__init__()
|
||
self.seq_len = seq_len
|
||
path = Path(data_path)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"Data file not found: {path}")
|
||
# Memory-map for zero-copy random access.
|
||
self.data: np.ndarray = np.memmap(path, dtype="uint16", mode="r")
|
||
# Hint OS to preload entire file into page cache (2.2TB RAM available)
|
||
import mmap as _mmap
|
||
try:
|
||
self.data._mmap.madvise(_mmap.MADV_SEQUENTIAL)
|
||
except (AttributeError, OSError):
|
||
pass # madvise not available on all platforms
|
||
if len(self.data) < seq_len + 1:
|
||
raise ValueError(
|
||
f"Data file has only {len(self.data)} tokens, "
|
||
f"need at least {seq_len + 1}."
|
||
)
|
||
|
||
def __len__(self) -> int:
|
||
# Each window needs seq_len tokens plus one extra for the target shift.
|
||
return len(self.data) - self.seq_len
|
||
|
||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
# Slice from the memmap (returns a uint16 numpy view).
|
||
chunk = self.data[idx : idx + self.seq_len + 1]
|
||
# Cast to int32 (not int64) to halve CPU worker memory usage:
|
||
# uint16 (2 B) → int32 (4 B) instead of uint16 → int64 (8 B, 4× bloat).
|
||
# int32 is sufficient for vocab_size=64000 (max token id 65535 fits in int32).
|
||
# The int32→int64 (long) promotion happens on GPU inside _step(), for free.
|
||
chunk = torch.from_numpy(chunk.astype(np.int32))
|
||
input_ids = chunk[:-1] # [seq_len]
|
||
targets = chunk[1:] # [seq_len]
|
||
return input_ids, targets
|
||
|
||
|
||
class PackedDataset(Dataset):
|
||
"""
|
||
Non-overlapping packed dataset over a memory-mapped uint16 binary token file.
|
||
|
||
Intended for data that has already been packed (documents concatenated with
|
||
EOS tokens). Windows do not overlap; stride = seq_len.
|
||
|
||
The target sequence is shifted by one token relative to input_ids. Because
|
||
the last token of a window shares its target with the *first* token of the
|
||
next window, the final target position is filled with -1 (the standard
|
||
``ignore_index`` for ``nn.CrossEntropyLoss``).
|
||
|
||
Args:
|
||
data_path: Path to the .bin file produced by data/prepare.py.
|
||
seq_len: Number of tokens per sample (context length).
|
||
"""
|
||
|
||
def __init__(self, data_path: Union[str, Path], seq_len: int) -> None:
|
||
super().__init__()
|
||
self.seq_len = seq_len
|
||
path = Path(data_path)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"Data file not found: {path}")
|
||
self.data: np.ndarray = np.memmap(path, dtype="uint16", mode="r")
|
||
# Optimize mmap for shuffled random access pattern (DistributedSampler)
|
||
import mmap as _mmap
|
||
try:
|
||
self.data._mmap.madvise(_mmap.MADV_RANDOM) # disable kernel read-ahead (random access)
|
||
self.data._mmap.madvise(_mmap.MADV_WILLNEED) # async prefault into page cache
|
||
except (AttributeError, OSError):
|
||
pass
|
||
if len(self.data) < seq_len:
|
||
raise ValueError(
|
||
f"Data file has only {len(self.data)} tokens, "
|
||
f"need at least {seq_len}."
|
||
)
|
||
|
||
def __len__(self) -> int:
|
||
return len(self.data) // self.seq_len
|
||
|
||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
start = idx * self.seq_len
|
||
end = start + self.seq_len
|
||
|
||
# Cast to int32 (not int64) to halve CPU worker memory usage.
|
||
# int32 is sufficient for vocab_size=64000; int32→long promotion on GPU.
|
||
input_ids = torch.from_numpy(
|
||
self.data[start:end].astype(np.int32)
|
||
) # [seq_len]
|
||
|
||
# Targets are shifted by one. If end < len(data) we can read the
|
||
# extra token normally; otherwise pad the last position with -1.
|
||
if end < len(self.data):
|
||
targets = torch.from_numpy(
|
||
self.data[start + 1 : end + 1].astype(np.int32)
|
||
) # [seq_len]
|
||
else:
|
||
# Last window: all but the final position can be computed.
|
||
# Use int32 for the filled portion; -1 fits in int32.
|
||
targets = torch.full((self.seq_len,), fill_value=-1, dtype=torch.int32)
|
||
if end - start - 1 > 0:
|
||
targets[: self.seq_len - 1] = torch.from_numpy(
|
||
self.data[start + 1 : end].astype(np.int32)
|
||
)
|
||
|
||
return input_ids, targets
|