kempnerforge.training

Training loop and optimization for KempnerForge.

Public API:

• run_training: Build and run a full training job from a JobConfig

• build_loss_fn / build_optimizer / build_scheduler: Component factories

• run_eval: Evaluation loop (loss + perplexity)

• maybe_no_sync: Gradient accumulation helper

kempnerforge.training.build_loss_fn(config)

Build a composed loss function from training config.

Follows the build_optimizer pattern: config in, callable out. Binds chunk_size for chunked CE and fp32 when train.fp32_cross_entropy is set, and composes z-loss, so the caller gets a clean (logits, labels) -> Tensor interface.

Return type:

Callable[[torch.Tensor, torch.Tensor], torch.Tensor]

kempnerforge.training.build_optimizer(model, config, module_patterns=None)

Construct an optimizer with per-parameter-group weight decay settings.

With config.lr_multipliers set, each listed module’s parameters get their own groups at config.lr times the module’s multiplier, named by "module"; the scheduler then scales every group’s own rate by the same factor. A module nested in another listed one keeps its own parameters.

Parameters:

• model (torch.nn.Module) – Model whose parameters to optimize.

• config (OptimizerConfig) – Optimizer configuration.

• module_patterns (Mapping[str, Sequence[str]] | None) – Module name to the fnmatch patterns of its parameter names, such as VLMConfig.module_patterns. Required when config.lr_multipliers is set.

Returns:

Configured optimizer instance.

Raises:

ValueError – config.lr_multipliers names a module module_patterns lacks or one that matches no parameter, or two of its modules match the same parameters or parameters that partly overlap.

Return type:

torch.optim.Optimizer

kempnerforge.training.build_scheduler(optimizer, config, max_steps)

Build a LR scheduler from config.

Parameters:

• optimizer (torch.optim.Optimizer) – Optimizer to schedule.

• config (SchedulerConfig) – Scheduler configuration.

• max_steps (int) – Total training steps (used to compute decay length).

Returns:

PyTorch LambdaLR scheduler.

Return type:

torch.optim.lr_scheduler.LambdaLR

kempnerforge.training.maybe_no_sync(model, micro_step, grad_accum_steps)

Skip gradient sync on intermediate accumulation steps.

On the last micro-step, gradients are synchronized normally. On earlier micro-steps, sync is skipped to avoid redundant all-reduces.

Works with FSDP2 (which implements no_sync() on the module). For non-distributed models, this is a no-op.

Parameters:

• model (torch.nn.Module)

• micro_step (int)

• grad_accum_steps (int)

Return type:

Generator[None, None, None]

kempnerforge.training.run_eval(model, eval_dataloader, loss_fn, device, eval_steps, *, pp_schedule=None, pp_rank=None, pp_size=None, pp_group=None, pp_n_microbatches=None)

Run evaluation and return metrics.

Parameters:

• model (torch.nn.Module) – The model (FSDP-wrapped, TP-sharded, or plain).

• eval_dataloader (torch.utils.data.DataLoader) – DataLoader yielding {“input_ids”, “labels”} batches.

• loss_fn (callable) – Loss function (logits, labels) -> scalar tensor.

• device (torch.device) – Device to move batches to.

• eval_steps (int) – Number of eval batches to process.

• pp_schedule – Pipeline parallel schedule (None for non-PP). It is the training schedule, run forward-only through its eval(). Its gradient reduction still runs after each call, with its gradient scaling off, so gradients left on the stage keep their values.

• pp_rank (int | None) – This rank’s PP stage index.

• pp_size (int | None) – Total number of PP stages.

• pp_group – Process group for PP loss broadcast.

• pp_n_microbatches (int | None) – Microbatches per call of pp_schedule. The batches go through it this many at a time, so each microbatch has the shape the stages were built for; eval_steps must be a multiple of it. None sends all eval_steps batches in one call.

Returns:

Dict with “eval/loss” and “eval/perplexity”.

Raises:

ValueError – Under PP, eval_steps is not a multiple of pp_n_microbatches, or eval_dataloader yields no batch.

Return type:

dict[str, float]

kempnerforge.training.run_training(config, *, step_fn=None, hooks=None)

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:

• config (JobConfig)

• step_fn (Callable[[TrainingSession, int], StepResult] | None)

• hooks (HookRunner | None)

Return type:

None

Modules

data_pipeline	Training/eval data pipeline construction and phase (annealing) scheduling.
entry	Library entry point for a KempnerForge training job.
eval	Evaluation utilities for KempnerForge.
freeze	Parameter freezing helpers.
grad	Gradient utilities for distributed training.
hooks	Training loop hooks for extensibility without forking train.py.
loop	The training step loop and the three step bodies it dispatches to.
loss	Loss function registry for KempnerForge.
optimizer	Optimizer construction for KempnerForge.
orthogonalize	Exact Newton-Schulz orthogonalization of sharded matrices for Muon.
runtime	Process-level context for a training job: ranks, device, mesh, pipeline.
scheduler	Learning rate schedulers for KempnerForge.