train: add shard_dataset (Phase 3/4 glue)

train/shard_dataset.py

@@ -0,0 +1,159 @@
+# Map-style reader over the built mix, designed so that resumption is an index advance, not a replay.
+
+# Why this file exists in this shape (docs/01-plan.md §5, "the resume trap"): `transformers.Trainer` can
+# resume an `IterableDataset` only by *decoding and discarding* the batches before the checkpoint step,
+# which costs O(steps) of I/O after every interruption -- up to 3,815 steps of nothing. A map-style
+# reader over a concatenated token store makes the same thing a one-line offset. §3.1's "same position"
+# is therefore a number in the checkpoint, not a feeling.
+#
+# Sampling policy, stated so it can be reproduced:
+#   stream   = the uint16 id arrays of manifest["shards"] in manifest order, concatenated
+#   sample i = stream[i*L : (i+1)*L]  with L = seq_len, no padding, documents packed back to back
+#   order    = a single global permutation of sample indices from numpy's PCG64(seed)
+#
+# Packing without a document boundary means a few windows per shard splice two unrelated documents and
+# let attention cross the seam. That is what TinyLlama, SmolLM and GPT-2-era pretraining all do; a
+# block-diagonal mask would cost activation memory we do not have at bs 2x2048 on a T4. Recorded as a
+# known, deliberate deviation rather than a silence.
+#
+# The permutation is over ~490k indices (1B tokens / 2048), i.e. ~2 MB, so a *global* shuffle is free
+# here -- no need for the windowed approximations that streaming readers settle for.
+#
+# torch is imported lazily-ish: the byte-level part of this module is usable in a CPU-only check.
+
+import bisect
+import json
+import os
+import struct
+
+import numpy as np
+
+
+def read_shard_index(path):
+    """(vocab, n_docs, n_tokens, offset_bytes, ids_bytes) for one shard without loading its ids."""
+    with open(path, "rb") as f:
+        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", f.read(16))
+    if flags != 0:
+        raise ValueError(f"{path}: unknown shard flags {flags}")
+    return vocab, n_docs, n_tokens, 16 + 4 * (n_docs + 1), 2 * n_tokens
+
+
+class PackedTokenStore:
+    """Random-access uint16 view of the concatenated mix, in manifest order."""
+
+    def __init__(self, root, manifest=None, files=None, vocab=None):
+        man = manifest if manifest is not None else json.load(
+            open(os.path.join(root, "manifest.json")))
+        self.vocab = vocab or man["tokenizer"]["vocab_size"]
+        names = files if files is not None else [s["file"] for s in man["shards"]]
+        self.shards, self.starts, self._maps = [], [], []
+        tot = 0
+        for name in names:
+            p = name if os.path.isabs(name) else os.path.join(root, name)
+            v, n_docs, n_tokens, ob, ib = read_shard_index(p)
+            if v != self.vocab:
+                raise ValueError(f"{p}: shard vocab {v} != manifest vocab {self.vocab}")
+            # memmap, never load: the whole mix is ~2.6 GB and a GPU session has 30 GiB shared with
+            # two CUDA contexts. Page faults are the point.
+            mm = np.memmap(p, dtype=np.uint16, mode="r", offset=ob, shape=(n_tokens,))
+            self.shards.append((p, mm, n_tokens))
+            self.starts.append(tot)
+            tot += n_tokens
+        self.total_tokens = tot
+        if tot == 0:
+            # Zero-length store must not look like a valid empty dataset (E-017's shape).
+            raise ValueError("no tokens in the mix -- nothing was built or nothing was restored")
+
+    def tokens(self, g0, n):
+        """Global token span [g0, g0+n) as a contiguous uint16 array, across shard boundaries."""
+        out = np.empty(n, dtype=np.uint16)
+        done = 0
+        while done < n:
+            si = bisect.bisect_right(self.starts, g0 + done) - 1
+            p, mm, nt = self.shards[si]
+            local = (g0 + done) - self.starts[si]
+            take = min(n - done, nt - local)
+            if take <= 0:
+                raise IndexError(f"span [{g0},{g0+n}) runs past the end of the mix ({self.total_tokens})")
+            out[done:done + take] = mm[local:local + take]
+            done += take
+        return out
+
+
+def make_dataset(store, seq_len, seed, start_sample=0, count=None, shuffle=True):
+    """A map-style torch Dataset of (input_ids, labels, attention_mask) with an exact resume cursor.
+
+    start_sample is the number of samples ALREADY consumed: the permutation is generated in full and
+    then sliced, so the sample at index j of the resumed run is the same global sample it would have
+    been in an uninterrupted run. That is the property Gate 3 has to demonstrate.
+    """
+    import torch
+
+    n_samples = store.total_tokens // seq_len
+    if drop_last is False:
+        raise ValueError("partial trailing windows are not supported; drop_last must stay True")
+    perm = None
+    if shuffle:
+        # Seed must be identical on every resume or the two runs read different data orders. Stored in
+        # the checkpoint (cursor.json) and asserted by the trainer.
+        perm = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
+        perm = perm[start_sample:]
+    else:
+        perm = np.arange(start_sample, n_samples, dtype=np.int64)
+    if count is not None:
+        perm = perm[:count]
+
+    class _DS(torch.utils.data.Dataset):
+        def __len__(self):
+            return int(len(perm))
+
+        def __getitem__(self, j):
+            s = int(perm[j])
+            ids = store.tokens(s * seq_len, seq_len).astype(np.int64)
+            x = torch.from_numpy(ids)
+            return {"input_ids": x,
+                    "labels": torch.cat((x[1:], x[:1])),
+                    "attention_mask": torch.ones_like(x)}
+
+        # Trainer calls remove_unused_columns on models whose forward signature it inspects; exposing
+        # these names keeps it from dropping anything.
+        column_names = ["input_ids", "labels", "attention_mask"]
+
+    _DS.perm = perm
+    return _DS()
+
+
+def cursor_dict(seq_len, start_sample, seed, dataset_files_sha=None):
+    """The whole of the data-position state. Written into every checkpoint, checked on every resume."""
+    return {"seq_len": seq_len, "samples_consumed": start_sample,
+            "tokens_consumed": start_sample * seq_len, "shuffle_seed": seed,
+            "dataset_files_sha": dataset_files_sha}
+
+
+def load_cursor(dirpath):
+    p = os.path.join(dirpath, "cursor.json")
+    if not os.path.exists(p):
+        return None
+    with open(p) as f:
+        c = json.load(f)
+    if not isinstance(c.get("samples_consumed"), int) or c["samples_consumed"] < 0:
+        raise ValueError(f"{p}: samples_consumed missing or nonsense -- refusing a silent restart")
+    return c
+
+
+def save_cursor(dirpath, cur):
+    os.makedirs(dirpath, exist_ok=True)
+    with open(os.path.join(dirpath, "cursor.json"), "w") as f:
+        json.dump(cur, f, indent=1, sort_keys=True)
+
+
+def mix_fingerprint(root, manifest=None):
+    """Cheap, stable identity for the exact data being trained on: sha of the manifest's shard list
+    plus each shard's size. Used by the trainer to refuse a resume onto a different corpus."""
+    import hashlib
+    man = manifest or json.load(open(os.path.join(root, "manifest.json")))
+    h = hashlib.sha256()
+    for s in man["shards"]:
+        p = os.path.join(root, s["file"])
+        h.update(f"{s['file']}:{os.path.getsize(p) if os.path.exists(p) else 'missing'}\n".encode())
+    return h.hexdigest()[:16]