nvalchemi.training.hooks.TrainingUpdateHook

class nvalchemi.training.hooks.TrainingUpdateHook

Base class for hooks that customize training-update phases.

Subclasses override __call__() and dispatch on stage to handle one or more claimed stages: BEFORE_BATCH, DO_BACKWARD, DO_OPTIMIZER_STEP, AFTER_OPTIMIZER_STEP. Compose via + to build a TrainingUpdateOrchestrator. See Training update hooks for the stage contract and restrictions each update hook must follow.

priority

Dispatch order within an orchestrator; lower runs first. Canonical buckets: 10 = gradient accumulation, 20 = mixed precision, 30 = gradient clipping, 40 = spike skipping. Default 50.

Type:

int

_exclusive_update_key

Optional key for hook families that must appear at most once inside an orchestrator.

Type:

str | None

Notes

TrainingUpdateHook is NOT directly compatible with the standard Hook Protocol – its __call__ signature includes a will_skip argument and returns (bool, torch.Tensor | None) rather than the Protocol’s __call__(ctx, stage) -> None. This is intentional: Hook is a structural Protocol so domain-specific hook families can use signatures suited to their semantics. Bare instances must be composed via + or wrapped by a TrainingUpdateOrchestrator (the strategy auto-wraps lone hooks); the orchestrator owns Protocol compliance.

will_skip is a stage-local cumulative veto signal. It is True when an earlier, higher-priority hook has already requested that the current stage’s gated operation be skipped. The orchestrator still calls later hooks after a veto so they can observe the decision, update bookkeeping, or emit diagnostics, but those hooks should avoid side effects that assume the gated operation will run. A hook may also return False to veto the operation for lower-priority hooks.

This signal is intended for composable pipeline behavior. For example, a gradient-accumulation hook can veto DO_OPTIMIZER_STEP on non-step microbatches; later hooks then receive will_skip=True and can skip work such as gradient clipping, scaler updates, or expensive parameter scans. will_skip is reset for each stage dispatch and should not be interpreted as a global training-step status unless the orchestrator also records that state on ctx.

Each __call__ returns (proceed, loss):

• proceed is a strict bool (int/None raise TypeError). On BEFORE_BATCH and DO_OPTIMIZER_STEP the orchestrator applies any-veto-wins composition: if any hook returns False the gated operation (zero_gradients or optimizer/scheduler.step) is skipped. On DO_BACKWARD and AFTER_OPTIMIZER_STEP the value is unused; return True.

• loss is the loss tensor the hook would use, transformed or not. Default is ctx.loss unchanged. The orchestrator threads it through hooks in priority order during DO_BACKWARD so each hook sees its predecessor’s transform; backward() runs once on the final loss. Hooks that run on stages other than DO_BACKWARD may return None for loss because the orchestrator ignores it there.

Examples

>>> import torch
>>> from nvalchemi.training._stages import TrainingStage
>>> class ClipGrads(TrainingUpdateHook):
...     priority = 30
...     def __init__(self, max_norm):
...         self.max_norm = max_norm
...     def __call__(self, ctx, stage, will_skip):
...         match stage:
...             case TrainingStage.DO_OPTIMIZER_STEP:
...                 if not will_skip:
...                     for opt in ctx.optimizers:
...                         params = (p for g in opt.param_groups for p in g["params"])
...                         torch.nn.utils.clip_grad_norm_(params, self.max_norm)
...                 return True, ctx.loss
...             case _:
...                 return True, ctx.loss