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).

  • Padding is structural, not training signal. The final packed row may be padded with pad_id to fill out block + 1 tokens once the stream runs out – bounded waste of at most one row, ever, never more. pack_documents() returns a per-position loss_mask alongside the rows, and StreamingCorpus.epoch_batches() threads it through every micro-batch. Callers MUST apply that mask (or otherwise exclude its False positions) before averaging next-token loss – otherwise the fabricated pad targets are silently trained on as if they were real observed labels.

  • Token ids and every iteration control (block, batch_size, rank/world_size, …) are validated eagerly wherever they enter this module – at a function call or at StreamingCorpus construction – rather than being cast/truncated silently or left to fail lazily once packing/iteration is already underway.

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.

num_documents, seed, and epoch are validated as nonnegative exact integers at entry – num_documents=0 is a legitimate empty-corpus case (returns an empty order); a NEGATIVE or fractional value for any of the three previously either raised deep inside RandomState/SeedSequence with a misleading message, or (for num_documents specifically) was silently accepted by RandomState.permutation as an empty order regardless of magnitude (e.g. -3 behaved like 0).

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.

Every document’s tokens are validated before packing – finite, exact-integer values in a lossless int64 range; fractional, non-finite, or out-of-range ids raise rather than being silently cast (MXR-080-0062). block, pad_id, and boundary_id are validated the same way at entry; block must additionally be positive (MXR-080-0063) – there is no meaningful zero-or-negative-block packing.

Returns a PackedCorpus whose loss_mask (shape (n_rows, block), boolean) marks which TARGET positions (row[1:]) are real observed next-token labels vs. fabricated padding. Boundary policy, decided explicitly here: padding never appears anywhere but the tail of the FINAL row (see above), so loss_mask is True everywhere except possibly that one tail; a boundary_id token is real signal like any other token and is never masked out. This function does not implement SFT-style prompt masking (weight-0 on prompt tokens) – that is an orthogonal concern this codebase already handles elsewhere via per-sample weights; loss_mask here answers exactly one question, “is this target a fabricated pad position.” Callers MUST apply this mask (or otherwise exclude its False positions) before averaging next-token loss, or the fabricated pad targets are silently trained on as if they were real labels (MXR-080-0061).

Parameters:
Return type:

PackedCorpus

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

Bases: object

Result of pack_documents(): the packed rows, the measured packing efficiency, and a per-target loss_mask marking which target positions are real observed next-token labels vs. fabricated padding (see pack_documents() for the exact boundary policy).

Parameters:

  • rows (np.ndarray)

  • real_tokens (int)

  • total_tokens (int)

  • loss_mask (np.ndarray)

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.

Every constructor argument is validated eagerly, here, at construction – rank/world_size cross-checked together (0 <= rank < world_size), block/batch_size required positive, seed required a nonnegative exact integer, pad_id/boundary_id required exact integers. Previously an invalid rank/world_size/block/batch_size was stored unchecked and only surfaced (if at all) on the first epoch_batches/rank_document_indices call (MXR-080-0063).

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) int64, targets (b, block) int64, loss_mask (b, block) bool) micro-batches for this rank.

loss_mask[i, j] is True iff targets[i, j] is a real observed next-token label, and False only for fabricated padding introduced to fill out the final row of the packed stream once real tokens run out (see pack_documents() for the exact boundary policy) – at most one row per call carries any False at all, regardless of corpus size. Callers MUST multiply the per-position loss by this mask (or otherwise exclude False positions) rather than average the raw per-position loss unconditionally: doing so silently teaches the model that pad_id is a real next-token target (MXR-080-0061). context/targets stay int64 – never downcast to float32 here, which would silently lose integer identity for ids at or above 2**24 (MXR-080-0062); a consumer that wants float input casts explicitly on its own side.

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, ndarray]]