Training engine

elfes.engine connects training, evaluation, checkpoint selection, and resumable stopping for CPU, CUDA, and replicated DDP. Workflows supply ordinary PyTorch models, optimizers, DataLoaders, and step functions.

Train a mean objective

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")

loss is a local mean. Its denominator defaults to num_samples; pass loss_weight when it is an atom count, number of labeled components, or another explicit weight. Each update divides the weighted loss sum by the actual global denominator across its accumulation window, including incomplete tails. Terms combined into one scalar mean must share the declared denominator. Use Objective for independently normalized terms.

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.

Independently normalized objectives

from elfes.engine import Objective

objective = Objective(
    weights={"primary": 1.0, "auxiliary": 0.2},
    counts=lambda batch: {
        "primary": batch["target"].numel(),
        "auxiliary": batch["mask"].sum(),
    },
)


def step(model, batch):
    primary, auxiliary = model(batch["inputs"])
    mask = batch["mask"]
    return StepOutput(
        loss={
            "primary": (primary - batch["target"]).square().sum(),
            "auxiliary": (auxiliary[mask] - batch["aux_target"][mask]).square().sum(),
        },
        num_samples=len(batch["inputs"]),
    )


evaluator = Evaluator(accelerator, {"val": val_loader}, step, objective=objective)
trainer = Trainer(
    accelerator, model, optimizer, train_loader, step,
    objective=objective, evaluator=evaluator,
)

With an Objective, loss values are differentiable sums, and coefficients live in objective.weights. The engine sums each term's counts over the actual global accumulation window, then optimizes the weighted sum of term means. counts(batch) must read deterministic, nonnegative, finite data counts without model execution, random sampling, or mutation. Counts may be fractional but cannot depend on trainable parameters.

This path buffers one window of input batches on the execution device, while retaining only one microbatch's forward graph. The scalar mean path streams batches directly. Mask out missing targets before forming the loss; an empty indexed error sum supplies a differentiable zero. Globally absent terms contribute no gradient. A wholly unlabelled window skips optimizer, EMA, and step scheduler, including momentum and weight decay. Evaluation reports zero counts and omits unobserved means instead of reporting perfect zero error.

Metrics include loss/<name>, loss_count/<name>, and the weighted total loss. Loss names and the loss_sum/ internal prefix are reserved. Additional physical metrics remain independent of the training objective.

Optimization recipes

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, objective=objective)
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 decay reaches its minimum at total_steps. No warmup 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.

Run records

Pass scientific configuration through trainer.train(..., metadata={...}) and configure RunLogger to persist UTC-stamped events. Each started/resumed call records the execution configuration, budget, world size, precision, optimizer groups, EMA, objective coefficients and supplied metadata. Workflows supply units, normalization, data splits, seeds, and the concrete scheduler recipe.

Train metrics report per-group lr/0, lr/1, etc. used by the most recent attempted update, pre-clipping grad_norm averaged over successful updates, actual last_update_samples, and samples_per_second excluding evaluation and checkpoint phases. Evaluation records raw/EMA selection and duration; checkpoint records the chosen position and best/final flags; export records the weight choice and path. Finished/stopped events summarize per-call train, evaluation, and checkpoint wall time. These diagnostic timings are not a substitute for controlled performance benchmarks.

Evaluate raw or EMA weights

raw_metrics = trainer.evaluate(weights="raw")
ema_metrics = trainer.evaluate(weights="ema")

Periodic training evaluation uses EMA when enabled, and best-model selection follows those metrics. Explicit evaluation and export APIs default to raw weights. An EMA request fails when EMA is absent; it never silently falls back to raw.

Evaluator requires drop_last=False and counts every real sample once, including uneven or empty rank shards. It restores model mode and training RNG after evaluation. Set requires_grad=True for predictions derived through autograd, such as conservative forces. Distributed evaluation unwraps DDP for independent forwards while retaining compilation and mixed precision; evaluation steps must not perform their own distributed collectives.

Load a historical checkpoint without a training setup

from elfes.engine import export_weights, load_weights

load_weights(model, run_dir / "checkpoints" / "best", weights="ema")
model = accelerator.prepare(model)
results = evaluator.evaluate(model)
export_weights(run_dir / "checkpoints" / "best", output_path, weights="ema")

Pass an unwrapped, uncompiled model to load_weights. No optimizer or training DataLoader is needed. Exports are ordinary model state dictionaries, with no DDP or compilation prefixes. EMA includes checkpoint-time raw model buffers. best selects the training position; weights selects the parameter set at that position.

Stop and resume

StopRequest.catch_signals() temporarily handles SIGINT/SIGTERM in the main thread. Any rank may also call stop.request(reason). The trainer finishes the current accumulation window, saves last, skips final evaluation, and returns with state.stop_reason. Signal handlers only set a flag; saving and communication run in the training loop.

state = trainer.train(
    TrainConfig(epochs=100),
    checkpoint_dir=run_dir / "checkpoints",
    resume_from=run_dir / "checkpoints" / "last",
)

Checkpoints atomically publish one complete PyTorch file containing named raw/EMA weights, optimizer/scheduler/scaler state, data progress, and per-rank RNG state. last, best, and final are relative symlinks. A requested pause updates last; a normal endpoint updates final. Retention preserves all three pointers' targets.

Exact continuation requires deterministic map-style Dataset reads and collation, unchanged data/batching/world size, and the same objective/count definition, optimizer, scheduler recipe, EMA and precision setup. Epoch-addressed shuffle and model randomness such as Dropout are restored. Worker-local random transforms, custom stateful samplers, and streaming data need a separate data-state protocol and are outside this guarantee. Increasing the cumulative epoch/step limits is allowed; using a shorter max_steps as a normal endpoint may add a final evaluation and alter a plateau schedule, so use StopRequest for a transparent pause.

This checkpoint format replaces the old Accelerate checkpoint directories. Historical runs continue with their frozen training code.

Interval(every: int, unit: Literal['step', 'epoch'])

Monitor(
    metric: str,
    mode: Literal["min", "max"] = "min",
    patience: int | None = None,
    min_delta: float = 0.0,
)

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",
)

StopRequest()

catch_signals() -> Iterator[StopRequest]

Evaluator(
    accelerator: Accelerator,
    dataloaders: Mapping[str, DataLoader[Any]],
    step: Step,
    *,
    requires_grad: bool = False,
    objective: Objective | None = None,
)

RunLogger(run_dir: Path, *, is_main_process: bool, console: Console | None = None)

log(event: Mapping[str, Any]) -> None

close() -> None

MaxMetric(value: Scalar)

MeanMetric(total: Scalar, count: Scalar)

RootMeanSquareMetric(sum_squared: Scalar, count: Scalar)

SumMetric(total: Scalar)

Objective(weights: Mapping[str, float], counts: Callable[[Any], Mapping[str, Scalar]])

WarmupDecay(
    total_steps: int,
    warmup_steps: int = 0,
    start_factor: float = 0.0,
    minimum_factor: float = 0.0,
    decay: Literal["cosine", "linear", "constant"] = "cosine",
)

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,
)

StepOutput(
    loss: Tensor | Mapping[str, Tensor],
    num_samples: int,
    metrics: dict[str, Metric] = dict(),
    loss_weight: Scalar | None = None,
)

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,
)

updates_per_epoch: int

planned_steps(config: TrainConfig) -> int

train(
    config: TrainConfig,
    *,
    checkpoint_dir: Path | None = None,
    resume_from: Path | None = None,
    stop: StopRequest | None = None,
    metadata: Mapping[str, Any] | None = None,
) -> TrainState

evaluate(*, weights: Weights = 'raw') -> dict[str, dict[str, float]]

save_model(path: Path, *, weights: Weights = 'raw') -> None

export_weights(checkpoint: Path, path: Path, *, weights: Weights = 'raw') -> None

load_weights(model: Module, checkpoint: Path, *, weights: Weights = 'raw') -> None

weight_decay_groups(
    model: Module,
    weight_decay: float,
    *,
    exclude: Callable[[str, Parameter], bool] = lambda name, parameter: ndim < 2,
) -> list[dict[str, Any]]