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. The manager records each rank’s data-parallel coordinates, the partition build_data_pipeline reads, with the per-rank loader/RNG state.

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 text warm start when no checkpoint supplied the model weights and the run starts at step 0, and re-applies the effective freeze at the resumed step.

A checkpoint in checkpoint.dir takes precedence over checkpoint.load_path, so a requeued warm start resumes rather than starting over; rank 0 then warns that load_path is not used, naming both paths, unless load_path names the resumed checkpoint and excludes nothing.

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. A custom step_fn must run its forward and loss under session.autocast(), which casts fp32 master weights to the compute dtype when no FSDP2 wrap does (one GPU, a pp-only mesh).

Parameters:
Return type:

None