kempnerforge.training.loop

The training step loop and the three step bodies it dispatches to.

run_training_loop owns everything that repeats: the step body, NaN handling, optimizer/scheduler advance, freeze and phase transitions, metrics, eval, checkpointing, and shutdown. The step body itself is a plain callable (StepFn) chosen once by select_step_fn(), so a model family with a different forward contract is a new function rather than another branch inside the loop.

Functions

add_moe_losses(loss, model, mc)

MoE auxiliary + router z-loss terms (only called when mc.is_moe).

checkpoint_extra(config, step, phases)

Metadata saved alongside every checkpoint so a resume is exact.

clip_grads(model, max_norm)

pipeline_step(session, step)

PP step: collect microbatches, hand them to the schedule, broadcast loss.

run_training_loop(session, *[, step, ...])

Run the step loop from step to train.max_steps.

select_step_fn(config)

Pick the step body for this job: PP wins, then VLM, then text-only.

text_step(session, step)

Standard step (no PP, text-only), with optional per-dataset mixture metrics.

vlm_step(session, step)

VLM step (no PP): pixel_values + input_ids forward, with a text-token count.

Classes

BatchStream

The training dataloader plus its live iterator.

StepResult

What one training step reports back to the loop.

TrainingSession

Everything one training run needs after the build phase.
class kempnerforge.training.loop.BatchStream[source]

Bases: object

The training dataloader plus its live iterator.

Wraps the epoch-boundary restart and the “no dataloader configured” case so the step bodies don’t each re-implement them. The loader is read off the pipeline on every access, so swapping pipeline.dataloader takes effect on the next step instead of leaving a stale snapshot.

__init__(pipeline)[source]
Parameters:

pipeline (DataPipeline)

Return type:

None

property dataloader: Any
property has_data: bool
ensure_started()[source]

Materialize the iterator, if it isn’t already, for the loader in use.

next_batch does this itself; the loop calls it up front so building an iterator (which spawns dataloader workers) stays outside the region MetricsTracker times.

Return type:

None

reset()[source]

Drop the iterator so the next step takes a fresh one from the loader.

A StatefulDataLoader re-applies its recorded skip, so this picks up inside the current epoch rather than restarting it.

Return type:

None

next_batch()[source]
Return type:

dict[str, torch.Tensor]

class kempnerforge.training.loop.StepResult[source]

Bases: object

What one training step reports back to the loop.

The per-dataset fields stay empty for step bodies that don’t produce them; text_tokens stays None, which a VLM run treats as an error rather than logging a zero.

loss: float
grad_norm: float
text_tokens: int | None = None
dataset_token_counts: dict[str, int]
dataset_loss_sums: dict[str, float]
dataset_loss_counts: dict[str, int]
__init__(loss, grad_norm, text_tokens=None, dataset_token_counts=<factory>, dataset_loss_sums=<factory>, dataset_loss_counts=<factory>)
Parameters:
Return type:

None

class kempnerforge.training.loop.TrainingSession[source]

Bases: object

Everything one training run needs after the build phase.

config: JobConfig
runtime: RuntimeContext
model: torch.nn.Module
optimizer: torch.optim.Optimizer
scheduler: Any
loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
step_fn: Callable[[TrainingSession, int], StepResult]
data: DataPipeline
phases: PhaseState
checkpointer: CheckpointManager
tracker: MetricsTracker
hooks: HookRunner
nan_detector: NaNDetector
shutdown_handler: ShutdownHandler
pipeline: PipelineBundle | None = None
eval_dataloader: Any | None = None
profiler: Any | None = None
property batches: BatchStream

The live batch iterator. Rebuilt if data is replaced wholesale.

__init__(config, runtime, model, optimizer, scheduler, loss_fn, step_fn, data, phases, checkpointer, tracker, hooks, nan_detector, shutdown_handler, pipeline=None, eval_dataloader=None, profiler=None)
Parameters:
Return type:

None

kempnerforge.training.loop.clip_grads(model, max_norm)[source]
Parameters:
Return type:

float

kempnerforge.training.loop.add_moe_losses(loss, model, mc)[source]

MoE auxiliary + router z-loss terms (only called when mc.is_moe).

Parameters:
Return type:

torch.Tensor

kempnerforge.training.loop.pipeline_step(session, step)[source]

PP step: collect microbatches, hand them to the schedule, broadcast loss.

step is unused — MoE (the only step-dependent branch) is rejected with PP in JobConfig.validate.

Parameters:
Return type:

StepResult

kempnerforge.training.loop.vlm_step(session, step)[source]

VLM step (no PP): pixel_values + input_ids forward, with a text-token count.

Parameters:
Return type:

StepResult

kempnerforge.training.loop.text_step(session, step)[source]

Standard step (no PP, text-only), with optional per-dataset mixture metrics.

Parameters:
Return type:

StepResult

kempnerforge.training.loop.select_step_fn(config)[source]

Pick the step body for this job: PP wins, then VLM, then text-only.

Parameters:

config (JobConfig)

Return type:

Callable[[TrainingSession, int], StepResult]

kempnerforge.training.loop.checkpoint_extra(config, step, phases)[source]

Metadata saved alongside every checkpoint so a resume is exact.

vlm_freeze reflects the post-transition state when a FreezeStage has fired (see freeze_meta_at_step).

Parameters:
Return type:

dict

kempnerforge.training.loop.run_training_loop(session, *, step=0, tokens_seen=0)[source]

Run the step loop from step to train.max_steps.

Returns the final (step, tokens_seen). Drains any pending async checkpoint before returning so the caller can tear down the process group safely. session.data is read once and must stay fixed for the duration; only its dataloader may be swapped mid-run.

Parameters:
Return type:

tuple[int, int]