kempnerforge.training.entry

Library entry point for a KempnerForge training job.

scripts/train.py is a thin CLI wrapper over run_training(): it parses argv, loads the config, and calls in. Each phase below is a separate function so an experiment can reuse the scaffold — or swap one phase — without editing the shared script.

Functions

build_checkpoint_manager(config, runtime, ...)

Create the checkpoint manager, scoping DCP to the PP stage when needed.

build_model(config, runtime, loss_fn)

Build the model with the full parallelism stack applied.

restore_checkpoint(config, model, scheduler, ...)

Auto-resume from latest (or checkpoint.load_path).

run_training(config, *[, step_fn, hooks])

Run a full training job: build every phase, run the loop, tear down.
kempnerforge.training.entry.build_model(config, runtime, loss_fn)[source]

Build the model with the full parallelism stack applied.

With distributed.pp > 1 this builds this rank’s pipeline stage plus its schedule (which is why loss_fn is needed here); otherwise it delegates to build_parallel_model and returns no pipeline.

Parameters:
Return type:

tuple[torch.nn.Module, PipelineBundle | None]

kempnerforge.training.entry.build_checkpoint_manager(config, runtime, model, optimizer, pipeline)[source]

Create the checkpoint manager, scoping DCP to the PP stage when needed.

With PP, each stage has different parameters — DCP needs a group scoped to ranks within the same PP stage (all non-PP mesh dimensions), and each stage saves DCP shards to its own subdirectory to avoid file collisions.

Parameters:
Return type:

CheckpointManager

kempnerforge.training.entry.restore_checkpoint(config, model, scheduler, ckpt_mgr)[source]

Auto-resume from latest (or checkpoint.load_path).

Returns (step, tokens_seen) — (0, 0) when there is nothing to resume from. Also restores the metrics run ids, runs the MoT warm-start hook, and re-applies the effective freeze at the resumed step.

Parameters:
Return type:

tuple[int, int]

kempnerforge.training.entry.run_training(config, *, step_fn=None, hooks=None)[source]

Run a full training job: build every phase, run the loop, tear down.

step_fn and hooks let an experiment own the step body or register hooks without copying the build phases. Both default to what scripts/train.py uses.

Parameters:
Return type:

None