Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / megatron / gpt_dataset.py: 17%
143 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""``GPTDataset`` — fixed-length samples drawn from a Megatron IndexedDataset.
17Behaviourally aligned with Megatron-LM's ``GPTDataset``:
191. ``document_index`` — a shuffled flat array repeating every document for
20 ``num_epochs`` epochs. Documents stay whole; only their visit order is
21 randomised.
222. ``sample_index`` — for sample ``i`` the entry ``[doc_id, doc_offset]``
23 points to the document and within-document token offset where that
24 sample starts. Built by sequentially walking the document stream and
25 slicing every ``seq_length`` tokens.
263. ``shuffle_index`` — a random permutation applied on top, so sample
27 ``__getitem__(i)`` actually reads sample ``shuffle_index[i]``.
29A sample may span multiple documents (Megatron's concatenated-doc style);
30in that case we glue the relevant slices together until we reach
31``seq_length + 1`` tokens (one extra so labels can be shifted, exactly
32Megatron's ``add_extra_token_to_sequence=1``). If the very last sample
33runs out of documents, the tail is right-padded with ``pad_token_id``.
35Pure Python+numpy. No Cython helpers — the index construction runs once
36per ``(seed, prefix, num_samples, seq_length, num_epochs)`` tuple and is
37cached to disk under ``<prefix>_cache/`` so subsequent runs reuse it.
38"""
39import functools
40import hashlib
41import logging
42import os
43import tempfile
44from typing import Any, Dict, Optional, Tuple
46import numpy as np
47import torch
48from torch.utils.data import Dataset
50from hyper_parallel.data.megatron.indexed_dataset import IndexedDataset
53logger = logging.getLogger(__name__)
56def _build_document_index(num_documents: int, num_epochs: int, rng: np.random.RandomState) -> np.ndarray:
57 """Return a globally-shuffled, ``num_epochs``-repeated array of document indices.
59 Matches Megatron-LM's ``separate_last_epoch=False`` default: tile the
60 document range across ``num_epochs`` and run a single global shuffle.
61 This intentionally lets a document appear several times before another
62 appears once — that randomness is part of the canonical training
63 distribution and the realised epoch-boundary noise is what Megatron's
64 schedulers expect.
65 """
66 doc_idx = np.tile(np.arange(num_documents, dtype=np.int32), num_epochs)
67 rng.shuffle(doc_idx)
68 return doc_idx
71def _build_sample_index(
72 sizes: np.ndarray,
73 doc_idx: np.ndarray,
74 seq_length: int,
75 tokens_per_epoch: int,
76 num_epochs: int,
77) -> np.ndarray:
78 """Construct the ``[num_samples + 1, 2]`` (doc_id, doc_offset) sample table.
80 Walks the per-document token stream in ``doc_idx`` order, marking the
81 start of every ``seq_length``-window. Each sample includes one extra
82 token so labels can be shifted by 1; the helper keeps the
83 pre-Megatron-Cython slow-path semantics (no boundary correction —
84 multi-document samples are reassembled by :meth:`GPTDataset._read_sample`).
85 """
86 tokens_to_consume = num_epochs * tokens_per_epoch - 1 # -1 to keep ``+1`` extra token in reach
87 sample_stride = seq_length
88 num_samples = tokens_to_consume // sample_stride
89 # int64: column 0 indexes ``doc_idx`` whose length is
90 # ``num_documents * num_epochs`` — int32 would wrap past ~2.1e9 slots on
91 # large corpora trained for many epochs, silently selecting the wrong
92 # document.
93 sample_idx = np.zeros((num_samples + 1, 2), dtype=np.int64)
95 cur_doc_index = 0
96 cur_doc_offset = 0
97 cur_doc_remaining = int(sizes[doc_idx[cur_doc_index]])
98 sample_idx[0, 0] = cur_doc_index
99 sample_idx[0, 1] = cur_doc_offset
100 for sample_id in range(1, num_samples + 1):
101 remaining = sample_stride
102 while remaining > 0:
103 if cur_doc_remaining > remaining:
104 cur_doc_offset += remaining
105 cur_doc_remaining -= remaining
106 remaining = 0
107 else:
108 remaining -= cur_doc_remaining
109 cur_doc_index += 1
110 if cur_doc_index >= doc_idx.shape[0]:
111 # Document stream exhausted — record where we landed
112 # and stop. Caller will pad the tail.
113 cur_doc_remaining = 0
114 break
115 cur_doc_offset = 0
116 cur_doc_remaining = int(sizes[doc_idx[cur_doc_index]])
117 sample_idx[sample_id, 0] = cur_doc_index
118 sample_idx[sample_id, 1] = cur_doc_offset
119 return sample_idx
122def _build_shuffle_index(num_samples: int, rng: np.random.RandomState) -> np.ndarray:
123 """Random permutation of ``[0, num_samples)``."""
124 shuffle = np.arange(num_samples, dtype=np.int64)
125 rng.shuffle(shuffle)
126 return shuffle
129def _layout_fingerprint(doc_sizes: np.ndarray) -> str:
130 """Hash the per-document token-count array.
132 ``doc_sizes`` is exactly what determines every ``sample_idx`` offset, so a
133 corpus regenerated at the same path with a different document layout — even
134 one preserving document count and total token count — produces a different
135 fingerprint and invalidates the stale cache.
136 """
137 return hashlib.sha1(np.ascontiguousarray(doc_sizes, dtype=np.int64).tobytes()).hexdigest()[:16]
140def _cache_key(
141 prefix: str,
142 seq_length: int,
143 num_samples: int,
144 num_epochs: int,
145 seed: int,
146 layout_fp: str,
147) -> str:
148 """Stable filename suffix for the ``(doc, sample, shuffle)`` index trio.
150 ``layout_fp`` fingerprints the underlying corpus layout (see
151 :func:`_layout_fingerprint`) so that regenerating the ``.bin``/``.idx``
152 under the same path invalidates a stale cache instead of silently
153 replaying indices computed against the old token stream.
154 """
155 h = hashlib.sha1(
156 f"{prefix}|{seq_length}|{num_samples}|{num_epochs}|{seed}|{layout_fp}".encode("utf-8"),
157 ).hexdigest()[:16]
158 return h
161class GPTDataset(Dataset):
162 """Megatron-style GPT dataset on top of an :class:`IndexedDataset`.
164 Args:
165 indexed_dataset: Underlying token-stream reader.
166 num_samples: How many fixed-length samples to expose. Typically
167 ``max_steps * global_batch_size`` so the dataloader never
168 runs dry mid-training.
169 seq_length: Tokens per ``input_ids`` (excluding the extra label token).
170 seed: RNG seed for the document / shuffle permutations.
171 pad_token_id: Token used to right-pad the very last sample when
172 the document stream is exhausted before ``seq_length + 1``
173 tokens are available.
174 eod_mask_loss: When ``True``, the EOD token's label is replaced by
175 ``-100`` so it does not contribute to the loss.
176 eod_token_id: EOD token id; ignored when ``eod_mask_loss`` is False.
177 cache_dir: Where to cache the (doc, sample, shuffle) index arrays.
178 Defaults to ``<indexed_dataset.path_prefix>_cache``.
180 Returns from ``__getitem__``:
181 ``{"input_ids": Tensor[seq_length], "labels": Tensor[seq_length]}``
182 — labels is a copy of input_ids (the model handles next-token
183 shifting internally, matching the rest of the trainer's
184 dataset contract).
185 """
187 def __init__(
188 self,
189 indexed_dataset: IndexedDataset,
190 num_samples: int,
191 seq_length: int,
192 seed: int = 1234,
193 *,
194 pad_token_id: int = 0,
195 eod_mask_loss: bool = False,
196 eod_token_id: Optional[int] = None,
197 cache_dir: Optional[str] = None,
198 ) -> None:
199 if num_samples <= 0:
200 raise ValueError(f"num_samples must be > 0, got {num_samples}")
201 if seq_length <= 0:
202 raise ValueError(f"seq_length must be > 0, got {seq_length}")
204 self.indexed_dataset = indexed_dataset
205 self.num_samples = int(num_samples)
206 self.seq_length = int(seq_length)
207 self.seed = int(seed)
208 self.pad_token_id = int(pad_token_id)
209 self.eod_mask_loss = bool(eod_mask_loss)
210 self.eod_token_id = eod_token_id
212 self.document_index, self.sample_index, self.shuffle_index = self._build_indices(cache_dir)
213 # Adjacent samples often touch the same document — cache the concat
214 # result for the last few requested doc ids so straddling reads
215 # don't re-walk the IndexedDataset.
216 self._cached_doc = functools.lru_cache(maxsize=128)(self._read_doc)
218 def _build_indices(self, cache_dir: Optional[str]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
219 """Build (or load from cache) the three Megatron-style indices."""
220 sizes = self.indexed_dataset.sequence_lengths
221 num_docs = self.indexed_dataset.num_documents
222 if num_docs == 0:
223 raise ValueError("IndexedDataset has no documents to sample from")
224 # Build a per-document length array using the existing
225 # ``document_indices`` boundaries.
226 doc_boundaries = self.indexed_dataset.document_indices
227 doc_sizes = np.zeros(num_docs, dtype=np.int64)
228 for d in range(num_docs):
229 start = int(doc_boundaries[d])
230 end = int(doc_boundaries[d + 1])
231 doc_sizes[d] = int(sizes[start:end].sum())
233 tokens_per_epoch = int(doc_sizes.sum())
234 if tokens_per_epoch == 0:
235 raise ValueError("IndexedDataset has zero total tokens")
236 # ``+1`` because each sample reads ``seq_length + 1`` tokens (input
237 # plus the next-token label). Ceil so we always cover ``num_samples``.
238 tokens_required = self.num_samples * self.seq_length + 1
239 num_epochs = max(1, int(np.ceil(tokens_required / tokens_per_epoch)))
241 cache_dir = self._resolve_cache_dir(cache_dir)
242 key = _cache_key(
243 self.indexed_dataset.path_prefix, self.seq_length,
244 self.num_samples, num_epochs, self.seed,
245 _layout_fingerprint(doc_sizes),
246 )
247 doc_path = os.path.join(cache_dir, f"doc_idx_{key}.npy")
248 sample_path = os.path.join(cache_dir, f"sample_idx_{key}.npy")
249 shuffle_path = os.path.join(cache_dir, f"shuffle_idx_{key}.npy")
251 if os.path.isfile(doc_path) and os.path.isfile(sample_path) and os.path.isfile(shuffle_path):
252 return np.load(doc_path), np.load(sample_path), np.load(shuffle_path)
254 rng = np.random.RandomState(self.seed)
255 doc_idx = _build_document_index(num_docs, num_epochs, rng)
256 sample_idx = _build_sample_index(
257 doc_sizes, doc_idx, self.seq_length, tokens_per_epoch, num_epochs,
258 )
259 # ``sample_idx`` has at most ``num_samples + 1`` rows; clip in case
260 # we built slightly more (rounded up by tokens_per_epoch / seq_len).
261 usable_samples = max(sample_idx.shape[0] - 1, 1)
262 shuffle_idx = _build_shuffle_index(min(self.num_samples, usable_samples), rng)
263 np.save(doc_path, doc_idx)
264 np.save(sample_path, sample_idx)
265 np.save(shuffle_path, shuffle_idx)
266 return doc_idx, sample_idx, shuffle_idx
268 def _resolve_cache_dir(self, user_cache_dir: Optional[str]) -> str:
269 """Return a writable cache directory, falling back to ``$TMPDIR``.
271 Defaulting to ``<prefix>_cache`` next to the corpus is convenient,
272 but production datasets often sit on read-only shared mounts where
273 ``os.makedirs`` would raise ``EROFS``. Detect and silently relocate.
274 """
275 candidate = user_cache_dir or (self.indexed_dataset.path_prefix + "_cache")
276 try:
277 os.makedirs(candidate, exist_ok=True)
278 # Probe writability — ``os.makedirs`` succeeds on directories
279 # that exist read-only.
280 probe = os.path.join(candidate, ".hp_cache_probe")
281 with open(probe, "w", encoding="utf-8") as f:
282 f.write("")
283 os.remove(probe)
284 return candidate
285 except OSError as exc:
286 fallback = os.path.join(tempfile.gettempdir(), "hp_megatron_cache")
287 os.makedirs(fallback, exist_ok=True)
288 logger.warning(
289 "GPTDataset cache dir %s not writable (%s); using fallback %s",
290 candidate, exc, fallback,
291 )
292 return fallback
294 def __len__(self) -> int:
295 return min(self.num_samples, int(self.shuffle_index.size))
297 def _read_sample(self, sample_pos: int) -> np.ndarray:
298 """Concatenate the ``seq_length + 1`` tokens for sample ``sample_pos``."""
299 start_doc_idx, start_offset = self.sample_index[sample_pos]
300 end_doc_idx, end_offset = self.sample_index[sample_pos + 1]
301 parts = []
302 if start_doc_idx == end_doc_idx:
303 doc_tokens = self._cached_doc(int(self.document_index[start_doc_idx]))
304 parts.append(doc_tokens[start_offset:end_offset + 1])
305 else:
306 doc_tokens = self._cached_doc(int(self.document_index[start_doc_idx]))
307 parts.append(doc_tokens[start_offset:])
308 for doc_id in range(start_doc_idx + 1, end_doc_idx):
309 parts.append(self._cached_doc(int(self.document_index[doc_id])))
310 tail_doc_tokens = self._cached_doc(int(self.document_index[end_doc_idx]))
311 parts.append(tail_doc_tokens[:end_offset + 1])
312 out = np.concatenate(parts) if len(parts) > 1 else parts[0]
313 needed = self.seq_length + 1
314 if out.size >= needed:
315 return out[:needed]
316 # Pad short tail (only happens for the very last sample).
317 pad = np.full(needed - out.size, self.pad_token_id, dtype=out.dtype)
318 return np.concatenate([out, pad])
320 def _read_doc(self, doc_id: int) -> np.ndarray:
321 """All tokens of ``doc_id`` as a single 1-D array (memoised)."""
322 return self.indexed_dataset.get_document(doc_id)
324 def __getitem__(self, idx: int) -> Dict[str, Any]:
325 actual = int(self.shuffle_index[idx])
326 tokens = self._read_sample(actual)
327 input_ids = torch.from_numpy(tokens[: self.seq_length].astype(np.int64))
328 labels = input_ids.clone()
329 if self.eod_mask_loss and self.eod_token_id is not None:
330 labels[input_ids == int(self.eod_token_id)] = -100
331 return {"input_ids": input_ids, "labels": labels}