nvalchemi.training.hooks.MixedPrecisionHook

pydantic model nvalchemi.training.hooks.MixedPrecisionHook

Automatic-mixed-precision hook driving autocast and GradScaler.

MixedPrecisionHook is a TrainingUpdateHook. When it is registered directly on TrainingStrategy, the strategy auto-wraps it in a TrainingUpdateOrchestrator. The orchestrator owns backward() and optimizer/scheduler stepping; this hook supplies a scaled loss, exposes ctx.grad_scaler for scaler-aware stepping, and unscales gradients immediately before an optimizer step proceeds so gradient accumulation can keep accumulating scaled gradients.

The first TrainingStage.BEFORE_BATCH lazily constructs the autocast region on the workflow’s primary device (ctx.workflow.devices[0]), so the hook need not know the device at construction time. For fp16, the same path also lazily constructs the torch.amp.GradScaler. The autocast region is released inside TrainingStage.DO_BACKWARD before the orchestrator calls backward(), while the scaler persists across batches. Force and stress predictions produced during the model forward, plus the configured training losses, are therefore inside the autocast region; backward is not.

Precision modes:

• torch.float32 — no autocast context or scaler is created; the hook is a functional no-op aside from participating in the orchestrated update path.

• torch.bfloat16 — autocast casts eligible ops to bfloat16. No gradient scaling because bf16’s exponent range matches fp32.

• torch.float16 — autocast casts eligible ops to float16 during forward and loss computation. The scaler scales the loss before the orchestrator calls backward(), unscales gradients just before optimizer stepping, and skips optimizer steps that would otherwise consume inf/nan gradients.

Raises:

pydantic.ValidationError – If precision is not one of the supported dtypes.

Examples

>>> import torch
>>> from nvalchemi.training.hooks import MixedPrecisionHook
>>> MixedPrecisionHook(precision=torch.bfloat16).precision
torch.bfloat16
>>> MixedPrecisionHook(precision="float16").precision
torch.float16
>>> MixedPrecisionHook(precision="bf16").precision
torch.bfloat16

Notes

• When multiple optimizers are configured, every optimizer in ctx.optimizers is unscaled in list order immediately before stepping. The orchestrator advances each scheduler in ctx.lr_schedulers only when its paired optimizer step was not skipped by the scaler.

• For gradient accumulation, accumulated gradients remain scaled until the effective batch is ready to step. Earlier-priority update hooks can veto TrainingStage.DO_OPTIMIZER_STEP to suppress unscale, scaler step, and scaler update for intermediate accumulation batches.

• A strategy may register only one MixedPrecisionHook. Multiple instances are rejected to prevent duplicated autocast/scaler operations.

• Under precision=torch.float16 on CPU, no warning is emitted and no exception is raised; the hook still drives backward() and step() through the same scaler path.

field precision: dtype [Required]

Autocast dtype and scaler policy. Accepts either a torch.dtype (e.g. torch.float16) or the canonical string name ("float32", "bfloat16", "float16"), or a shorthand alias ("fp32", "bf16", "fp16").

Constraints:

• func = <function _restrict_precision at 0xeb0b244b4f40>

• json_schema_input_type = PydanticUndefined

• return_type = PydanticUndefined

• when_used = always