nvalchemi.training.hooks.EMAHook

class nvalchemi.training.hooks.EMAHook(*, model_key='main', decay=0.999, update_every=1, start_step=0, use_buffers=True, num_updates=0)

Hook maintaining an exponential moving average of a training model.

Runs through TrainingUpdateOrchestrator and updates at TrainingStage.AFTER_OPTIMIZER_STEP. It lazily builds a AveragedModel wrapped around ctx.models[model_key] on the first eligible step, and updates it via get_ema_multi_avg_fn() — no manual parameter arithmetic. The hook is a pure observer: it never calls backward(), touches gradients, drives any optimizer / scheduler / GradScaler, or mutates ctx.models. If an earlier update hook vetoes TrainingStage.DO_OPTIMIZER_STEP, the orchestrator passes will_skip=True and EMA does not update on that batch.

Access the averaged wrapper via get_averaged_model(), which raises a RuntimeError if no eligible step has yet triggered lazy initialization. A device/dtype field is omitted by design; after AveragedModel deep-copies the source, EMAHook aligns each averaged parameter and buffer to the corresponding source tensor’s device and floating-point dtype. This keeps generated or monkey-patched modules whose deepcopy/load path materializes registered tensors on CPU or in a default dtype usable without model-specific hooks.

Note

If the copied module defines modify_ema_methods(), the hook calls it once immediately after constructing the averaged model. Model wrappers can use this optional interface to restore runtime methods discarded by deepcopy.

Parameters:

• model_key (str, optional) – Key identifying the source model inside ctx.models. Default "main".

• decay (float, optional) – EMA decay factor in [0.0, 1.0). Default 0.999.

• update_every (int, optional) – Positive step stride for averaging updates. Default 1.

• start_step (int, optional) – Non-negative minimum completed step before updates begin. Default 0.

• use_buffers (bool, optional) – Forwarded to AveragedModel; when True also averages module buffers. Default True.

• num_updates (int, optional) – Non-negative count of EMA updates already performed. Settable at construction so a checkpoint can restore the update counter; normally left at its default and advanced internally as updates occur. Default 0.

Raises:

• pydantic.ValidationError – If any field violates its declared bounds or an unknown kwarg is passed.

• KeyError – On first eligible call, if model_key is missing from ctx.models.

• RuntimeError – From get_averaged_model() when called before lazy init.

See also

torch.optim.swa_utils.AveragedModel

Underlying averaging wrapper.

torch.optim.swa_utils.get_ema_multi_avg_fn

Factory for the EMA averaging function.

Examples

Checkpoint recipe for inference / eval reload of the EMA-averaged weights. Save hook.get_averaged_model().module alongside the base model and rebuild the AveragedModel wrapper after loading, because create_model_spec() only reconstructs plain Module objects:

>>> from torch import nn
>>> from torch.optim.swa_utils import AveragedModel
>>> from nvalchemi.training import (
...     EMAHook, create_model_spec, load_checkpoint, save_checkpoint,
... )
>>> base = nn.Linear(4, 2)
>>> hook = EMAHook(model_key="main", decay=0.99)
>>> # ... training loop drives `hook` via TrainingStrategy ...
>>> spec = create_model_spec(nn.Linear, in_features=4, out_features=2)
>>> save_checkpoint(
...     "ckpt/",
...     models={
...         "main": (base, spec),
...         "main_ema": (hook.get_averaged_model().module, spec),
...     },
... )
>>> loaded = load_checkpoint("ckpt/")
>>> reconstructed_ema = AveragedModel(loaded.models["main_ema"][0])

To resume training with EMA continuing from a checkpoint, use state_dict() / load_state_dict(), which round-trip num_updates and the averaged weights into a freshly constructed hook.

Notes

The default deepcopy-based construction does not support fully_shard (FSDP2) / DTensor models; override _build_averaged_model() to supply a pre-built sharded copy.

get_averaged_model()

Return the AveragedModel wrapper or raise if uninitialized.

Raises:

RuntimeError – If neither setup nor an eligible training step has initialized EMA.

Return type:

AveragedModel

load_state_dict(state)

Restore hook counters and averaged weights from a prior snapshot.

Parameters:

state (Mapping[str, Any]) – Mapping produced by state_dict(). Missing config keys and num_updates are ignored. Missing averaged_model_state clears any prior live or pending averaged state. Any present config key must equal the corresponding constructor field.

Raises:

ValueError – If a config field in state differs from this hook’s current field.

Return type:

None

Notes

Before lazy init, averaged_model_state is stashed and applied during _ensure_initialized(). Clearing on absence prevents stale averaged state from surviving a config-only reload. Checkpoint loaders may still choose a map_location, but EMAHook reapplies per-tensor device and floating-point dtype placement after loading averaged state so registered tensors remain usable for validation.

model_config = {'arbitrary_types_allowed': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context, /)

This function is meant to behave like a BaseModel method to initialise private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

• self (BaseModel) – The BaseModel instance.

• context (Any) – The context.

Return type:

None

state_dict()

Return a serializable snapshot of hook state.

Returns:

Contains the config fields, num_updates, and — if available — averaged_model_state sourced from the live AveragedModel or, before lazy init, from any stashed pending state. No device key is emitted.

Return type:

dict[str, Any]