kempnerforge.training.data_pipeline

Training/eval data pipeline construction and phase (annealing) scheduling.

The builders here own every [data] / [eval] branch that used to sit inline in the training entry point: pre-tokenized mmap, HuggingFace (eager or streaming), multi-dataset mixtures, and the VLM image/video paths.

Functions

advance_phases(state, data, step)

Activate every phase whose start_step has been reached.

build_data_pipeline(config, runtime)

Build the training data pipeline for whichever [data] source is set.

build_eval_dataloader(config, runtime)

Build the eval dataloader, or None when eval is off or unsupported.

build_phase_state(config, data, step)

Resolve data-annealing phases and re-derive the active one on resume.

Classes

DataPipeline

The training dataset/loader plus the mixture state phases need.

PhaseState

Data-annealing phases and the currently active scaling.
class kempnerforge.training.data_pipeline.DataPipeline[source]

Bases: object

The training dataset/loader plus the mixture state phases need.

dataloader is None means no data source was configured; the text and PP step bodies then fall back to random tokens. dp_rank / dp_size record the data-parallel partition the samplers were built for.

dataloader: Any | None = None
dp_rank: int = 0
dp_size: int = 1
mixture_dataset: MixtureDataset | None = None
mixture_sampler: MixtureSampler | None = None
mixture_weights: dict[str, float]
__init__(dataloader=None, dp_rank=0, dp_size=1, mixture_dataset=None, mixture_sampler=None, mixture_weights=<factory>)
Parameters:
Return type:

None

class kempnerforge.training.data_pipeline.PhaseState[source]

Bases: object

Data-annealing phases and the currently active scaling.

phases: list[TrainingPhase]
original_weights: dict[str, float]
temperature: float = 1.0
next_idx: int = 0
lr_scale: float = 1.0
__init__(phases=<factory>, original_weights=<factory>, temperature=1.0, next_idx=0, lr_scale=1.0)
Parameters:
Return type:

None

kempnerforge.training.data_pipeline.build_data_pipeline(config, runtime)[source]

Build the training data pipeline for whichever [data] source is set.

Returns an empty pipeline when no source is configured — the text and PP step bodies then run on random tokens.

Parameters:
Return type:

DataPipeline

kempnerforge.training.data_pipeline.build_eval_dataloader(config, runtime)[source]

Build the eval dataloader, or None when eval is off or unsupported.

VLM + eval is out of scope on this branch: run_eval calls model(input_ids), which does not match VLMWrapper.forward.

Parameters:
Return type:

Any | None

kempnerforge.training.data_pipeline.build_phase_state(config, data, step)[source]

Resolve data-annealing phases and re-derive the active one on resume.

Parameters:
Return type:

PhaseState

kempnerforge.training.data_pipeline.advance_phases(state, data, step)[source]

Activate every phase whose start_step has been reached.

Returns True when at least one phase fired, so the caller can drop the materialized data iterator and pick up the new sampler weights.

Parameters:
Return type:

bool