File size: 8,416 Bytes
d5147ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """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
# Optional marker written by 111fineweb / prepare_document_data
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 # autoregressive shift
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
# Each doc of length L contributes L-1 supervised positions.
lengths = np.maximum(self._data_indices.doc_len.astype(np.int64) - 1, 1)
# seed≠0: permute docs so resume-without-cursor does not replay the same prefix.
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"]}
# prefix_lens = 0 for every doc → full causal path in PrefixLM kernels
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"]:
# Degenerate batch: emit a fully-padded empty step (should be rare).
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
|