Skip to content

training_stage

training_stage

Callback for epoch-based training-stage transitions.

TrainingStageCallback 

TrainingStageCallback(stages, *, learning_rate_schedule_active=False)

Bases: Callback

Apply ordered training stages on top of the cached base regime.

The callback assumes stage ordering, optimizer-group references, and loss_weights paths were already validated in versatil.validation. At runtime it only:

  • caches the base optimizer / trainability / loss state once
  • picks the stage active for trainer.current_epoch
  • applies that stage as a delta over the cached base regime
  • restores the base regime during gaps between stages

Loss overrides are routed purely through the generic loss_module.weights / set_weights / update_weights API. The callback does not know about concrete loss classes or composite internals.

When group_lrs is used together with a scheduler, staged learning rates are interpreted as new scheduler base rates. The current scheduler factor is preserved; the callback does not reset scheduler progress.

Initialize the training stage callback.

Parameters:

Name Type Description Default
stages list[TrainingStage]

Ordered runtime stages from training.stages.

 required
learning_rate_schedule_active bool

Whether the Lightning module configured a scheduler. Used to fail fast when staged learning rates cannot update scheduler base rates.

 False
Source code in src/versatil/training/callbacks/training_stage.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def __init__(
    self,
    stages: list[TrainingStage],
    *,
    learning_rate_schedule_active: bool = False,
) -> None:
    """Initialize the training stage callback.

    Args:
        stages: Ordered runtime stages from ``training.stages``.
        learning_rate_schedule_active: Whether the Lightning module configured
            a scheduler. Used to fail fast when staged learning rates cannot
            update scheduler base rates.
    """
    super().__init__()
    if not stages:
        raise ValueError("TrainingStageCallback requires a non-empty stage list.")
    self.stages = stages
    self.learning_rate_schedule_active = learning_rate_schedule_active
    self._uses_group_learning_rates = any(stage.group_lrs for stage in self.stages)
    self._last_applied_stage_key: str | None = None
    self._validated = False
    self._base_group_learning_rates: dict[str, float] = {}
    self._base_group_weight_decays: dict[str, float] = {}
    self._base_loss_weights: WeightsDictionary | None = None
    self._parameter_group_by_id: dict[int, str] = {}
    self._base_parameter_trainability: dict[int, bool] = {}
    self._frozen_modules: tuple[torch.nn.Module, ...] = ()
    self._trainable_modules: tuple[torch.nn.Module, ...] = ()

on_train_start 

on_train_start(trainer, pl_module)

Cache the base regime and apply the stage active at resume epoch.

Source code in src/versatil/training/callbacks/training_stage.py 
72
73
74
75
76
77
78
def on_train_start(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Cache the base regime and apply the stage active at resume epoch."""
    lightning_policy = self._require_lightning_policy(pl_module)
    self._ensure_initialized(trainer=trainer, pl_module=lightning_policy)
    self._apply_active_stage(trainer=trainer, pl_module=lightning_policy)

on_train_epoch_start 

on_train_epoch_start(trainer, pl_module)

Apply a newly active stage, or restore modes for the current stage.

Source code in src/versatil/training/callbacks/training_stage.py 
80
81
82
83
84
85
86
def on_train_epoch_start(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Apply a newly active stage, or restore modes for the current stage."""
    lightning_policy = self._require_lightning_policy(pl_module)
    self._ensure_initialized(trainer=trainer, pl_module=lightning_policy)
    self._apply_active_stage(trainer=trainer, pl_module=lightning_policy)