sparse_data#

Sparse teacher-policy dataset for external draft training.

An alternative to caching base hidden states: the TVD objectives need only the teacher’s deployment policy (top-k then top-p at the serving temperature), which is far smaller than a hidden state per token and is captured during generation rather than by a second forward pass.

Records are gzipped JSONL, one conversation per line:

{
    "id": str,
    "prompt_ids": [int],  # context, not scored
    "gen_ids": [int],  # response tokens, scored
    "topk_tok": [[int]],  # per response position, teacher support
    "topk_lp": [[float]],
}  # matching log-probabilities

They are decoded once into a memmapped cache beside the shards, stored as int32 ids and float16 probabilities; holding them as Python objects does not fit in memory for a corpus of any size.

Classes

SparsePolicyCollator

Pad or truncate a batch of sparse-policy records to train_len.

SparsePolicyDataset

Conversations with the teacher's sparse deployment policy per response token.

Functions

discover_shards

Every sparse-policy shard under root, in the order the cache is keyed on.
class SparsePolicyCollator#

Bases: object

Pad or truncate a batch of sparse-policy records to train_len.

__init__(train_len)#

Set the fixed sequence length every batch is padded or truncated to.

Parameters:

train_len (int)

class SparsePolicyDataset#

Bases: Dataset

Conversations with the teacher’s sparse deployment policy per response token.

__init__(files, max_length=4096, top_k=None, cache_dir=None, limit=None, source_weights=None)#

Open (building if needed) the memmapped cache for these shards.

top_k defaults to the width found in the data. A narrower value truncates the teacher policy and a wider one pads it with zero-probability entries; both change the objective silently.

source_weights maps a record’s source field to a sampling weight, oversampling under-represented sources by repeating their index entries. Weight by tokens, not rows: sources differ by an order of magnitude in response length, so balanced row counts can still leave one source holding most of the trained positions.

Parameters:

  • files (list[str])

  • max_length (int)

  • top_k (int | None)

  • cache_dir (str | None)

  • limit (int | None)

  • source_weights (dict[str, float] | None)

discover_shards(root)#

Every sparse-policy shard under root, in the order the cache is keyed on.

Shared with the training entrypoint, and named in the pre-build instructions, so a cache built by hand is keyed identically to the one training would build. A command that globbed differently – one pattern, or non-recursively – would produce a different _cache_signature, and training would rebuild from scratch under torchrun having apparently pre-built nothing.

Both suffixes are combined rather than preferred: an or would silently drop every .jsonl shard from a directory that also holds a .jsonl.gz.

Return type:

list[str]