| """Document-level Causal LM dataset (traditional pretraining packing). |
| |
| On-disk layout (mirrors V1 QA layout, but with document indices):: |
| |
| <dataset_path>/ |
| metadata.json # same schema as V1 (tokenizer_info, max_seq_len, total_length) |
| tokens.npy # int token stream (mmap) |
| epoch_0/ |
| doc_start.npy # int64, start offset into tokens.npy |
| doc_len.npy # int32/int64, length in tokens (>= 2) |
| epoch_1/ ... |
| |
| Each document is packed as a standard next-token LM example: |
| inputs = doc[:-1], labels = doc[1:] |
| prefix_lens = 0, causal_lens = len(inputs) # full causal attention |
| |
| Use with ``data.format=document`` (and preferably ``lm_mode=causal``). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, fields |
| from pathlib import Path |
| from typing import Optional, Any |
| import os |
| import json |
|
|
| import numpy as np |
| import pydantic |
|
|
| import torch |
| from torch.utils.data import IterableDataset, get_worker_info |
|
|
| from models.common import IGNORE_LABEL_ID |
| from models.prefixlm_attention import compute_aux_seq_tensors_scalars |
| from models.layers import find_multiple |
| from multipack_sampler import MultipackDistributedBatchSampler |
|
|
|
|
| class DocumentDatasetConfig(pydantic.BaseModel): |
| seed: int |
| dataset_path: str |
| batch_max_length: int |
| drop_last_batch: bool |
|
|
| rank: int |
| num_replicas: int |
| fixed_epoch: Optional[int] = None |
| start_data_epoch: int = 0 |
| resume_sampler_start_index: int = 0 |
|
|
|
|
| class DocumentDatasetMeta(pydantic.BaseModel): |
| tokenizer_info: dict[str, Any] = {} |
| vocab_size: Optional[int] = None |
| max_seq_len: int |
| total_length: int |
| |
| layout: Optional[str] = None |
|
|
|
|
| @dataclass |
| class DocumentDatasetIndices: |
| doc_start: np.ndarray |
| doc_len: np.ndarray |
|
|
|
|
| class DocumentDataset(IterableDataset): |
| """Packed document Causal LM loader (drop-in sibling of ``V1Dataset``).""" |
|
|
| def __init__(self, config: DocumentDatasetConfig): |
| super().__init__() |
| self.config = config |
| self.metadata = self._load_metadata() |
|
|
| self._data: Optional[np.ndarray] = None |
| self._data_indices: Optional[DocumentDatasetIndices] = None |
| self._sampler: Optional[MultipackDistributedBatchSampler] = None |
| self._epoch = int(config.start_data_epoch) |
| self._resume_sampler_start = int(config.resume_sampler_start_index) |
| self._active_data_epoch: Optional[int] = None |
|
|
| def _load_metadata(self) -> DocumentDatasetMeta: |
| with open(os.path.join(self.config.dataset_path, "metadata.json"), "r") as f: |
| metadata = DocumentDatasetMeta(**json.load(f)) |
| metadata.max_seq_len -= 1 |
| assert metadata.vocab_size is None |
| metadata.vocab_size = find_multiple(metadata.tokenizer_info.pop("vocab_size"), 256) |
| return metadata |
|
|
| def _count_data_epochs(self) -> int: |
| cached = getattr(self, "_num_data_epochs_cached", None) |
| if cached is not None: |
| return int(cached) |
| root = Path(self.config.dataset_path) |
| n = sum(1 for p in root.glob("epoch_*") if p.is_dir()) |
| if n <= 0: |
| raise FileNotFoundError(f"no epoch_* dirs under {root}") |
| self._num_data_epochs_cached = n |
| return n |
|
|
| def _load_dataset_before_epoch_begin(self): |
| if self._data is None: |
| self._data = np.load(os.path.join(self.config.dataset_path, "tokens.npy"), mmap_mode="r") |
|
|
| if self.config.fixed_epoch is not None: |
| epoch_idx = int(self.config.fixed_epoch) |
| else: |
| epoch_idx = int(self._epoch) % self._count_data_epochs() |
| self._active_data_epoch = int(epoch_idx) |
| epoch_dir = os.path.join(self.config.dataset_path, f"epoch_{epoch_idx}") |
| self._data_indices = DocumentDatasetIndices( |
| **{ |
| f.name: np.load(os.path.join(epoch_dir, f"{f.name}.npy"), mmap_mode="r") |
| for f in fields(DocumentDatasetIndices) |
| } |
| ) |
| if self.config.fixed_epoch is None: |
| self._epoch += 1 |
|
|
| |
| lengths = np.maximum(self._data_indices.doc_len.astype(np.int64) - 1, 1) |
| |
| if int(self.config.seed) != 0: |
| rng = np.random.RandomState(int(self.config.seed) + int(epoch_idx) * 10007) |
| perm = rng.permutation(int(lengths.size)) |
| lengths = lengths[perm] |
| self._data_indices = DocumentDatasetIndices( |
| doc_start=self._data_indices.doc_start[perm], |
| doc_len=self._data_indices.doc_len[perm], |
| ) |
| self._sampler = MultipackDistributedBatchSampler( |
| lengths=lengths, |
| batch_max_length=self.config.batch_max_length, |
| drop_last_batch=self.config.drop_last_batch, |
| rank=self.config.rank, |
| num_replicas=self.config.num_replicas, |
| ) |
|
|
| def _load_batch(self, indices: np.ndarray): |
| assert self._data is not None and self._data_indices is not None |
|
|
| docs = [] |
| starts = self._data_indices.doc_start[indices] |
| lens = self._data_indices.doc_len[indices] |
| for s, n in zip(starts, lens): |
| n = int(n) |
| if n < 2: |
| continue |
| docs.append(self._data[int(s) : int(s) + n].astype(np.int32)) |
|
|
| batch = {k: [] for k in ["inputs", "labels", "position_ids"]} |
| |
| seqlens_prefix = [] |
| seqlens_causal = [] |
| for doc in docs: |
| inputs = doc[:-1] |
| labels = doc[1:] |
| batch["inputs"].append(inputs) |
| batch["labels"].append(labels) |
| batch["position_ids"].append(np.arange(len(inputs), dtype=np.int32)) |
| seqlens_prefix.append(0) |
| seqlens_causal.append(len(inputs)) |
|
|
| if not batch["inputs"]: |
| |
| pad_len = self.config.batch_max_length |
| empty = { |
| "inputs": np.zeros(pad_len, dtype=np.int32), |
| "labels": np.full(pad_len, IGNORE_LABEL_ID, dtype=np.int32), |
| "position_ids": np.zeros(pad_len, dtype=np.int32), |
| } |
| tensors, scalars = compute_aux_seq_tensors_scalars( |
| np.zeros(0, dtype=np.int32), |
| np.zeros(0, dtype=np.int32), |
| self.config.batch_max_length, |
| ) |
| return {k: torch.from_numpy(v) for k, v in (empty | tensors).items()}, scalars |
|
|
| batch = {k: np.concatenate(v, dtype=np.int32) for k, v in batch.items()} |
| pad_len = self.config.batch_max_length - batch["inputs"].shape[0] |
| if pad_len > 0: |
| pad_values = { |
| "inputs": 0, |
| "labels": IGNORE_LABEL_ID, |
| "position_ids": 0, |
| } |
| for k, fill in pad_values.items(): |
| batch[k] = np.pad(batch[k], (0, pad_len), mode="constant", constant_values=fill) |
|
|
| seqlen_i = np.array(seqlens_prefix, dtype=np.int32) |
| seqlen_o = np.array(seqlens_causal, dtype=np.int32) |
| tensors, scalars = compute_aux_seq_tensors_scalars(seqlen_i, seqlen_o, self.config.batch_max_length) |
| return {k: torch.from_numpy(v) for k, v in (batch | tensors).items()}, scalars |
|
|
| def __iter__(self): |
| worker_info = get_worker_info() |
| assert worker_info is None or worker_info.num_workers == 1 |
|
|
| self._load_dataset_before_epoch_begin() |
| assert self._sampler is not None |
| start_index = int(self._resume_sampler_start) |
| self._resume_sampler_start = 0 |
| for indices in self._sampler.iter(start_index=start_index): |
| batch, scalars = self._load_batch(indices) |
| scalars = { |
| **scalars, |
| "_resume_data_epoch": int(self._active_data_epoch if self._active_data_epoch is not None else 0), |
| "_resume_sampler_start": int(getattr(self._sampler, "cursor", 0)), |
| "_resume_epoch_num_samples": int(getattr(self._sampler, "lengths").size), |
| } |
| yield batch, scalars |
|
|
