ema

ema

Exponential moving average callback.

EMACallback

EMACallback(power=0.75, update_after_step=0, inv_gamma=1.0, min_value=0.0, max_value=0.9999)

Bases: Callback

Exponential Moving Average callback for model weights.

Maintains a moving average of model weights during training. The EMA model is used for validation and can provide more stable predictions.

Based on @crowsonkb's notes on EMA Warmup: If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps), gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at 215.4k steps).

Initialize EMA callback.

Parameters:

Name	Type	Description	Default
power	float	Exponential factor of EMA warmup (default: 0.75 for shorter training)	0.75
update_after_step	int	Start EMA updates after this many steps	0
inv_gamma	float	Inverse multiplicative factor of EMA warmup	1.0
min_value	float	Minimum EMA decay rate	0.0
max_value	float	Maximum EMA decay rate	0.9999

Source code in src/versatil/training/callbacks/ema.py

def __init__(
    self,
    power: float = 0.75,
    update_after_step: int = 0,
    inv_gamma: float = 1.0,
    min_value: float = 0.0,
    max_value: float = 0.9999,
):
    """Initialize EMA callback.

    Args:
        power: Exponential factor of EMA warmup (default: 0.75 for shorter training)
        update_after_step: Start EMA updates after this many steps
        inv_gamma: Inverse multiplicative factor of EMA warmup
        min_value: Minimum EMA decay rate
        max_value: Maximum EMA decay rate
    """
    super().__init__()
    self.power = power
    self.update_after_step = update_after_step
    self.inv_gamma = inv_gamma
    self.min_value = min_value
    self.max_value = max_value
    self.decay = 0.0
    self.ema_model: torch.nn.Module | None = None
    self._resume_ema_state: dict[str, torch.Tensor] | None = None

on_fit_start

on_fit_start(trainer, pl_module)

Create EMA model copy at start of training.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module (LightningPolicy)	required

Source code in src/versatil/training/callbacks/ema.py

def on_fit_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
    """Create EMA model copy at start of training.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module (LightningPolicy)
    """
    self.ema_model = copy.deepcopy(pl_module.policy)
    if self._resume_ema_state is not None:
        self.ema_model.load_state_dict(self._resume_ema_state, strict=False)
        self._resume_ema_state = None
    self.ema_model.eval()
    self.ema_model.requires_grad_(False)

on_train_batch_end

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)

Update EMA model after each training batch.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required
outputs	Tensor | dict[str, Any] | None	Training step outputs	required
batch	Any	Current batch	required
batch_idx	int	Batch index	required

Source code in src/versatil/training/callbacks/ema.py

def on_train_batch_end(
    self,
    trainer: pl.Trainer,
    pl_module: pl.LightningModule,
    outputs: torch.Tensor | dict[str, Any] | None,
    batch: Any,
    batch_idx: int,
) -> None:
    """Update EMA model after each training batch.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
        outputs: Training step outputs
        batch: Current batch
        batch_idx: Batch index
    """
    if self.ema_model is None:
        return
    window_incomplete = (batch_idx + 1) % trainer.accumulate_grad_batches != 0
    if window_incomplete and not trainer.is_last_batch:
        # Track optimizer steps, not micro-batches: updating on every
        # accumulation micro-batch would decay the average N times faster
        # than the configured warmup schedule. Lightning also flushes a
        # partial accumulation window on the epoch's last batch, so that
        # optimizer step must enter the average too.
        return

    self.decay = self._get_decay(trainer.global_step)

    with torch.no_grad():
        for module, ema_module in zip(
            pl_module.policy.modules(), self.ema_model.modules(), strict=True
        ):
            for buffer, ema_buffer in zip(
                module.buffers(recurse=False),
                ema_module.buffers(recurse=False),
                strict=True,
            ):
                ema_buffer.copy_(
                    buffer.to(device=ema_buffer.device, dtype=ema_buffer.dtype)
                )

            for param, ema_param in zip(
                module.parameters(recurse=False),
                ema_module.parameters(recurse=False),
                strict=True,
            ):
                if isinstance(param, dict):
                    raise RuntimeError("Dict parameter not supported")

                if isinstance(module, _BatchNorm) or not param.requires_grad:
                    ema_param.copy_(param.to(dtype=ema_param.dtype).data)
                else:
                    ema_param.mul_(self.decay)
                    ema_param.add_(
                        param.data.to(dtype=ema_param.dtype), alpha=1 - self.decay
                    )

    if trainer.global_step % 100 == 0:
        pl_module.log("ema_decay", self.decay, on_step=True, on_epoch=False)

on_save_checkpoint

on_save_checkpoint(trainer, pl_module, checkpoint)

Inject EMA weights into the checkpoint, stashing the raw weights.

The checkpoint's state_dict intentionally holds the EMA weights so every saved checkpoint is deployment-ready. The raw fast weights are kept under a separate key so resumed training does not restart from the average with optimizer moments that belong to the raw weights.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required
checkpoint	dict	Checkpoint dictionary being saved	required

Source code in src/versatil/training/callbacks/ema.py

def on_save_checkpoint(
    self,
    trainer: pl.Trainer,
    pl_module: pl.LightningModule,
    checkpoint: dict,
) -> None:
    """Inject EMA weights into the checkpoint, stashing the raw weights.

    The checkpoint's ``state_dict`` intentionally holds the EMA weights so
    every saved checkpoint is deployment-ready. The raw fast weights are
    kept under a separate key so resumed training does not restart from
    the average with optimizer moments that belong to the raw weights.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
        checkpoint: Checkpoint dictionary being saved
    """
    if self.ema_model is None:
        return
    raw_policy_state: dict[str, torch.Tensor] = {}
    ema_state = self.ema_model.state_dict()
    for key, value in ema_state.items():
        ckpt_key = f"policy.{key}"
        if ckpt_key in checkpoint["state_dict"]:
            raw_policy_state[ckpt_key] = checkpoint["state_dict"][ckpt_key].clone()
            checkpoint["state_dict"][ckpt_key] = value.clone()
    checkpoint[CheckpointKey.RAW_POLICY_STATE_DICT.value] = raw_policy_state

on_load_checkpoint

on_load_checkpoint(trainer, pl_module, checkpoint)

Restore raw training weights and queue the EMA state for rebuild.

Lightning loads state_dict (the EMA weights) into the module before this hook runs, so the raw weights are written back onto the policy here and the EMA weights are stashed until on_fit_start recreates the EMA model.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required
checkpoint	dict	Loaded checkpoint dictionary	required

Source code in src/versatil/training/callbacks/ema.py

def on_load_checkpoint(
    self,
    trainer: pl.Trainer,
    pl_module: pl.LightningModule,
    checkpoint: dict,
) -> None:
    """Restore raw training weights and queue the EMA state for rebuild.

    Lightning loads ``state_dict`` (the EMA weights) into the module
    before this hook runs, so the raw weights are written back onto the
    policy here and the EMA weights are stashed until ``on_fit_start``
    recreates the EMA model.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
        checkpoint: Loaded checkpoint dictionary
    """
    raw_policy_state = checkpoint.get(CheckpointKey.RAW_POLICY_STATE_DICT.value)
    if not raw_policy_state:
        return
    policy_prefix = "policy."
    self._resume_ema_state = {
        key.removeprefix(policy_prefix): checkpoint["state_dict"][key].clone()
        for key in raw_policy_state
        if key in checkpoint["state_dict"]
    }
    pl_module.policy.load_state_dict(
        {
            key.removeprefix(policy_prefix): value
            for key, value in raw_policy_state.items()
        },
        strict=False,
    )

on_validation_start

on_validation_start(trainer, pl_module)

Temporarily replace policy with EMA model for validation.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/ema.py

def on_validation_start(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Temporarily replace policy with EMA model for validation.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if self.ema_model is not None:
        self._original_policy = pl_module.policy
        self._sync_policy_runtime_state(
            source_policy=pl_module.policy,
            target_policy=self.ema_model,
        )
        pl_module.policy = self.ema_model

on_validation_end

on_validation_end(trainer, pl_module)

Restore original policy after validation.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/ema.py

def on_validation_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Restore original policy after validation.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if hasattr(self, "_original_policy"):
        pl_module.policy = self._original_policy
        delattr(self, "_original_policy")