API reference¶

trainloop ¶

CLASS	DESCRIPTION
`BaseHook`	Lifecycle hooks for `BaseTrainer`.
`BaseTrainer`	Minimal training loop that orchestrates builds, accumulation, retries, and hooks.
`CUDAMaxMemoryHook`	Record peak CUDA memory per step into `trainer.step_info`.
`CheckpointingHook`	Save and optionally restore checkpoints at regular intervals.
`EMAHook`	Maintain an exponential moving average of model weights.
`ImageFileLoggerHook`	Persist logged images to `workspace/visualizations` on rank 0.
`LossNoneWarning`	Warning raised when `forward` returns `None` in distributed contexts.
`ProgressHook`	Log progress to stdout with optional metrics, ETA, and memory.
`StatsHook`	Aggregate training stats and pass them to a callback.
`TensorBoardHook`	Log metrics and images to TensorBoard (rank 0 only).
`WandbHook`	Log metrics and images to Weights & Biases (rank 0 only).

FUNCTION	DESCRIPTION
`map_nested_tensor`	Apply `f` to every tensor contained in a nested structure.

BaseHook ¶

Lifecycle hooks for BaseTrainer.

BaseTrainer ¶

Minimal training loop that orchestrates builds, accumulation, retries, and hooks.

Subclasses provide component factories and a forward pass; the base class handles sequencing, mixed precision, accumulation, state management, and hook dispatch.

PARAMETER	DESCRIPTION
`max_steps`	Number of training steps to run. TYPE: `int`
`grad_clip`	Max gradient norm; if set, gradients are clipped before stepping. TYPE: `float \| None` DEFAULT: `None`
`max_non_finite_grad_retries`	Number of retries when encountering non-finite gradients (scaler disabled). TYPE: `int \| None` DEFAULT: `None`
`mixed_precision`	`"fp16"` or `"bf16"` to enable autocast; `None` disables it. TYPE: `str \| None` DEFAULT: `None`
`gradient_accumulation_steps`	Number of microsteps to accumulate before stepping. TYPE: `int \| None` DEFAULT: `None`
`workspace`	Optional working directory used by hooks (e.g., checkpoints, logs). TYPE: `Path \| str \| None` DEFAULT: `None`
`device`	Device for the model and tensors. TYPE: `device \| str \| int \| None` DEFAULT: `None`
`no_sync_accumulate`	Whether to call `no_sync` on distributed modules during accumulation. TYPE: `bool` DEFAULT: `True`
`state_dict_options`	Torch distributed checkpoint options. TYPE: `StateDictOptions \| None` DEFAULT: `None`
`logger`	Logger instance; a default logger is created when omitted. TYPE: `Logger \| None` DEFAULT: `None`

METHOD	DESCRIPTION
`build_data_loader`	Return the training data iterator.
`build_model`	Construct and return the model.
`build_optimizer`	Create the optimizer for the model.
`build_lr_scheduler`	Optionally create a learning-rate scheduler.
`build_hooks`	Return hooks to run during training.
`build_grad_scaler`	Create the gradient scaler used for mixed precision.
`train`	Run the training loop until `max_steps` are completed.
`forward`	Perform a forward pass and return loss plus records for logging.
`log_scalars`	Dispatch nested scalar records to logging hooks.
`log_images`	Dispatch image records to hooks.

build_data_loader ¶

build_data_loader() -> Iterable

Return the training data iterator.

build_model ¶

build_model() -> nn.Module

Construct and return the model.

build_optimizer ¶

build_optimizer() -> torch.optim.Optimizer

Create the optimizer for the model.

build_lr_scheduler ¶

build_lr_scheduler() -> torch.optim.lr_scheduler.LRScheduler | None

Optionally create a learning-rate scheduler.

build_hooks ¶

build_hooks() -> list[BaseHook]

Return hooks to run during training.

build_grad_scaler ¶

build_grad_scaler() -> torch.amp.GradScaler

Create the gradient scaler used for mixed precision.

train ¶

train()

Run the training loop until max_steps are completed.

forward ¶

forward(input: Any) -> tuple[torch.Tensor | None, Records]

Perform a forward pass and return loss plus records for logging.

PARAMETER	DESCRIPTION
`input`	Batch yielded by the data loader. TYPE: `Any`

RETURNS	DESCRIPTION
`Tensor \| None`	The loss (`None` skips backward/step; if using DDP/FSDP, avoid invoking the wrapped module's `forward` in that case).
`Records`	A nested dict of numeric metrics that will be averaged and emitted to hooks.

log_scalars ¶

log_scalars(records: dict[str, Any], dry_run: bool = False)

Dispatch nested scalar records to logging hooks.

Dictionary keys are treated as path components. Logging hooks may join nested components with backend-specific namespace separators, such as /. Avoid those separators within individual keys because distinct nested paths may otherwise map to the same backend metric name.

PARAMETER	DESCRIPTION
`records`	Nested dict of scalar metrics to log. TYPE: `dict[str, Any]`
`dry_run`	If True, hooks should avoid side effects and only report intent. TYPE: `bool` DEFAULT: `False`

log_images ¶

log_images(records: dict[str, Any], dry_run: bool = False)

Dispatch image records to hooks.

PARAMETER	DESCRIPTION
`records`	Nested dict of images to log. TYPE: `dict[str, Any]`
`dry_run`	If True, hooks should avoid side effects and only report intent. TYPE: `bool` DEFAULT: `False`

CUDAMaxMemoryHook ¶

Bases: BaseHook

Record peak CUDA memory per step into trainer.step_info.

CheckpointingHook ¶

Bases: BaseHook

Save and optionally restore checkpoints at regular intervals.

PARAMETER	DESCRIPTION
`interval`	Save every `interval` steps. TYPE: `int`
`keep_previous`	Keep the last N checkpoints in addition to the latest. TYPE: `int` DEFAULT: `0`
`keep_interval`	Save and keep checkpoints every `keep_interval` steps. TYPE: `int` DEFAULT: `0`
`keep_steps`	Save and keep checkpoints at these explicit step numbers. TYPE: `Sequence[int] \| None` DEFAULT: `None`
`path`	Directory (relative to workspace unless absolute) for checkpoints. TYPE: `Path \| str` DEFAULT: `'checkpoints'`
`load`	Path to load at startup or `"latest"` to auto-resume. TYPE: `Path \| str \| Literal['latest'] \| None` DEFAULT: `'latest'`
`exit_signals`	Signals that trigger a checkpoint then exit. TYPE: `list[Signals] \| Signals \| None` DEFAULT: `None`
`exit_code`	Exit code after handling an exit signal. TYPE: `int \| Literal['128+signal']` DEFAULT: `'128+signal'`
`exit_wait`	Optional sleep before exit (useful for schedulers). TYPE: `timedelta \| float` DEFAULT: `0.0`

EMAHook ¶

Bases: BaseHook

Maintain an exponential moving average of model weights.

PARAMETER	DESCRIPTION
`decay`	EMA decay rate. TYPE: `float` DEFAULT: `0.999`
`use_buffers`	Whether to include model buffers in the EMA. TYPE: `bool` DEFAULT: `False`
`name`	Name used for the EMA state in the trainer state dict. TYPE: `str` DEFAULT: `'ema'`

ImageFileLoggerHook ¶

Bases: BaseHook

Persist logged images to workspace/visualizations on rank 0.

PARAMETER	DESCRIPTION
`image_format`	File extension or callable taking the flattened image key. TYPE: `str \| Callable[[tuple[str, ...]], str]` DEFAULT: `'png'`

LossNoneWarning ¶

Bases: UserWarning

Warning raised when forward returns None in distributed contexts.

ProgressHook ¶

Bases: StatsHook

Log progress to stdout with optional metrics, ETA, and memory.

PARAMETER	DESCRIPTION
`interval`	Log every N steps. TYPE: `int` DEFAULT: `1`
`with_records`	Include per-step records in the log line. TYPE: `bool` DEFAULT: `False`
`sync`	If True, aggregate across distributed ranks. TYPE: `bool` DEFAULT: `False`
`eta_warmup`	Steps to warm up ETA calculation. TYPE: `int` DEFAULT: `10`
`show_units`	Whether to print units (s, GiB) alongside values. TYPE: `bool` DEFAULT: `True`

StatsHook ¶

Bases: BaseHook

Aggregate training stats and pass them to a callback.

PARAMETER	DESCRIPTION
`interval`	Emit stats every N steps. TYPE: `int` DEFAULT: `10`
`sync`	If True, aggregate stats across distributed ranks. TYPE: `bool` DEFAULT: `True`
`callback`	Function called with the trainer and aggregated stats. TYPE: `Callable[[BaseTrainer, TrainingStats], None]`

TensorBoardHook ¶

Bases: BaseHook

Log metrics and images to TensorBoard (rank 0 only).

Nested namespace components are joined with namespace_separator and the final name is separated with /. Avoid these separators within individual keys because TensorBoard uses them to organize tags.

PARAMETER	DESCRIPTION
`texts`	Optional text values to log when training starts. TYPE: `dict[str, Any] \| None` DEFAULT: `None`
`namespace_separator`	Separator for nested tag prefixes; the final TensorBoard name boundary is always "/". TYPE: `str` DEFAULT: `'/'`

WandbHook ¶

Bases: BaseHook

Log metrics and images to Weights & Biases (rank 0 only).

Nested scalar key components are joined with /. Avoid / within an individual component because it is also W&B's panel namespace separator.

PARAMETER	DESCRIPTION
`project`	W&B project name. TYPE: `str`
`config`	Optional config dict or JSON file path to log. TYPE: `dict[str, Any] \| str \| None` DEFAULT: `None`
`tags`	Optional tag list. TYPE: `Sequence[str] \| None` DEFAULT: `None`
`image_format`	File format for images or a callable taking the flattened image key and returning the format. TYPE: `str \| None \| Callable[[tuple[str, ...]], str \| None]` DEFAULT: `'png'`
`**wandb_kwargs`	Extra arguments forwarded to `wandb.init`. DEFAULT: `{}`

METHOD	DESCRIPTION
`on_log_images`	Note that the final component of each flattened key becomes the image caption,

on_log_images ¶

on_log_images(trainer: BaseTrainer, records: dict, dry_run: bool = False)

Note that the final component of each flattened key becomes the image caption, and the remaining components become the W&B key. For example, foo/bar/name1 and foo/bar/name2 are logged as two images under foo/bar, so W&B displays them together in one panel.

The list is replaced rather than extended when the same W&B key is logged again at the same step. Callers should therefore collect all images for a panel and pass them in a single call per step.

map_nested_tensor ¶

map_nested_tensor(f: Callable[[Tensor], Any], obj: Any)

Apply f to every tensor contained in a nested structure.