nvalchemi.training.hooks.TrainingUpdateOrchestrator

class nvalchemi.training.hooks.TrainingUpdateOrchestrator(*hooks)

Composes TrainingUpdateHook instances and drives updates.

Claims the training-update stages BEFORE_BATCH, DO_BACKWARD, DO_OPTIMIZER_STEP, AFTER_OPTIMIZER_STEP. The strategy also calls the orchestrator during SETUP so child hooks can initialize runtime state before the first batch. Per-stage behavior is selected by direct TrainingStage comparisons to avoid per-batch multiple-dispatch overhead. See Training update hooks for the stage contract enforced by the orchestrator.

Parameters:

*hooks (TrainingUpdateHook or TrainingUpdateOrchestrator) – Hooks to compose. Any orchestrator argument is flattened into its children. Members are sorted by priority ascending; ties preserve insertion order (Python’s stable sort).

frequency

Required by the Hook Protocol; always 1.

Type:

int

stage

Set to None so the registry consults _runs_on_stage.

Type:

None

Raises:

TypeError – If any positional argument is not a TrainingUpdateHook or TrainingUpdateOrchestrator.

Parameters:

hooks (TrainingUpdateHook | TrainingUpdateOrchestrator)

Notes

TrainingUpdateOrchestrator IS compatible with the standard Hook Protocol – it is the registry-facing wrapper around one or more TrainingUpdateHook instances. Concrete training update hooks (EMAHook, GradientClipHook, etc.) are NOT directly Protocol-compliant on their own; they must be composed into an orchestrator before registration. The training strategy auto-wraps a bare TrainingUpdateHook for convenience.

On DO_BACKWARD each hook returns (_, loss); the orchestrator assigns ctx.loss = loss between hooks so the next hook sees the transformed value. backward() is called once on the final ctx.loss. Example: a *0.5 hook followed by a *2.0 hook leaves ctx.loss equal to the original loss before backward.

close()

Close child update hooks that expose close.

Return type:

None

iter_hooks()

Yield child update hooks in orchestrator dispatch order.

Return type:

Iterator[TrainingUpdateHook]

property optimizer_step_skipped: bool

Whether the most recent optimizer-step stage was vetoed.