ounce100m-code / train /shard_dataset.py

perm_sha: identify the visit order by hashing the permutation, not just its seed (E-035 residual) -- cursor gains shuffle_perm_sha, asserted on resume; session repins the two files

398d16e verified about 1 month ago

Raw History Blame Contribute Delete

9.61 kB

# 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 perm_sha(n_samples, seed, base=None):
    """Identity of the visit order itself, not just the seed that produced it.
    `seed` pins the order only for a fixed numpy. Kaggle updates its notebook image underneath us, and a
    session that resumed after such a bump would read a different permutation from the same seed -- some
    windows twice, others never -- while every other guard (weights, cursor arithmetic, fingerprint,
    batch size) still agreed. Hashing the full permutation is ~10 ms and makes that state impossible to
    enter quietly. Computed before the resume slice, so two sessions at different offsets agree on it.
    `base` lets a caller that already generated the permutation avoid generating it twice.
    """
    import hashlib
    if base is None:
        base = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
    return hashlib.sha256(np.ascontiguousarray(base).tobytes()).hexdigest()[:16]


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
    perm = None
    order_sha = "unshuffled"
    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 -- alongside perm_sha(), which catches
        # the case the seed alone cannot: the same seed on a different numpy.
        base = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
        order_sha = perm_sha(n_samples, seed, base=base)
        perm = base[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 = _DS()
    # Instance attributes, not class attributes: `_DS.perm = perm` would let a second make_dataset() call
    # in the same process overwrite the first one's order.
    ds.perm = perm
    ds.order_sha = order_sha
    return ds


def cursor_dict(seq_len, samples_consumed, seed, dataset_files_sha=None, step=None, order_sha=None):
    """The whole of the data-position state. Written into every checkpoint, checked on every resume.
    `step` is recorded alongside `samples_consumed` so the two can be cross-checked on resume: they are
    redundant by construction (samples = step x sequences-per-step), and a cursor where they disagree is a
    corrupted or hand-edited file, not a position to continue from. `shuffle_perm_sha` is the identity of
    the visit order, which the seed alone does not pin (see perm_sha).
    """
    return {"seq_len": seq_len, "samples_consumed": samples_consumed,
            "tokens_consumed": samples_consumed * seq_len, "shuffle_seed": seed,
            "shuffle_perm_sha": order_sha, "step": step, "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. Used by the trainer to refuse a resume
    onto a different corpus, so it has to be computed the same way in every session.
    Names alone (or names and sizes, as this function did until the E-031 review) are too weak: a rebuilt
    mix whose shards happen to land on the same sizes is a different corpus, and the whole point of the
    check is to catch a corpus swap between sessions. The manifest already certifies a sha256 per shard,
    so folding it in costs nothing and needs no file reads; the local size is included as well so a
    truncated file on disk cannot hide behind a manifest that still describes the good one.
    """
    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']}:{s.get('sha256')}:"
                 f"{os.path.getsize(p) if os.path.exists(p) else 'missing'}\n".encode())
    return h.hexdigest()[:16]

	# 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[iL : (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 perm_sha(n_samples, seed, base=None):
	"""Identity of the visit order itself, not just the seed that produced it.

	`seed` pins the order only for a fixed numpy. Kaggle updates its notebook image underneath us, and a
	session that resumed after such a bump would read a different permutation from the same seed -- some
	windows twice, others never -- while every other guard (weights, cursor arithmetic, fingerprint,
	batch size) still agreed. Hashing the full permutation is ~10 ms and makes that state impossible to
	enter quietly. Computed before the resume slice, so two sessions at different offsets agree on it.
	`base` lets a caller that already generated the permutation avoid generating it twice.
	"""
	import hashlib
	if base is None:
	base = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
	return hashlib.sha256(np.ascontiguousarray(base).tobytes()).hexdigest()[:16]


	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
	perm = None
	order_sha = "unshuffled"
	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 -- alongside perm_sha(), which catches
	# the case the seed alone cannot: the same seed on a different numpy.
	base = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
	order_sha = perm_sha(n_samples, seed, base=base)
	perm = base[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 = _DS()
	# Instance attributes, not class attributes: `_DS.perm = perm` would let a second make_dataset() call
	# in the same process overwrite the first one's order.
	ds.perm = perm
	ds.order_sha = order_sha
	return ds


	def cursor_dict(seq_len, samples_consumed, seed, dataset_files_sha=None, step=None, order_sha=None):
	"""The whole of the data-position state. Written into every checkpoint, checked on every resume.

	`step` is recorded alongside `samples_consumed` so the two can be cross-checked on resume: they are
	redundant by construction (samples = step x sequences-per-step), and a cursor where they disagree is a
	corrupted or hand-edited file, not a position to continue from. `shuffle_perm_sha` is the identity of
	the visit order, which the seed alone does not pin (see perm_sha).
	"""
	return {"seq_len": seq_len, "samples_consumed": samples_consumed,
	"tokens_consumed": samples_consumed * seq_len, "shuffle_seed": seed,
	"shuffle_perm_sha": order_sha, "step": step, "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. Used by the trainer to refuse a resume
	onto a different corpus, so it has to be computed the same way in every session.

	Names alone (or names and sizes, as this function did until the E-031 review) are too weak: a rebuilt
	mix whose shards happen to land on the same sizes is a different corpus, and the whole point of the
	check is to catch a corpus swap between sessions. The manifest already certifies a sha256 per shard,
	so folding it in costs nothing and needs no file reads; the local size is included as well so a
	truncated file on disk cannot hide behind a manifest that still describes the good one.
	"""
	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']}:{s.get('sha256')}:"
	f"{os.path.getsize(p) if os.path.exists(p) else 'missing'}\n".encode())
	return h.hexdigest()[:16]