ViperEkura
d7cd69fef5
- StreamingSeqDataset yields windows sequentially through each shard - Shard-level shuffle, distributed and multi-worker shard partitioning - __len__ returns total window count for scheduler total_steps - Better OS page-cache locality than random-access map-style datasets
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
"""Streaming IterableDataset for pre-training with shard-level shuffle.
|
|
|
|
Unlike the map-style datasets, the streaming dataset yields windows
|
|
sequentially through each data shard — no random access, no sampler.
|
|
Each DataLoader worker independently streams its assigned shard subset,
|
|
giving better OS page-cache locality for large-scale (TB+) datasets.
|
|
|
|
Key properties:
|
|
- Implements ``torch.utils.data.IterableDataset``.
|
|
- ``__len__`` returns total window count so ``compute_total_steps`` works.
|
|
- Shard-level shuffle with deterministic seed (reproducible across runs).
|
|
- Distributed: each rank gets a disjoint subset of shards.
|
|
- Multi-worker: each worker within a rank gets a disjoint subset.
|
|
"""
|
|
|
|
import random
|
|
from typing import Iterator, Optional
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch import Tensor
|
|
from torch.utils.data import IterableDataset
|
|
|
|
from astrai.dataset.storage import Store
|
|
|
|
|
|
def _resolve_rank_and_world_size() -> tuple[int, int]:
|
|
if dist.is_available() and dist.is_initialized():
|
|
return dist.get_rank(), dist.get_world_size()
|
|
return 0, 1
|
|
|
|
|
|
def _total_windows(token_count, window_size, stride):
|
|
if token_count <= window_size:
|
|
return 0
|
|
return (token_count - 1 - window_size) // stride + 1
|
|
|
|
|
|
class StreamingSeqDataset(IterableDataset):
|
|
"""Streaming next-token prediction dataset.
|
|
|
|
Yields ``{"input_ids": [L], "target_ids": [L]}`` dicts by sliding a
|
|
window sequentially through each data shard. Shards are shuffled
|
|
deterministically. Distributed and multi-worker DataLoader modes are
|
|
supported: each consumer gets a disjoint shard subset.
|
|
|
|
Args:
|
|
store: Already-loaded Store with a ``"sequence"`` key.
|
|
window_size: Context length per sample.
|
|
stride: Step between consecutive windows (default: window_size).
|
|
shuffle: Shuffle shard order.
|
|
seed: Base seed for deterministic shard shuffle.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
store: Store,
|
|
window_size: int,
|
|
stride: Optional[int] = None,
|
|
shuffle: bool = True,
|
|
seed: int = 42,
|
|
rank: Optional[int] = None,
|
|
world_size: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
if window_size <= 0:
|
|
raise ValueError("window_size must be positive")
|
|
self.store = store
|
|
self.window_size = window_size
|
|
self.stride = stride if stride is not None else window_size
|
|
self.shuffle = shuffle
|
|
self.seed = seed
|
|
self._rank, self._world_size = (
|
|
rank,
|
|
world_size if rank is not None else _resolve_rank_and_world_size(),
|
|
)
|
|
|
|
if "sequence" not in store.keys:
|
|
raise KeyError(
|
|
f"Store is missing required key 'sequence'; "
|
|
f"available keys: {sorted(store.keys)}"
|
|
)
|
|
|
|
@property
|
|
def num_samples(self) -> int:
|
|
return _total_windows(self.store.token_count, self.window_size, self.stride)
|
|
|
|
def __len__(self) -> int:
|
|
return self.num_samples
|
|
|
|
def __iter__(self) -> Iterator[dict[str, Tensor]]:
|
|
segments = self.store._data["sequence"]
|
|
n_shards = len(segments)
|
|
|
|
indices = list(range(n_shards))
|
|
if self.shuffle:
|
|
rng = random.Random(self.seed)
|
|
rng.shuffle(indices)
|
|
|
|
worker_info = torch.utils.data.get_worker_info()
|
|
if worker_info is None:
|
|
num_consumers = self._world_size
|
|
consumer_id = self._rank
|
|
else:
|
|
num_consumers = self._world_size * worker_info.num_workers
|
|
consumer_id = self._rank * worker_info.num_workers + worker_info.id
|
|
|
|
my_shards = [
|
|
i for idx, i in enumerate(indices) if idx % num_consumers == consumer_id
|
|
]
|
|
|
|
for shard_idx in my_shards:
|
|
segment = segments[shard_idx]
|
|
seq_len = segment.shape[0]
|
|
for begin in range(0, seq_len - self.window_size, self.stride):
|
|
end = begin + self.window_size
|
|
yield {
|
|
"input_ids": torch.as_tensor(segment[begin:end], dtype=torch.long),
|
|
"target_ids": torch.as_tensor(
|
|
segment[begin + 1 : end + 1], dtype=torch.long
|
|
),
|
|
}
|