Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / megatron / indexed_dataset.py: 27%
133 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"""Megatron-LM ``.bin`` / ``.idx`` reader and writer.
17Wire format (matches Megatron-LM ``mmap_indexed_dataset.py``):
19``<prefix>.idx``:
20- 9-byte magic ``b"MMIDIDX\\x00\\x00"``
21- ``uint64`` version (=1)
22- ``uint8`` dtype code (1=uint8, 2=int8, 3=int16, 4=int32, 5=int64,
23 6=float64, 7=float32, 8=uint16)
24- ``uint64`` sequence count ``N``
25- ``uint64`` document count ``D``
26- ``int32 [N]`` sequence lengths (in tokens)
27- ``int64 [N]`` sequence pointers (byte offset into ``.bin``)
28- ``int64 [D + 1]`` document indices (sequence-index boundaries; first is 0,
29 last is ``N``)
30- Optional ``int8 [N]`` sequence modes for multimodal datasets — ignored
31 here (LM-only path)
33``<prefix>.bin``: raw concatenation of all sequences in the declared dtype.
35We expose:
37- :func:`get_bin_path` / :func:`get_idx_path` — canonical suffixing logic.
38- :class:`IndexedDataset` — memory-mapped ``.bin`` payload with in-memory
39 ``.idx`` metadata; ``__getitem__`` returns ``np.ndarray``.
40- :class:`IndexedDatasetBuilder` — sequential writer for tests, fixtures,
41 and ad-hoc preprocessing.
43Pure ``numpy`` — no Cython helpers — so this works wherever ``numpy`` runs.
44The data path stays zero-copy: token slices view the memory-mapped ``.bin``;
45the smaller ``.idx`` arrays are read into RAM once at open time.
46"""
47import logging
48import os
49import struct
50from typing import Dict, List, Optional, Tuple
52import numpy as np
55logger = logging.getLogger(__name__)
58# Megatron magic header — kept byte-identical so a ``.idx`` produced by
59# Megatron-LM ``preprocess_data.py`` loads transparently here.
60_MAGIC = b"MMIDIDX\x00\x00"
61_VERSION = 1
63_DTYPES: Dict[int, np.dtype] = {
64 1: np.dtype(np.uint8),
65 2: np.dtype(np.int8),
66 3: np.dtype(np.int16),
67 4: np.dtype(np.int32),
68 5: np.dtype(np.int64),
69 6: np.dtype(np.float64),
70 7: np.dtype(np.float32),
71 8: np.dtype(np.uint16),
72}
73_DTYPE_CODES: Dict[np.dtype, int] = {v: k for k, v in _DTYPES.items()}
76def strip_suffix(prefix: str) -> str:
77 """Return ``prefix`` with a trailing ``.bin``/``.idx`` removed (idempotent).
79 The single source of truth for path-prefix normalisation — :func:`get_bin_path`,
80 :func:`get_idx_path` and the dataset builder all route through here so the
81 suffix rule never diverges across call sites.
82 """
83 return prefix[:-4] if prefix.endswith((".bin", ".idx")) else prefix
86def get_bin_path(prefix: str) -> str:
87 """Return ``<prefix>.bin`` (strips an already-present ``.bin``/``.idx``)."""
88 return strip_suffix(prefix) + ".bin"
91def get_idx_path(prefix: str) -> str:
92 """Return ``<prefix>.idx``."""
93 return strip_suffix(prefix) + ".idx"
96def dtype_code(dtype: np.dtype) -> int:
97 """Return the on-disk dtype code for ``dtype``.
99 Raises:
100 ValueError: When ``dtype`` is not part of the Megatron dtype table.
101 """
102 dtype = np.dtype(dtype)
103 if dtype not in _DTYPE_CODES:
104 raise ValueError(
105 f"Unsupported Megatron indexed-dataset dtype: {dtype}. "
106 f"Allowed: {sorted(d.name for d in _DTYPE_CODES)}"
107 )
108 return _DTYPE_CODES[dtype]
111def code_dtype(code: int) -> np.dtype:
112 """Inverse of :func:`dtype_code`."""
113 try:
114 return _DTYPES[code]
115 except KeyError as exc:
116 raise ValueError(f"Unknown Megatron dtype code: {code}") from exc
119class _IndexReader:
120 """In-memory metadata parsed from the ``.idx`` companion file.
122 Holds the four arrays Megatron-LM stores: lengths, pointers,
123 document indices and optional modes. The index payload is read into
124 compact numpy arrays; the much larger ``.bin`` remains memory-mapped by
125 :class:`IndexedDataset`.
126 """
128 def __init__(self, idx_path: str, *, multimodal: bool = False) -> None:
129 if not os.path.isfile(idx_path):
130 raise FileNotFoundError(f"Megatron index file not found: {idx_path}")
131 with open(idx_path, "rb") as f:
132 magic = f.read(len(_MAGIC))
133 if magic != _MAGIC:
134 raise ValueError(
135 f"Not a Megatron indexed dataset (magic mismatch): "
136 f"{idx_path}. Expected {_MAGIC!r}, got {magic!r}"
137 )
138 version = struct.unpack("<Q", f.read(8))[0]
139 if version != _VERSION:
140 raise ValueError(
141 f"Unsupported indexed-dataset version: {version} "
142 f"(only v{_VERSION} is recognised)"
143 )
144 code = struct.unpack("<B", f.read(1))[0]
145 self.dtype = code_dtype(code)
146 self.sequence_count = struct.unpack("<Q", f.read(8))[0]
147 self.document_count = struct.unpack("<Q", f.read(8))[0]
149 # The index payload is small (4-12 bytes per sequence) — read it
150 # into RAM rather than memmap. Avoids the fragility of slicing a
151 # memmap with offsetted ``np.frombuffer`` while still being fast.
152 self.sequence_lengths = np.frombuffer(
153 f.read(self.sequence_count * 4), dtype=np.int32,
154 )
155 self.sequence_pointers = np.frombuffer(
156 f.read(self.sequence_count * 8), dtype=np.int64,
157 )
158 self.document_indices = np.frombuffer(
159 f.read(self.document_count * 8), dtype=np.int64,
160 )
161 if multimodal:
162 self.sequence_modes = np.frombuffer(
163 f.read(self.sequence_count), dtype=np.int8,
164 )
165 else:
166 self.sequence_modes = None
169class IndexedDataset:
170 """Read-only Megatron ``.bin``/``.idx`` dataset.
172 Args:
173 path_prefix: Path without the ``.bin``/``.idx`` suffix. Either
174 suffix is also accepted and stripped.
175 mmap: When ``True`` (default), the ``.bin`` is memory-mapped; reads
176 return zero-copy views. Set to ``False`` to load the entire
177 ``.bin`` into RAM — useful only for sub-MB fixtures in tests.
178 multimodal: Whether to parse the optional ``sequence_modes`` array
179 in the ``.idx`` (LM training leaves this off).
181 Example:
182 >>> ds = IndexedDataset("/data/wiki") # noqa: SKIP
183 >>> tokens = ds[0] # noqa: SKIP
184 >>> doc_0 = ds.get_document(0) # noqa: SKIP
186 Note:
187 Equivalent to Megatron-LM's ``MMapIndexedDataset`` but rewritten
188 for the HyperParallel coding conventions; the on-disk format is
189 byte-compatible.
190 """
192 def __init__(self, path_prefix: str, *, mmap: bool = True, multimodal: bool = False) -> None:
193 idx_path = get_idx_path(path_prefix)
194 bin_path = get_bin_path(path_prefix)
195 if not os.path.isfile(bin_path):
196 raise FileNotFoundError(f"Megatron .bin not found: {bin_path}")
197 self._index = _IndexReader(idx_path, multimodal=multimodal)
198 if mmap:
199 self._bin_buffer = np.memmap(bin_path, dtype=np.uint8, mode="r", order="C")
200 else:
201 with open(bin_path, "rb") as f:
202 data = f.read()
203 self._bin_buffer = np.frombuffer(data, dtype=np.uint8)
204 self.path_prefix = path_prefix
206 # ------------------------------------------------------------------
207 # Sequence-level access
208 # ------------------------------------------------------------------
210 def __len__(self) -> int:
211 return int(self._index.sequence_count)
213 @property
214 def dtype(self) -> np.dtype:
215 """Token dtype (numpy)."""
216 return self._index.dtype
218 @property
219 def sequence_lengths(self) -> np.ndarray:
220 """int32[N] token-counts for each sequence."""
221 return self._index.sequence_lengths
223 @property
224 def document_indices(self) -> np.ndarray:
225 """int64[D+1] sequence-index boundaries marking documents."""
226 return self._index.document_indices
228 def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> np.ndarray:
229 """Read a slice of sequence ``idx``.
231 Args:
232 idx: Sequence index in ``[0, len)``.
233 offset: Token offset within the sequence.
234 length: Tokens to read; ``None`` reads to the end.
236 Returns:
237 A ``np.ndarray`` of the configured dtype. The returned array is
238 a view into the mmap when ``mmap=True``; callers must not write
239 to it.
240 """
241 seq_len = int(self._index.sequence_lengths[idx])
242 if length is None:
243 length = seq_len - offset
244 if offset < 0 or length < 0 or offset + length > seq_len:
245 raise ValueError(
246 f"Slice [{offset}:{offset + length}] out of bounds for "
247 f"sequence {idx} of length {seq_len}"
248 )
249 item_size = self._index.dtype.itemsize
250 byte_ptr = int(self._index.sequence_pointers[idx]) + offset * item_size
251 # ``sequence_pointers`` are accumulated ``nbytes`` of items, so
252 # ``byte_ptr`` is always a multiple of ``item_size`` and a typed
253 # view is alignment-safe; the slice is a zero-copy view into the
254 # mmap (or a plain ndarray slice in the RAM path).
255 return self._bin_buffer[byte_ptr: byte_ptr + length * item_size].view(self._index.dtype)
257 def __getitem__(self, idx: int) -> np.ndarray:
258 return self.get(idx)
260 # ------------------------------------------------------------------
261 # Document-level access
262 # ------------------------------------------------------------------
264 @property
265 def num_documents(self) -> int:
266 """Document count (``len(document_indices) - 1``).
268 On disk ``document_count`` is ``D + 1`` because Megatron stores the
269 full boundary array including the trailing ``N`` terminator.
270 """
271 return int(self._index.document_count) - 1
273 def get_document(self, doc_idx: int) -> np.ndarray:
274 """Concatenate all sequences belonging to ``doc_idx``."""
275 if doc_idx < 0 or doc_idx >= self.num_documents:
276 raise IndexError(
277 f"document index {doc_idx} out of range [0, {self.num_documents})"
278 )
279 start = int(self._index.document_indices[doc_idx])
280 end = int(self._index.document_indices[doc_idx + 1])
281 parts = [self.get(i) for i in range(start, end)]
282 if not parts:
283 return np.array([], dtype=self._index.dtype)
284 return np.concatenate(parts)
287class IndexedDatasetBuilder:
288 """Write a ``.bin``/``.idx`` pair sequentially.
290 Usage:
291 builder = IndexedDatasetBuilder("/path/to/wiki", dtype=np.int32)
292 builder.add_item(np.array([1, 2, 3, 4], dtype=np.int32))
293 builder.end_document()
294 builder.add_item(np.array([5, 6, 7], dtype=np.int32))
295 builder.end_document()
296 builder.finalize()
298 Args:
299 path_prefix: Output path without suffix.
300 dtype: Token dtype; must be in the Megatron dtype table.
302 Note:
303 Designed for fixture generation and small-scale preprocessing.
304 Streams tokens straight to the ``.bin`` and keeps lightweight
305 Python lists for the index metadata until :meth:`finalize`.
306 """
308 def __init__(self, path_prefix: str, *, dtype: np.dtype = np.int32) -> None:
309 self.path_prefix = path_prefix
310 self.dtype = np.dtype(dtype)
311 # Make sure the dtype is supported before opening any file.
312 dtype_code(self.dtype)
313 self._bin_path = get_bin_path(path_prefix)
314 self._idx_path = get_idx_path(path_prefix)
315 os.makedirs(os.path.dirname(self._bin_path) or ".", exist_ok=True)
316 self._bin_file = open(self._bin_path, "wb") # pylint: disable=R1732
317 self._sequence_lengths: List[int] = []
318 self._sequence_pointers: List[int] = []
319 self._document_indices: List[int] = [0]
320 self._byte_offset = 0
322 def add_item(self, tokens: np.ndarray) -> None:
323 """Append one sequence of tokens to the dataset."""
324 tokens = np.asarray(tokens, dtype=self.dtype)
325 self._bin_file.write(tokens.tobytes(order="C"))
326 self._sequence_lengths.append(int(tokens.size))
327 self._sequence_pointers.append(self._byte_offset)
328 self._byte_offset += int(tokens.nbytes)
330 def end_document(self) -> None:
331 """Mark the current sequence boundary as a document boundary."""
332 self._document_indices.append(len(self._sequence_lengths))
334 def finalize(self) -> Tuple[str, str]:
335 """Close ``.bin`` and write the companion ``.idx``.
337 Returns:
338 ``(idx_path, bin_path)`` for caller convenience.
340 Raises:
341 ValueError: When no document boundary was emitted.
342 """
343 self._bin_file.close()
344 # ``end_document`` must have been called at least once OR the writer
345 # is empty; an empty dataset still needs a 1-element boundary list.
346 if self._document_indices[-1] != len(self._sequence_lengths):
347 self.end_document()
348 seq_lengths = np.asarray(self._sequence_lengths, dtype=np.int32)
349 seq_pointers = np.asarray(self._sequence_pointers, dtype=np.int64)
350 doc_indices = np.asarray(self._document_indices, dtype=np.int64)
351 with open(self._idx_path, "wb") as f:
352 f.write(_MAGIC)
353 f.write(struct.pack("<Q", _VERSION))
354 f.write(struct.pack("<B", dtype_code(self.dtype)))
355 f.write(struct.pack("<Q", seq_lengths.size))
356 f.write(struct.pack("<Q", doc_indices.size))
357 f.write(seq_lengths.tobytes())
358 f.write(seq_pointers.tobytes())
359 f.write(doc_indices.tobytes())
360 logger.info(
361 "IndexedDatasetBuilder finalized: %d sequences, %d documents -> %s + %s",
362 seq_lengths.size, doc_indices.size, self._bin_path, self._idx_path,
363 )
364 return self._idx_path, self._bin_path