Training update hooks#

Training update hooks are for policies that need to participate in the weight-update portion of a training batch. They are intentionally narrower than general Hook objects: a TrainingUpdateHook only runs on the stages owned by TrainingUpdateOrchestrator, and the orchestrator performs the actual backward(), optimizer step, scheduler step, and gradient zeroing calls.

Use this hook family when multiple update policies need to coordinate around the same batch update. Typical examples include gradient accumulation, mixed precision, gradient clipping, spike skipping, and post-step model averaging. Use a standard training hook for read-only observation or lifecycle logic that does not need to own backward or optimizer-step behavior.

Most users register concrete TrainingUpdateHook instances, such as MixedPrecisionHook or EMAHook, directly on the strategy. The TrainingUpdateOrchestrator is the coordination object that runs those update hooks in one ordered update path; strategies create it automatically for bare update hooks. Construct it yourself only when you need to pre-compose update hooks before passing them to a strategy.

Fine-tuning helpers such as ModulePatchHook and TrainableParameterHook are registration-time hooks. They adapt the model tree and optimizer parameter set before training starts, but they do not own any batch update stages. See Fine-Tuning Pretrained Models and Fine-tuning API for those workflows.

ctx.step_count tracks completed optimizer/scheduler steps on this worker, and ctx.global_step_count tracks completed optimizer/scheduler steps across all data-parallel workers. If an update hook vetoes DO_OPTIMIZER_STEP for gradient accumulation or spike skipping, the batch still advances ctx.batch_count and ctx.epoch_step_count but does not advance either step counter.

TrainingStage#

TrainingStage enumerates the eighteen hook-firing points within a training run:

TrainingStage hook firing points across a training run.#

Training stages reference#
Stage	Value	When it fires
`SETUP`	1	Once before optimizer construction; setup hooks may mutate model wrappers and dataloaders before training begins.
`BEFORE_TRAINING`	2	Once before the epoch loop, after the model is on device and optimizers are constructed.
`BEFORE_EPOCH`	3	Start of each epoch, before the first batch.
`BEFORE_BATCH`	4	Start of each batch, before gradient zeroing.
`BEFORE_FORWARD`	5	Before the model forward pass; right stage for input transforms such as neighbor-list construction.
`AFTER_FORWARD`	6	After the model forward pass; predictions are available.
`BEFORE_LOSS`	7	Before the loss computation.
`AFTER_LOSS`	8	After the loss computation; the loss tensor is populated.
`BEFORE_BACKWARD`	9	Before the backward pass; typical slot for observers that need the pre-backward state.
`DO_BACKWARD`	10	Replacement slot for the backward pass. At most one hook may claim this stage; observers should use `BEFORE_BACKWARD` / `AFTER_BACKWARD`.
`AFTER_BACKWARD`	11	After the backward pass; gradients are available. Typical slot for gradient clipping and gradient-norm logging.
`BEFORE_OPTIMIZER_STEP`	12	Immediately before the optimizer step; last pre-step point for observers that need unscaled gradients.
`DO_OPTIMIZER_STEP`	13	Replacement slot for the optimizer and LR-scheduler step. At most one hook may claim this stage; observers should use `BEFORE_OPTIMIZER_STEP` / `AFTER_OPTIMIZER_STEP`.
`AFTER_OPTIMIZER_STEP`	14	After the optimizer and scheduler step path completes; typical slot for EMA updates and post-step logging.
`AFTER_BATCH`	15	End of each batch; generic batch cleanup.
`AFTER_EPOCH`	16	End of each epoch, after the last batch.
`AFTER_TRAINING`	17	Once after the final epoch.
`AFTER_VALIDATION`	18	Event-based; fires inside `TrainingStrategy.validate()` after a validation pass produces its summary. Reliable slot for loggers that need the latest validation results.

Distributed data parallel#

DDPHook wraps optimized models in torch.nn.parallel.DistributedDataParallel during TrainingStage.SETUP. This setup stage runs after distributed rank/device resolution and before optimizer construction, so optimizers are built from the DDP-wrapped model parameters. See Distributed Training for the workflow-level DistributedManager guide.

from nvalchemi.distributed import DistributedManager
from nvalchemi.training.hooks import DDPHook, MixedPrecisionHook
from nvalchemi.training.strategy import TrainingStrategy

DistributedManager.initialize()
manager = DistributedManager()

strategy = TrainingStrategy(
    ...,
    distributed_manager=manager,
    hooks=[
        DDPHook(find_unused_parameters=False),
        MixedPrecisionHook(precision="bf16"),
    ],
)

Launch single-node distributed training with torchrun:

torchrun --nproc_per_node=2 train.py

DDPHook uses TrainingStrategy.distributed_manager when one is provided, falling back to torch.distributed and torchrun environment variables. Sampler injection is automatic: a DistributedSampler is added when the dataloader does not already have one, and sampler.set_epoch() is called each epoch when available.

DDPHook is not a training-update hook, so it does not participate in DO_BACKWARD or DO_OPTIMIZER_STEP. Register it alongside MixedPrecisionHook normally; DDP wrapping happens before AMP opens its per-batch autocast/update path.

Field	Type	Description
`model_keys`	`tuple[str, ...] \| None`	Named models to wrap. `None` wraps all models that have optimizer configs.
`find_unused_parameters`	`bool \| None`	Forwarded to `DistributedDataParallel`. `None` uses the external manager’s setting when present, otherwise `False`.
`broadcast_buffers`	`bool \| None`	Forwarded to `DistributedDataParallel`. `None` uses the external manager’s setting when present, otherwise `False`.
`static_graph`	`<class 'bool'>`	Forwarded to `DistributedDataParallel`.
`process_group`	`typing.Any \| None`	Explicit process group. Defaults to a process group exposed by the external distributed manager or PyTorch’s default group.
`backend`	`str \| None`	Backend used when this hook initializes `torch.distributed`.
`auto_init`	`<class 'bool'>`	If `True`, initialize `torch.distributed` when `WORLD_SIZE > 1` and no manager/process group has already initialized communication.
`sampler_cls`	`collections.abc.Callable[..., typing.Any]`	Sampler class or factory used for supported dataloaders. The callable is invoked as `sampler_cls(dataset, **sampler_kwargs)`. The default is `torch.utils.data.DistributedSampler`.
`sampler_kwargs`	`dict[str, typing.Any]`	Keyword arguments forwarded to `sampler_cls`. For the default `DistributedSampler` and sampler callables that accept PyTorch’s distributed sampler keywords, missing `num_replicas`, `rank`, `shuffle`, `seed`, and `drop_last` values are inferred from the manager and dataloader before user-provided kwargs are applied.

PyTorch profiler traces#

TorchProfilerHook captures PyTorch profiler Chrome traces through PhysicsNeMo’s profiler wrapper. In training workflows it starts at TrainingStage.BEFORE_TRAINING, advances the profiler schedule at TrainingStage.AFTER_BATCH, and finalizes at TrainingStage.AFTER_TRAINING or when the strategy context exits. Standalone train_batch() calls start lazily at TrainingStage.BEFORE_BATCH and still finalize when the context closes.

from torch.profiler import ProfilerActivity, schedule

from nvalchemi.training.hooks import TorchProfilerHook
from nvalchemi.training.strategy import TrainingStrategy

profile_hook = TorchProfilerHook(
    output_dir="profiles/train-run",
    activities=(ProfilerActivity.CPU, ProfilerActivity.CUDA),
    schedule=schedule(wait=2, warmup=2, active=5, repeat=1),
    record_shapes=True,
    profile_memory=True,
    with_flops=True,
)

strategy = TrainingStrategy(..., hooks=[profile_hook])
strategy.run(train_loader)

Each process writes to profiles/train-run/rank_<global_rank>/torch/ unless PhysicsNeMo’s distributed manager is active and already owns rank suffixing.

Mixed precision#

MixedPrecisionHook enables torch.amp.autocast for the forward/loss portion of the batch and uses torch.amp.GradScaler when precision is torch.float16. The precision argument is required so configs must choose one of the supported policies explicitly:

import torch

from nvalchemi.training.hooks import MixedPrecisionHook
from nvalchemi.training.strategy import TrainingStrategy

strategy = TrainingStrategy(
    ...,
    hooks=[MixedPrecisionHook(precision=torch.bfloat16)],
)

precision accepts the dtype objects torch.float32, torch.bfloat16, and torch.float16, the canonical strings "float32", "bfloat16", and "float16", or the shorthand aliases "fp32", "bf16", and "fp16".

The policies are:

torch.float32: no autocast context is created and no scaler is used.
torch.bfloat16: eligible ops run under bf16 autocast and no scaler is used.
torch.float16: eligible forward/loss ops run under fp16 autocast, the hook scales the loss before backward, unscales gradients immediately before an optimizer step proceeds, and lets the scaler skip steps with inf or nan gradients.

Register at most one MixedPrecisionHook per strategy. The strategy rejects multiple mixed-precision hooks so that autocast, loss scaling, unscale, scaler step, and scaler update cannot be applied twice in one batch update.

Autocast scope#

The autocast context spans BEFORE_BATCH through DO_BACKWARD and is released before backward(). Model components or loss functions that require full precision for a specific subregion should use a local torch.amp.autocast(..., enabled=False) block within that region.

Gradient accumulation#

Veto DO_OPTIMIZER_STEP on intermediate microbatches to accumulate gradients across a window of K batches. Under torch.float16, MixedPrecisionHook suppresses GradScaler.unscale, GradScaler.step, and GradScaler.update for vetoed batches; gradients remain scaled until the veto is lifted on the Kth batch. Schedulers advance only when the paired optimizer actually stepped.

Zero-gradient policy is set by the BEFORE_BATCH return value: return proceed=False on intermediate microbatches to skip zeroing and accumulate into existing gradients.

Validation#

TrainingStrategy.validate() honors a registered MixedPrecisionHook automatically; the use_mixed_precision field on ValidationConfig controls whether autocast is active during inference. The standalone ValidationLoop is hook-agnostic and takes an explicit autocast callable instead. See Validation.

Stage constraints#

Training update hooks always receive (ctx, stage, will_skip) and return (proceed, loss). The meaning of those values depends on the stage:

Training update hook stage contract#
Stage	Hook responsibility	Return contract	Restrictions and expectations
`BEFORE_BATCH`	Decide whether the orchestrator should call `zero_gradients()`.	`proceed` must be a strict `bool`. Any `False` vetoes gradient zeroing. `loss` is ignored.	Do not call `backward()`, `optimizer.step()`, or `scheduler.step()`. Use this stage for zero-grad policy, per-batch update bookkeeping, or resetting state that is safe before the forward pass.
`DO_BACKWARD`	Transform or replace `ctx.loss` before the orchestrator calls `backward()` once.	`loss` must be a `torch.Tensor`. `proceed` is ignored.	Do not call `backward()` directly. Return the loss tensor the next update hook should see. This is the stage for loss scaling and other loss-space transforms.
`DO_OPTIMIZER_STEP`	Decide whether the orchestrator should call `step_optimizers()` and `step_lr_schedulers()`.	`proceed` must be a strict `bool`. Any `False` vetoes both the optimizer and scheduler step. `loss` is ignored.	Do not call `backward()`. Avoid side effects that assume a step will run when `will_skip` is `True`. This is the stage for pre-step logic such as gradient clipping, scaler updates, and accumulation/spike-skip decisions.
`AFTER_OPTIMIZER_STEP`	Observe the final step decision and run post-step bookkeeping.	`proceed` and `loss` are ignored. `will_skip` tells the hook whether the optimizer/scheduler step was vetoed.	Do not call `backward()` or perform another optimizer/scheduler step. Use this stage for work that should happen after the step path, such as EMA updates, diagnostics, and state cleanup.

Composition rules#

All update hooks for a strategy are composed into one orchestrator. Lower priority values run first, and registration order breaks ties. The orchestrator keeps calling later hooks after a veto so they can observe will_skip=True and update their own state consistently.

Only one object may own DO_BACKWARD or DO_OPTIMIZER_STEP in a TrainingStrategy. For convenience, the strategy auto-wraps bare TrainingUpdateHook instances into one TrainingUpdateOrchestrator. Passing stage=... while registering an update hook is not supported because update hooks declare their stages through the orchestrator.

Checkpointing#

CheckpointHook saves model weights, optimizer state, scheduler state, training counters, and the state of any CheckpointableHook implementors to disk at a configured cadence.

Field	Type	Description
`checkpoint_dir`	`<class 'pathlib._local.Path'>`	Root directory for restartable training checkpoints.
`step_interval`	`int \| None`	Completed-step save interval.
`epoch_interval`	`int \| None`	Completed-epoch save interval.
`async_save`	`<class 'bool'>`	Write checkpoint snapshots on a background thread.
`rank_zero_only`	`<class 'bool'>`	Restrict checkpoint writes to distributed rank 0.
`last_checkpoint_index`	`int \| None`	Most recent checkpoint index known to have been written. In async mode, this updates when the background future completes.

Pass the same hook instance to TrainingStrategy.load_checkpoint() when resuming so its internal state is restored alongside the model and optimizer:

from nvalchemi.training import CheckpointHook, TrainingStrategy

checkpoint = CheckpointHook(
    "runs/my-model/checkpoints",
    step_interval=500,
)

# Resume from step 500
strategy = TrainingStrategy.load_checkpoint(
    "runs/my-model/checkpoints/step_500",
    hooks=[checkpoint],
)
strategy.run(train_loader)

EMA model averaging#

EMAHook maintains an exponential moving average of one model’s weights. The averaged weights are updated at AFTER_OPTIMIZER_STEP only when the optimizer step was not vetoed, so the EMA step count stays in sync with the actual optimizer step count.

Field	Type	Description
`model_key`	`<class 'str'>`	Key identifying the source model in ctx.models.
`decay`	`<class 'float'>`	EMA decay factor in [0.0, 1.0).
`update_every`	`<class 'int'>`	Completed-step interval between EMA updates (global-modulo).
`start_step`	`<class 'int'>`	First completed step eligible for EMA updates.
`use_buffers`	`<class 'bool'>`	If True, also average module buffers (e.g. BN running stats).
`num_updates`	`<class 'int'>`	Number of EMA updates performed; restored from checkpoints.

Access the averaged model wrapper via ema.get_averaged_model() after training.

from nvalchemi.training.hooks import EMAHook
from nvalchemi.training.strategy import TrainingStrategy

ema = EMAHook(model_key="main", decay=0.999)
strategy = TrainingStrategy(..., hooks=[ema])
strategy.run(train_loader)

averaged_model = ema.get_averaged_model()

AveragedModel constructs its module with deepcopy. If a model wrapper defines a callable modify_ema_methods() method, EMAHook calls it once on the copied module immediately after construction and before loading pending EMA checkpoint weights. This optional interface is intended for wrappers whose third-party models install runtime methods that deepcopy does not preserve; ordinary models do not need to implement it.

Restartable update hooks#

If a training hook owns state that changes resumed training behavior — EMA weights, a learned update policy, accumulated statistics — implement CheckpointableHook by adding state_dict() and load_state_dict(). The strategy checkpoint loader restores state only into hooks that satisfy this protocol.

from nvalchemi.hooks import CheckpointableHook
from nvalchemi.training.hooks import TrainingUpdateHook

class StatefulHook(TrainingUpdateHook):
    def __init__(self):
        self.step_count = 0

    def state_dict(self):
        return {"step_count": self.step_count}

    def load_state_dict(self, state):
        self.step_count = int(state["step_count"])

    def __call__(self, ctx, stage, will_skip):
        ...
        return True, ctx.loss

assert isinstance(StatefulHook(), CheckpointableHook)

For Pydantic update hooks, call model_dump() inside state_dict() to capture field values before appending non-field runtime tensors. Tensor state must remain in state_dict(); use model_dump_json() only for configuration records or diagnostics.

API reference#

`DDPHook`	Wrap training models with `DistributedDataParallel` at setup time.
`MixedPrecisionHook`	Automatic-mixed-precision hook driving autocast and `GradScaler`.
`TrainingUpdateHook`	Base class for hooks that customize training-update phases.
`TrainingUpdateOrchestrator`	Composes `TrainingUpdateHook` instances and drives updates.
`EMAHook`	Hook maintaining an exponential moving average of a training model.
`CheckpointHook`	Periodically save restartable training strategy checkpoints.

The general-purpose profiling hooks StageTimingHook and TorchProfilerHook also work with training and are documented in Hooks — Core Framework.