Training and optimization
A training run starts with a model, a DataLoader, an optimizer, and a step function that computes the task's loss. Trainer manages gradient accumulation, parameter updates, and the timing of evaluation and checkpoints. The same setup supports CPU, CUDA, and replicated distributed data parallel (DDP) execution through Accelerate.
A training run
This example assumes a model, optimizer, training and validation loaders, and a Path named run_dir. The step returns a mean squared error and an additional error metric. Training evaluates averaged parameters, saves checkpoints, and can pause at an optimizer-update boundary.
from accelerate import Accelerator
from elfes.engine import (
Evaluator, MeanMetric, Monitor, RunLogger, StepOutput,
StopRequest, TrainConfig, Trainer,
)
def step(model, batch):
inputs, targets = batch
error = model(inputs) - targets
return StepOutput(
loss=error.square().mean(),
num_samples=len(inputs),
metrics={"mae": MeanMetric(error.abs().sum(), error.numel())},
)
accelerator = Accelerator(gradient_accumulation_steps=2)
evaluator = Evaluator(accelerator, {"validation": validation_loader}, step)
with RunLogger(run_dir, is_main_process=accelerator.is_main_process) as logger:
trainer = Trainer(
accelerator, model, optimizer, train_loader, step,
evaluator=evaluator, ema_decay=0.999, logger=logger,
)
stop = StopRequest()
with stop.catch_signals():
state = trainer.train(
TrainConfig(epochs=100, monitor=Monitor("validation/loss")),
checkpoint_dir=run_dir / "checkpoints",
stop=stop,
)
if stop.reason is None:
trainer.save_model(run_dir / "model.pt", weights="ema")
The meaning of the loss denominator and additional metrics is explained in Objectives and metrics. Evaluation covers EMA and best-model selection; Checkpoints and resuming describes the stop and export operations. Run logging explains the recorded events.
Updates and data progress
Training enters model.train(). Nonfinite losses and non-AMP nonfinite gradients fail before updating parameters. AMP overflow skips the optimizer update and EMA while advancing the consumed-data position. TrainState.step counts successful updates; samples counts globally consumed samples.
The default Accelerate even_batches=True policy may pad distributed training shards. Exact sampling requires even_batches=False, split_batches=False, and a batch sampler whose total batch count is divisible by the world size; ELFES verifies the resulting per-rank lengths before training begins.
Optimization recipes
The following setup uses the same mean-loss step. The scheduler is constructed after the trainer so its budget follows the prepared training loader.
import torch
from torch.optim.lr_scheduler import LambdaLR
from elfes.engine import WarmupDecay, weight_decay_groups
optimizer = torch.optim.AdamW(
weight_decay_groups(model, 0.01), lr=1e-3,
)
trainer = Trainer(accelerator, model, optimizer, train_loader, step)
config = TrainConfig(epochs=100)
schedule = WarmupDecay(
total_steps=trainer.planned_steps(config),
warmup_steps=5 * trainer.updates_per_epoch,
start_factor=0.1,
minimum_factor=0.01,
decay="cosine",
)
trainer.scheduler = LambdaLR(optimizer, schedule)
updates_per_epoch uses the distributed prepared loader and includes incomplete accumulation tails. planned_steps is the cumulative budget, capped by max_steps, assuming no skipped updates or early stopping. Build the scheduler before loading a checkpoint; its state is restored alongside the optimizer.
WarmupDecay supports cosine, linear, and constant post-warmup stages. Its argument counts completed successful updates: the first update uses factor(0), warmup reaches one at warmup_steps, and cosine or linear decay reaches its minimum at total_steps. With warmup_steps=0, the schedule starts at one. The remaining decay budget excludes warmup. LambdaLR checkpoints preserve the callable's named fields. Other native schedulers, including SequentialLR and ReduceLROnPlateau, remain directly usable.
weight_decay_groups excludes frozen parameters and deduplicates tied parameters. Its default exempts scalar/vector parameters from decay; pass exclude(name, parameter) for model-specific policies. Returned groups are ordinary mutable optimizer dictionaries with readable names and parameter names, so callers may set separate learning rates or use hand-written groups.
Training API
TrainConfig sets cumulative epoch and step limits, clipping, and evaluation, checkpoint, and logging frequency. Interval expresses repeated evaluation or checkpoint triggers in successful optimizer steps or completed epochs. TrainState records the resulting data position and update count. Best-model and early-stopping settings are described with Monitor.
Trainer(
accelerator: Accelerator,
model: Module,
optimizer: Optimizer,
train_dataloader: DataLoader,
step: Step,
*,
evaluator: Evaluator | None = None,
objective: Objective | None = None,
scheduler: LRScheduler | ReduceLROnPlateau | None = None,
ema_decay: float | None = None,
logger: RunLogger | None = None,
)
Train one model/optimizer on CPU, CUDA, or replicated DDP.
The step returns a mean loss and its denominator. Gradients are normalized over the actual global accumulation window, including its incomplete tail. Deterministic map-style datasets and the same data/batching/world size are required for exact continuation. Model RNG is checkpointed per rank.
property
updates_per_epoch: int
Number of accumulation windows after distributed loader preparation.
planned_steps(config: TrainConfig) -> int
Total update budget, assuming no AMP overflow or early stopping.
train(
config: TrainConfig,
*,
checkpoint_dir: Path | None = None,
resume_from: Path | None = None,
stop: StopRequest | None = None,
metadata: Mapping[str, Any] | None = None,
) -> TrainState
Continue until the epoch/step limit, early stopping, or a stop request.
A requested stop saves last and skips final evaluation. final denotes
a normally finished train call. Signals are opt-in via StopRequest.
evaluate(*, weights: Weights = 'raw') -> dict[str, dict[str, float]]
Evaluate current raw/EMA weights without changing training parameters.
save_model(path: Path, *, weights: Weights = 'raw') -> None
Export current raw/EMA state_dict without a compiled/DDP prefix.
dataclass
TrainConfig(
epochs: int,
max_steps: int | None = None,
log_every: int = 10,
evaluate_every: Interval | None = Interval(1, "epoch"),
checkpoint_every: Interval | None = Interval(1, "epoch"),
keep_checkpoints: int | None = None,
max_grad_norm: float | None = None,
monitor: Monitor | None = None,
scheduler_interval: Literal["step", "epoch"] = "step",
)
Policies for one call to Trainer.train().
dataclass
Interval(every: int, unit: Literal['step', 'epoch'])
A repeated trigger measured in optimizer steps or completed epochs.
dataclass
TrainState(
epoch: int = 0,
next_batch: int = 0,
step: int = 0,
batches: int = 0,
samples: int = 0,
best_metric: float | None = None,
best_step: int | None = None,
evaluations_without_improvement: int = 0,
stop_reason: str | None = None,
)
Data progress and successful optimizer updates are separate counters.
epoch counts completed epochs; next_batch is the next local microbatch. batches counts consumed local microbatches, including AMP-skipped updates. samples counts consumed samples globally, including training padding when the Accelerator DataLoader configuration enables even batches.
Optimization API
dataclass
WarmupDecay(
total_steps: int,
warmup_steps: int = 0,
start_factor: float = 0.0,
minimum_factor: float = 0.0,
decay: Literal["cosine", "linear", "constant"] = "cosine",
)
LR multiplier for LambdaLR, indexed by completed optimizer updates.
Linear warmup goes from start_factor at zero to one at warmup_steps. The remaining budget decays to minimum_factor at total_steps and stays there. With no warmup, the initial factor is one. Use scheduler_interval='step'.
weight_decay_groups(
model: Module,
weight_decay: float,
*,
exclude: Callable[[str, Parameter], bool] = lambda name, parameter: ndim < 2,
) -> list[dict[str, Any]]
Group trainable parameters once, including tied parameters only once.
By default scalar/vector parameters (including typical biases and norm scales) have no decay. Override exclude for the model's actual parameter semantics. Returned ordinary optimizer groups may be edited to set separate learning rates.