nvalchemi.hooks.TrainContext

class nvalchemi.hooks.TrainContext(*, batch, model=None, global_rank=0, workflow=None, step_count=0, global_step_count=0, batch_count=0, epoch_step_count=0, epoch=0, loss=None, losses=None, models=None, optimizers=<factory>, lr_schedulers=<factory>, gradients=None, grad_scaler=None, validation=None)

Context object passed to training hooks.

Parameters:

• batch (Batch | None)

• model (BaseModelMixin | None)

• global_rank (int)

• workflow (Any)

• step_count (int)

• global_step_count (int)

• batch_count (int)

• epoch_step_count (int)

• epoch (int)

• loss (torch.Tensor | None)

• losses (dict[str, torch.Tensor] | None)

• models (dict[str, BaseModelMixin] | ModuleDict | None)

• optimizers (list[torch.optim.Optimizer])

• lr_schedulers (list[LRScheduler | None])

• gradients (dict[str, torch.Tensor] | None)

• grad_scaler (torch.amp.GradScaler | None)

• validation (dict[str, Any] | None)

step_count

Current optimizer step number on this worker.

Type:

int

global_step_count

Current optimizer step number across all data-parallel workers.

Type:

int

batch_count

Number of training batches consumed, including batches whose optimizer step was skipped by update hooks.

Type:

int

epoch_step_count

Number of batches consumed within the current training epoch.

Type:

int

epoch

Current training epoch.

Type:

int

loss

Aggregate loss for the current step.

Type:

torch.Tensor | None

losses

Named loss components for the current step.

Type:

dict[str, torch.Tensor] | None

models

Models participating in the training step; this differs from the model attribute which is intended to represent a ‘main’ model in multi-model workflows. The key/model mapping should be semantic, e.g. ‘student’ and ‘teacher’ in distillation workflows, with ‘student’ being the intended ‘main’ model.

Type:

dict[str, BaseModelMixin] | ModuleDict | None

optimizers

Optimizers participating in the training step. Empty when no optimizer is attached (e.g. eval-only or manually-driven hook contexts); TrainingUpdateOrchestrator and similar consumers treat an empty list as a no-op.

Type:

list[torch.optim.Optimizer]

lr_schedulers

Learning rate schedulers participating in the training step. Aligned positionally with optimizers when populated; entries may be None when an optimizer has no scheduler. Empty when no scheduler is attached.

Type:

list[torch.optim.lr_scheduler.LRScheduler | None]

gradients

Parameter gradients for the current step.

Type:

dict[str, torch.Tensor] | None

grad_scaler

AMP gradient scaler for mixed-precision training; None when AMP is not in use.

Type:

torch.amp.GradScaler | None

validation

Latest validation summary produced by the training strategy’s validation checkpoint (TrainingStrategy.validate()). None until validation has run or after the latest summary is consumed by metric-driven schedulers. In distributed runs, the reduced summary is available on every rank.

Type:

dict[str, Any] | None