training

training

ParameterGroupConfig dataclass

ParameterGroupConfig(name, lr, weight_decay=None, params_pattern=None)

Configuration for a parameter group with specific learning rate.

Attributes:

Name	Type	Description
name	str	e.g., "backbone", "encoder", "decoder", "router".
lr	float	Learning rate.
weight_decay	float | None	Override global weight decay.
params_pattern	str | None	Pattern to match parameter names.

__post_init__

__post_init__()

Reject names owned by the optimizer grouping runtime.

Source code in src/versatil/configs/training.py

def __post_init__(self) -> None:
    """Reject names owned by the optimizer grouping runtime."""
    if self.name == OPTIMIZER_UNMATCHED_GROUPS_NAME:
        raise ValueError(
            f"'{OPTIMIZER_UNMATCHED_GROUPS_NAME}' is reserved for unmatched "
            "parameters."
        )

TrainingStageConfig dataclass

TrainingStageConfig(_target_='versatil.training.stage.TrainingStage', name=MISSING, start_epoch=MISSING, end_epoch=None, trainable_groups=list(), frozen_groups=list(), group_lrs=dict(), group_weight_decays=dict(), loss_weights=dict(), eval_frozen_modules=True)

Hydra schema for one declarative multi-stage training snapshot.

training.stages is ordered by start_epoch and interpreted as a sequence of deltas layered on top of the base training config. A stage may independently override parameter trainability, optimizer hyperparameters, and loss weights, while any omitted field falls back to the cached base regime.

loss_weights is a nested patch that must match the structure exposed by policy.loss_module.weights. Cross-object checks such as stage ordering, optimizer-group existence, and loss-weight path validation run later in versatil.validation.validate_experiment once the full policy and optimizer layout exist. The instantiated runtime object validates only self-contained invariants.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
name	str	Human-readable name.
start_epoch	int	First epoch of the stage.
end_epoch	int | None	Epoch the stage ends before, or null for the rest of training.
trainable_groups	list[str]	Parameter group names unfrozen during the stage.
frozen_groups	list[str]	Parameter group names frozen during the stage.
group_lrs	dict[str, float]	Learning-rate override per parameter group.
group_weight_decays	dict[str, float]	Weight-decay override per parameter group.
loss_weights	dict[str, Any]	Loss-weight override per loss module during the stage.
eval_frozen_modules	bool	Whether frozen modules run in eval mode during the stage.

OptimizerConfig dataclass

OptimizerConfig(target_class=MISSING, lr=0.0001, param_groups=list())

Base optimizer configuration.

Attributes:

Name	Type	Description
target_class	str	Torch optimizer class instantiated for training.
lr	float	Base learning rate (required by all optimizers).
param_groups	list[ParameterGroupConfig]	Parameter groups with different learning rates.

__post_init__

__post_init__()

Validate named optimizer groups used by training stages.

Source code in src/versatil/configs/training.py

def __post_init__(self) -> None:
    """Validate named optimizer groups used by training stages."""
    group_names = [group.name for group in self.param_groups]
    duplicates = _duplicates(group_names)
    if duplicates:
        raise ValueError(
            f"Optimizer parameter group names must be unique: {duplicates}."
        )
    if OPTIMIZER_UNMATCHED_GROUPS_NAME in group_names:
        raise ValueError(
            f"'{OPTIMIZER_UNMATCHED_GROUPS_NAME}' is reserved for unmatched "
            "parameters."
        )

AdamWConfig dataclass

AdamWConfig(target_class='torch.optim.AdamW', lr=0.0001, param_groups=list(), weight_decay=0.0001, betas=(0.9, 0.999), eps=1e-08, amsgrad=False)

Bases: OptimizerConfig

Configuration for torch.optim.AdamW optimizer.

Attributes:

Name	Type	Description
target_class	str	Torch optimizer class instantiated for training.
lr	float	Learning rate.
weight_decay	float	L2 weight decay coefficient.
betas	tuple[float, float]	Adam beta coefficients.
eps	float	Numerical stability epsilon.
amsgrad	bool	Whether the AMSGrad variant is used.

AdamConfig dataclass

AdamConfig(target_class='torch.optim.Adam', lr=0.0001, param_groups=list(), betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, amsgrad=False)

Bases: OptimizerConfig

Configuration for torch.optim.Adam optimizer.

Attributes:

Name	Type	Description
target_class	str	Torch optimizer class instantiated for training.
lr	float	Learning rate.
betas	tuple[float, float]	Adam beta coefficients.
eps	float	Numerical stability epsilon.
weight_decay	float	L2 weight decay coefficient.
amsgrad	bool	Whether the AMSGrad variant is used.

SGDConfig dataclass

SGDConfig(target_class='torch.optim.SGD', lr=0.01, param_groups=list(), momentum=0.0, weight_decay=0.0, dampening=0.0, nesterov=False)

Bases: OptimizerConfig

Configuration for torch.optim.SGD optimizer.

Attributes:

Name	Type	Description
target_class	str	Torch optimizer class instantiated for training.
lr	float	Learning rate.
momentum	float	Momentum factor.
weight_decay	float	L2 weight decay coefficient.
dampening	float	Dampening applied to the momentum.
nesterov	bool	Whether Nesterov momentum is used.

TrainingConfig dataclass

TrainingConfig(num_epochs=100, gradient_accumulate_every=1, optimizer=AdamWConfig(), clip_gradient_norm=False, clip_max_norm=0.1, lr_schedule=None, lr_warmup_steps=5000, lr_scheduler_kwargs=dict(), use_ema=True, ema_power=0.75, swa_lrs=None, swa_epoch_start=0.8, swa_annealing_epochs=10, compile=False, compile_mode=value, tune_lr=False, early_stopping_patience=10, reduce_lr_on_plateau=False, reduce_lr_patience=10, reduce_lr_cooldown=10, stages=list())

Training hyperparameters.

The optional stages list enables declarative multi-stage training. Each stage is applied as a delta over the init-time base regime cached by TrainingStageCallback. Epochs that belong to no stage explicitly fall back to that base regime.

Attributes:

Name	Type	Description
num_epochs	int	Total training epochs.
gradient_accumulate_every	int	Batches accumulated per optimizer step.
optimizer	OptimizerConfig	Optimizer (defaults to AdamW).
clip_gradient_norm	bool	Gradient clipping.
clip_max_norm	float	Gradient-norm clipping threshold, or null to disable.
lr_schedule	str | None	Learning rate schedule name accepted by transformers.get_scheduler, or null for a constant rate. See https://huggingface.co/docs/transformers/main_classes/optimizer_schedules#transformers.get_scheduler
lr_warmup_steps	int	Optimizer steps of linear learning-rate warmup.
lr_scheduler_kwargs	dict[str, float]	Extra keyword arguments for the scheduler.
use_ema	bool	Exponential Moving Average (EMA) of model parameters.
ema_power	float	Decay power of the exponential moving average.
swa_lrs	float | None	If not None, enables SWA with this learning rate.
swa_epoch_start	float	Start SWA at this fraction of total epochs (default: 80% through training).
swa_annealing_epochs	int	Epochs annealing into the SWA learning rate.
compile	bool	Whether the policy is compiled with torch.compile.
compile_mode	str	torch.compile mode.
tune_lr	bool	Whether the Lightning tuner searches a learning rate before training.
early_stopping_patience	int | None	Epochs without val improvement before stopping, or null to disable.
reduce_lr_on_plateau	bool	If True, reduce LR when val_loss plateaus.
reduce_lr_patience	int	Epochs without improvement before the plateau scheduler reduces the LR.
reduce_lr_cooldown	int	Number of epochs to wait after LR reduction before resuming normal operation.
stages	list[TrainingStageConfig]	Ordered training stage "delta" regimes applied on top of the base training regime.

__post_init__

__post_init__()

Validate training knobs that are incompatible with staged control.

Source code in src/versatil/configs/training.py

def __post_init__(self) -> None:
    """Validate training knobs that are incompatible with staged control."""
    if self.stages and self.reduce_lr_on_plateau:
        raise ValueError("training.stages does not support reduce_lr_on_plateau.")
    if self.lr_schedule is not None and self.reduce_lr_on_plateau:
        raise ValueError(
            "reduce_lr_on_plateau cannot be combined with lr_schedule: the "
            "per-step LR scheduler recomputes group learning rates every "
            "step, silently undoing plateau reductions."
        )