MegatronDataset: Megatron Core indexed dataset (.bin/.idx) token IDs

[albertz]

New dataset MegatronDataset for the Megatron Core indexed dataset format (MMIDIDX version 1),
as written by Megatron's IndexedDatasetBuilder
(reference).
A .bin file holds the token IDs, its .idx file the seq lengths, byte pointers and document boundaries.
This format is used by many LM corpora (e.g. Nemotron), not specific to one corpus.

train = {
    "class": "MegatronDataset",
    "data_file": ["/data/tokens-0.bin", "/data/tokens-1.bin"],
    "index_file": ["/data/tokens-0.idx", "/data/tokens-1.idx"],
    "vocab_size": 131072,
    "seq_ordering": "random",
    "partition_epoch": 10,
}

• NumPy and the standard library only: no megatron-core, PyTorch, CUDA or tokenizer needed to read the IDs.
• data_file and index_file: a file or a list of files, one index file per data file.
  Separate options, so another index (e.g. a resegmentation) can refer to an unchanged data file.
  A list makes a sub epoch of DistributeFilesDataset with several shards one MegatronDataset.
• One seq per index entry, in file order. Data key data: token IDs as stored (incl. EOD and special IDs), sparse, int32.
• Seq tags <name>/doc-<d>/seq-<j>: entry j of document d of the index file, <name> its base name without .idx.
  So a tag does not depend on the other files or their order (e.g. per sub epoch of DistributeFilesDataset).
  Unique (document indices are checked, index file base names must be unique), and the document is directly in the tag.
  The format has no other identifiers (no document IDs or names, Megatron's preprocessing drops them).
• vocab_size is required (the index does not store it, and the payload is not scanned for it).
  Optional vocab must have vocab_size labels. Token IDs outside [0, vocab_size) raise when the seq is loaded.
• Integer dtype codes 1, 2, 3, 4, 5, 8 (uint8, int8, int16, int32, int64, uint16), converted to int32, per file.
  Float codes, unknown codes, wrong magic or version, truncated or multimodal index files,
  document indices not from 0 to num seqs or not monotonic, negative lengths and spans outside the data file raise ValueError.
  The checks are vectorized and chunked, once on the first access.
• The files are memory-mapped on the first access, so construction and pickling are cheap, and workers reopen them.
• Seq ordering, partition_epoch and sharding (_num_shards / set_sharding) via the generic get_seq_order_for_epoch
  (a NumPy permutation over the seqs of this dataset, no Python lists).
• No packing, sampling or blending of Megatron's GPTDataset.

Tests (tests/test_Dataset.py, CPU, py3.12, torch 2.12 env at RWTH):

• test_MegatronDataset: uint16 (with ID 65535) and int32, empty document, zero-length seq, pickling before and after the first access,
  tags and seq_list (also invalid tags), the same seqs under the same tags over different file lists and orders,
  duplicate index base names, a separately written split index over the unchanged data file
  (sorted ordering, documents reassemble losslessly).
• test_MegatronDataset_invalid: all the error cases above, the vocab size mismatch, and mismatching file list lengths.
• test_MegatronDataset_partition_epoch_sharding: 3 sub-epochs x 2 shards cover every seq exactly once, deterministic, new order next epoch.
• test_MegatronDataset_DistributeFilesDataset: 6 shard files over 3 sub-epochs, 2 files per sub epoch dataset.
• All 4 pass on head (1.1 s together), all 4 fail on master (Dataset class 'MegatronDataset' not found).
  Full test_Dataset.py: 49 passed, 9 skipped.
• Probe outside the tests: files written by megatron-core's own IndexedDatasetBuilder (add_document, add_item, end_document, empty document, zero-length seq),
  uint16 and int32: IDs, lengths, pointers and document indices equal megatron-core's IndexedDataset, tags match its document indices.