mixle.data.streaming_corpus module

Streaming, sharded, packed, deterministic token-corpus pipeline (roadmap F3).

Scope note: this builds the pipeline – sharded corpus streaming with per-rank residency, sequence packing, and deterministic global ordering given (seed, epoch) – over a pre-tokenized corpus (plain integer token-id arrays, one array per document). A real BPE/tokenizer library is not a dependency of this codebase and is out of scope here; tokenization is a separate, later concern. F1 (the distributed trainer this pipeline would ultimately feed “at target tokens/s”) does not exist yet either, so “saturates F1” is not measurable from this module alone – everything else in the F3 acceptance list (per-rank sharding, packing efficiency, determinism, curriculum hooks) is real and tested here.

Composes with existing machinery rather than duplicating it:

  • Per-rank residency extends MPEncodedData’s existing sharding contract: that handle splits an input sequence round-robin, rank i keeping data[i], data[i + world_size], data[i + 2*world_size], ... (disjoint, complete-coverage). shard_documents_for_rank() applies the identical round-robin split, but to the shuffled global document order from global_document_order() rather than to the raw input order – so shuffling changes which rank sees which document while the sharding contract itself (disjoint, complete, index mod world_size) is unchanged and remains bitwise reproducible.

  • The packed output rows are block + 1 tokens: row[:-1] is the model input and row[1:] is the shifted next-token target at every position, matching the dense all-position teacher-forcing shape mixle.models.language_model._forward_all_positions / LM.fit_pairs already consume (as opposed to mixle.data.stream_token_source’s one-target-per-window shape, which is the right shape for a single unbounded sliding stream but wastes a factor of block of compute once sequences are packed).

global_document_order(num_documents, *, seed, epoch, sequence_selector=None)[source]

The single, authoritative order documents are consumed in across ALL ranks combined for this epoch.

Same (seed, epoch) -> bitwise-identical order on every call/process (pure function of its inputs, RandomState seeded deterministically). Different epoch with the same seed -> a different, still fully deterministic order.

sequence_selector is the curriculum hook for E7 (“a bandit over length buckets” rationing ultra-long examples, per the roadmap): an optional (order, seed, epoch) -> order callback that may reorder, filter, or otherwise transform the base permutation before it is handed out to ranks. Left None, the base shuffle passes through unchanged. This module does not implement any curriculum policy itself – only the extension point.

Parameters:
Return type:

ndarray

shard_documents_for_rank(order, rank, world_size)[source]

Round-robin split of a (possibly shuffled/filtered) global document order across ranks.

Rank i gets order[i], order[i + world_size], order[i + 2*world_size], ... – the same round-robin residency contract MPEncodedData already uses for its worker shards, so the two compose: disjoint across ranks, complete coverage of order.

Parameters:
Return type:

ndarray

pack_documents(documents, indices, block, *, pad_id=0, boundary_id=None)[source]

Concatenate documents[indices] (in the given order) into one token stream and chunk it into fixed-length block + 1 rows – the standard “concatenate-and-chunk with document boundaries” packing scheme used in LM pretraining.

row[:-1] (length block) is the model input; row[1:] (length block) is the shifted next-token target at every position. Padding (pad_id) is only ever needed to fill out the final row once the stream runs out, so waste is bounded by block tokens total, not block tokens per document – packing efficiency (the real-token fraction) climbs toward 1.0 as the corpus grows relative to block. boundary_id (e.g. an EOS id), if given, is inserted between consecutive documents so the model can see document edges within a packed row; it counts as a real (non-pad) token.

Parameters:
Return type:

PackedCorpus

class PackedCorpus(rows, real_tokens, total_tokens)[source]

Bases: object

Result of pack_documents(): the packed rows plus the measured packing efficiency.

Parameters:

  • rows (np.ndarray)

  • real_tokens (int)

  • total_tokens (int)

class StreamingCorpus(documents, *, rank, world_size, block, batch_size, seed=0, pad_id=0, boundary_id=None, sequence_selector=None)[source]

Bases: object

Per-rank streaming view over a sharded, tokenized corpus: shuffled deterministically by (seed, epoch), packed into fixed-length dense-teacher-forcing rows, and batched.

documents is the full corpus (every rank sees the same list; only its OWN shard is ever packed or batched – no gather, no materializing another rank’s tokens). For a real out-of-core corpus, documents is a lazy/mmap-backed sequence of per-shard-file document lists; the contract here is unchanged since sharding and packing only ever index it, never buffer the whole thing.

Parameters:

  • documents (Sequence[Any])

  • rank (int)

  • world_size (int)

  • block (int)

  • batch_size (int)

  • seed (int)

  • pad_id (int)

  • boundary_id (int | None)

  • sequence_selector (SequenceSelector | None)

rank_document_indices(epoch)[source]

This rank’s document indices for epoch, in the exact order they will be packed.

Parameters:

epoch (int)

Return type:

ndarray

epoch_batches(epoch)[source]

Yield (context (b, block) float32, targets (b, block) int64) micro-batches for this rank.

Deterministic given (seed, epoch, rank, world_size): re-running with the same inputs yields bitwise-identical batches. Sets last_packing_efficiency as a side effect (the real-token fraction of the rows this call packed).

Parameters:

epoch (int)

Return type:

Iterator[tuple[ndarray, ndarray]]