kempnerforge.metrics.tracker

Metrics collection, accumulation, and reporting.

MetricsTracker aggregates per-step metrics (loss, grad norm, throughput, MFU, memory) and dispatches them to configured logging backends (stdout, WandB, TensorBoard, MLflow) at a configurable interval.

Classes

MetricsTracker	Collects, smooths, and reports training metrics.
StepMetrics	Metrics for a single training step.
TensorBoardBackend	TensorBoard logging backend.
WandBBackend	Weights & Biases logging backend.

class kempnerforge.metrics.tracker.StepMetrics

Bases: object

Metrics for a single training step.

loss: float = 0.0

grad_norm: float = 0.0

lr: float = 0.0

tokens_per_sec: float = 0.0

mfu: float = 0.0

step_time_sec: float = 0.0

allocated_gb: float = 0.0

peak_gb: float = 0.0

reserved_gb: float = 0.0

total_gb: float = 0.0

mem_utilization: float = 0.0

__init__(loss=0.0, grad_norm=0.0, lr=0.0, tokens_per_sec=0.0, mfu=0.0, step_time_sec=0.0, allocated_gb=0.0, peak_gb=0.0, reserved_gb=0.0, total_gb=0.0, mem_utilization=0.0)

Parameters:

• loss (float)

• grad_norm (float)

• lr (float)

• tokens_per_sec (float)

• mfu (float)

• step_time_sec (float)

• allocated_gb (float)

• peak_gb (float)

• reserved_gb (float)

• total_gb (float)

• mem_utilization (float)

Return type:

None

class kempnerforge.metrics.tracker.MetricsTracker

Bases: object

Collects, smooths, and reports training metrics.

Timing is handled internally — call start_step() before and end_step() after each training step. Metrics are logged to all configured backends at the configured interval.

Parameters:

• config – Full job config (used for MFU calculation and backend selection).

• num_gpus – Number of GPUs for MFU denominator.

• gpu_peak_tflops – Per-GPU peak TFLOPS. If None, auto-detected.

__init__(config, num_gpus=1, gpu_peak_tflops=None)

Parameters:

• config (JobConfig)

• num_gpus (int)

• gpu_peak_tflops (float | None)

Return type:

None

start_step()

Mark the beginning of a training step.

Return type:

None

end_step(step, loss, grad_norm, lr, tokens_in_step)

Mark the end of a training step and optionally log metrics.

A logging step reads GPU memory once, from the current device. The tracker never resets the peak, so peak_gb is the most allocated since the process started, or since a reset_peak_memory elsewhere (DeviceMemoryMonitor.report makes one).

Parameters:

• step (int) – Current training step number.

• loss (float) – Loss value for this step.

• grad_norm (float) – Total gradient norm before clipping (what clip_grad_norm_ returns); under pipeline parallelism, the whole model’s norm.

• lr (float) – Current learning rate.

• tokens_in_step (int) – Total tokens processed in this step (across all GPUs).

Returns:

StepMetrics if this step was a logging step, None otherwise.

Return type:

StepMetrics | None

log_eval(metrics, step)

Log eval metrics to all backends and stdout.

Parameters:

• metrics (dict[str, float])

• step (int)

Return type:

None

init_backends(config)

Initialize logging backends (call after distributed setup).

Parameters:

config (JobConfig)

Return type:

None

tracking_entries()

One record per backend with a live run, for a checkpoint’s metadata.json.

Empty once close() has begun. Reads only what the backends already hold, so it never waits on a server. Each entry is copied once into a plain dict; a backend that cannot describe itself, or whose entry is not strings under a named backend, is skipped with a warning and costs no other backend its entry.

Return type:

list[dict[str, str]]

checkpoint_fields(step)

What each backend with a live run adds to its tracking entry of the save at step.

Keyed by backend name; this tracker is what run_training hands to CheckpointManager.set_checkpoint_registrar. Reads only what the backends hold, so it never waits on a server; a backend that fails costs only its own fields. Empty once close() has begun.

Parameters:

step (int)

Return type:

dict[str, dict[str, str]]

register_checkpoint(step, directory, record)

Hand a committed checkpoint to each backend with a live run; never blocks.

Each backend gets record bound to its own name, so it can amend only its own tracking entry, then or later from another thread. A backend that fails costs only its own registration. Nothing is registered once close() has begun.

Parameters:

• step (int)

• directory (Path)

• record (Callable[[str, dict[str, str]], None])

Return type:

None

close()

Flush and close all logging backends.

Return type:

None

class kempnerforge.metrics.tracker.WandBBackend

Bases: _LoggingBackend

Weights & Biases logging backend.

Initializes a WandB run on first log call.

__init__(config)

Parameters:

config (MetricsConfig)

Return type:

None

log(metrics, step)

Parameters:

• metrics (dict[str, float])

• step (int)

Return type:

None

tracking_entry()

This backend’s record of its live run for a checkpoint.

None until its run starts, after a failed init, and once it is closed.

Return type:

dict[str, str] | None

close()

Return type:

None

class kempnerforge.metrics.tracker.TensorBoardBackend

Bases: _LoggingBackend

TensorBoard logging backend.

__init__(config)

Parameters:

config (MetricsConfig)

Return type:

None

log(metrics, step)

Parameters:

• metrics (dict[str, float])

• step (int)

Return type:

None

tracking_entry()

This backend’s record of its live run for a checkpoint.

None until its run starts, after a failed init, and once it is closed.

Return type:

dict[str, str] | None

close()

Return type:

None