BaseTrainer

class core.trainer.BaseTrainer(cfg, rank, world_size)

Bases: object

The epoch loop, and everything around it that depends on neither task nor model.

Logging, checkpointing, DDP, early stopping and best-model selection live here; a subclass supplies setup() and the epoch internals.

train()

Run the full training loop.

Iterates over epochs, calling train_epoch() each epoch and val_epoch() every val.every_n_epochs epochs and at the final epoch. Checkpoints are saved after each epoch. Supports early stopping via the return value of val_epoch().

Return type:

dict | None

Returns:

Final metrics dict, or None if no metrics were logged.

abstract train_epoch()

Perform a training epoch.

abstract val_epoch()

Perform a validation epoch.

Called every val.every_n_epochs epochs, and at the end of training.

Returns:

bool, True if early stopping should be triggered, False (or falsey) otherwise.

Return type:

early_stop

test()

Perform a test epoch.

Called at the end of training if test.enable is True.

log_training_results()

Log the training results, at the end of training.

abstract setup(ckpt)

Set up the trainer.

This method is called after setting up distributed, loading the checkpoint, and setting up logging.

Parameters:

ckpt (dict | None) – The checkpoint dictionary, if any.

add_checkpoint_items(**kwargs)

Add items to the checkpoint dictionary, saved and restored by state_dict.

Registering is what makes an object survive a resume, so every rank does it; writing the file stays rank zero’s job.

Parameters:

kwargs – The items to add to the checkpoint dictionary.

make_ddp(module, find_unused_parameters=False)

Wrap a module in DistributedDataParallel if running in distributed mode.

Parameters:

• module (Module) – The module to wrap.

• find_unused_parameters (bool) – Passed to DDP. Set to True if some parameters are not used in the forward pass. Defaults to False.

Returns:

The DDP-wrapped module, or the original module if not distributed.

Return type:

torch.nn.Module

barrier()

Synchronize all processes in the distributed group.

reduce_mean(total, count)

Mean over the whole split rather than over this rank’s shard of it.

A distributed sampler hands every rank a different slice, so a locally divided sum leaves the ranks disagreeing on the score they select on. Every rank must call this the same number of times, in the same order.

Parameters:

• total (float) – This rank’s summed quantity.

• count (float) – This rank’s number of terms in that sum.

Return type:

float

Returns:

The pooled mean, NaN if no rank contributed a term.

reset_best_tracking(minimize=None)

Initialize the best-score bookkeeping used by save_if_best().

save_if_best() calls this itself on first use, so a trainer only needs it to pin a direction that config does not carry: pass minimize when the trainer fixes it in code, or call it from setup after deciding val.minimize there. Direction otherwise comes from val.minimize, defaulting to maximizing.

A trainer that never tracks a best never gets these attributes, so it stays free to use those names itself. Only best_model and best_metrics are always defined, since code outside this method reads them.

save_if_best(score, metrics=None)

Keep this epoch’s weights if its score is the best seen so far.

On an improvement: records the score and epoch, snapshots the weights into best_model, and writes best.pt. The two have different readers: best_model is what test reloads in this same process, and it survives ckpt.enable=false; best.pt is what a later run loads. Pass metrics to also record the best/val/* dict for logging. Feed the result to step_patience() to early-stop as well; a trainer that only wants best-model selection can call this alone.

Call it on every rank, since every rank’s test reloads its own best_model; the best.pt write self-gates to rank 0 on its own. A trainer whose score only exists on rank 0 may call it under a rank guard, but then owes step_patience() a verdict broadcast to the other ranks.

Parameters:

• score (float) – The scalar being tracked, in val.minimize direction. A non-finite score never counts as an improvement.

• metrics (Optional[dict]) – The epoch’s val metrics, if they should be recorded too.

Return type:

bool

Returns:

True if the score improved, i.e. the weights were kept.

step_patience(improved)

Advance the early-stopping counter with this epoch’s outcome.

An improvement refills the counter; otherwise it ticks down, but only once the epoch reaches val.start_patience.

Every rank has to reach the same verdict, or one leaves the epoch loop while the others block in the next collective. That holds on its own when the score is rank-invariant, i.e. a torchmetrics compute() or a reduce_mean() result; a rank-0-only score has to be broadcast by the caller instead.

Parameters:

improved (bool) – What save_if_best() returned for this epoch.

Return type:

bool

Returns:

True once patience is exhausted, i.e. training should stop.

get_param_groups(*modules)

Split parameters into a weight-decayed and an undecayed optimizer group.

A parameter skips weight decay if it is 1-D (biases, norm and other scale vectors) or if its name contains one of cfg.no_weight_decay. The default list covers biases, norms and every *_emb lookup table; a model overrides it only when a name lies about what it is.

Parameters:

modules (Module) – Modules to take parameters from, self.model when none are given. Pass more when parameters outside the model are trained by the same optimizer, such as an objective with its own projection head.

Return type:

list[dict]

log_lr()

Log the current learning rate for each optimizer parameter group.

clip_and_log_grad_norm(*modules)

Clip gradients in place and log the pre-clip norm as train/grad_norm.

The norm is measured even when grad_clip is unset so the metric does not vanish on unclipped runs. Call once per optimizer step, after backward().

Parameters:

modules (Module) – Modules to clip over, self.model by default. Pass the same ones as get_param_groups(), so the norm spans every parameter stepped.

Return type:

Tensor

log_epoch_grad_norm()

Log the epoch-mean grad norm, the counterpart to the epoch-mean train loss.

log_param_grad_stats(*modules)

Log per-parameter weight norm, grad norm, and grad-to-weight ratio.

Each gate costs one W&B key per parameter, so both are throttled to every log_stats_every_n_steps steps. Call after backward(), before any clip.

push_logs()

Flush accumulated logs to W&B, including current epoch and step.

get_best_pbar_metrics()

Return the best validation metrics formatted for the epoch progress bar.

Returns:

Best validation metrics with the best/val/ prefix stripped,

or an empty dict if no validation has been run yet.

Return type:

dict