Metrics tracker¶

MetricsTracker is the central per-step metric collector. One instance is created once in training/entry.py::run_training; the training loop bookends every step with start_step() / end_step() and the tracker handles timing, smoothing, stdout formatting, and backend dispatch.

Wiring¶

# training/entry.py::run_training
tracker = MetricsTracker(config, num_gpus=runtime.world_size)
tracker.init_backends(config)     # builds the rank-0 backends
ckpt_mgr.set_tracking_source(tracker.tracking_entries)
ckpt_mgr.set_checkpoint_registrar(tracker)
...
while step < tc.max_steps:
    tracker.start_step()
    # ... forward, backward, optimizer step ...
    tracker.end_step(step, loss=loss, grad_norm=gn, lr=lr,
                     tokens_in_step=tokens_per_step)
    ...
tracker.close()                   # in a finally

init_backends is separate from __init__ so that constructing a tracker touches neither distributed state nor any backend: the rank check needs the process group, so run_training calls it after distributed init, once checkpoint restore has already written the resumed run ids into the config.

Per-step flow¶

end_step runs four things in order:

Timing: step_time = perf_counter() - _step_start; derives tokens_per_sec = tokens_in_step / step_time.
MFU: calls compute_mfu(model_config, tokens_per_sec, num_gpus, gpu_peak_tflops, seq_len). See MFU.
EMA smoothing: updates four smoothed keys (loss, tokens_per_sec, mfu, step_time) with alpha=0.1.
Log if due: if step % log_interval == 0 or step == 1, read GPU memory once with get_memory_stats(), from the current device, then dispatch to stdout + all backends and return a StepMetrics. Otherwise return None without reading memory.

The tracker never resets the CUDA peak counter, so peak_gb (logged as gpu/peak_gb and mem=) is the high-water mark of the run, and mem_utilization is that peak over total_gb (see Memory monitor).

`StepMetrics`¶

@dataclass
class StepMetrics:
    loss: float           # raw, not smoothed
    grad_norm: float
    lr: float
    tokens_per_sec: float # global (across all GPUs)
    mfu: float
    step_time_sec: float
    allocated_gb: float
    peak_gb: float
    reserved_gb: float
    total_gb: float
    mem_utilization: float

Returned from end_step() only on logging steps. The raw values are written to backends; the EMA-smoothed copies go under a smoothed/ prefix (see below).

EMA smoothing¶

# kempnerforge/metrics/tracker.py — _update_smoothed
if key not in self._smoothed:
    self._smoothed[key] = value      # bootstrap with first sample
else:
    self._smoothed[key] = alpha * value + (1 - alpha) * self._smoothed[key]

alpha = 0.1 is hardcoded. Four metrics get smoothed: loss, tokens_per_sec, mfu, step_time. Smoothed values are reported to backends as smoothed/loss, smoothed/tokens_per_sec, etc. Stdout reports the raw values.

Smoothing runs every step regardless of log_interval — the EMA is built from every data point, even ones that aren’t logged.

Log interval¶

# kempnerforge/metrics/tracker.py: end_step, after the smoothing
if step % self.metrics_config.log_interval != 0 and step != 1:
    return None

mem_stats = get_memory_stats()   # read on logging steps only
...
self._log_step(step, metrics)
return metrics

Default log_interval = 10. Step 1 is always logged (useful sanity check that the pipeline is alive). Step 0 is never reached — the training loop starts at step = 1.

Stdout format¶

2026-11-04 14:03:12.345 [rank 0] INFO     [step 1000] loss=2.3400 | lr=3.00e-04 | grad_norm=1.250 | tok/s=125,000 | mfu=52.3% | mem=71.2/80GB | step_time=1.25s

The part from [step 1000] on is produced by format_metrics(step, {...}) from kempnerforge/metrics/logger.py. The tracker formats each value of the step line itself: loss to 4 decimals, lr in scientific notation, grad_norm to 3 decimals, tok/s with thousands separators, mfu as a percentage, mem as peak over total GB and step_time in seconds. format_metrics applies _format_number only to numbers, such as the eval metrics log_eval passes: int gets k/M/B suffix, float gets scientific notation below 0.01 or above 1e6, otherwise 4 decimals.

The same logger writes every other logger.info in the framework, and each line starts with the local time to the millisecond (YYYY-MM-DD HH:MM:SS.mmm), [rank N] and the level. With ANSI color, used if stdout.isatty() and NO_COLOR isn’t set, the time and rank are dimmed and the level is colored. A record logged with logger.exception(...) or stack_info=True prints its traceback or stack on the lines below it (see _RankFormatter).

Backend dispatch¶

end_step emits a numeric dict to every configured backend:

backend_dict = {
    "train/loss": ..., "train/grad_norm": ..., "train/lr": ...,
    "train/tokens_per_sec": ..., "train/mfu": ..., "train/step_time_sec": ...,
    "gpu/allocated_gb": ..., "gpu/peak_gb": ..., "gpu/reserved_gb": ...,
    "gpu/mem_utilization": ...,
    # plus smoothed/*
}
for backend in self._backends:
    backend.log(backend_dict, step=step)

Backends are built by init_backends, on rank 0 only, and each starts its run on its first log(). They are independent toggles, so any combination can be enabled at once:

# kempnerforge/metrics/tracker.py: _init_backends
if dist.is_initialized() and dist.get_rank() != 0:
    return
if mc.enable_wandb:
    self._backends.append(WandBBackend(mc))
if mc.enable_tensorboard:
    self._backends.append(TensorBoardBackend(mc))
if mc.enable_mlflow:
    from kempnerforge.metrics.mlflow.backend import MLflowBackend

    self._backends.append(MLflowBackend(mc, job_config=config))

Non-rank-0 workers keep _backends empty; their _log_step call iterates over an empty list and only the stdout line is emitted (and even that is filtered by the _RankFilter — see Logger).

Eval path¶

def log_eval(self, metrics: dict[str, float], step: int) -> None:
    logger.info(format_metrics(step, metrics))
    for backend in self._backends:
        backend.log(metrics, step=step)

Eval metrics don’t touch the smoothing table — they’re reported raw. The training loop calls log_eval after run_eval, and separately for per-module learning rates, MoE and per-dataset breakdowns:

# training/loop.py
tracker.log_eval(lr_metrics, step)      # lr/<module> (lr_adamw/<module> under muon), with optimizer.lr_multipliers
tracker.log_eval(moe_metrics, step)     # moe/aux_loss, moe/expert_balance
tracker.log_eval(ds_metrics, step)      # per-source loss
tracker.log_eval(eval_metrics, step)    # eval loss

Each call logs with the same step, so a single training step can produce several side-by-side rows in WandB / TensorBoard.

Close¶

def close(self) -> None:
    self._closed = True
    for backend in self._backends:
        backend.close()

Called from a finally in run_training, so a run that raises still closes its backends. WandB finishes the run, MLflow drains its queued calls, checkpoint registrations included, and ends the run (each bounded, see MLflow), and TensorBoard flushes the event file. Marking the tracker closed first makes tracking_entries(), checkpoint_fields() and register_checkpoint() return nothing afterwards.

Tracking entries¶

def tracking_entries(self) -> list[dict[str, str]]:
    ...

One record per backend that has a live run, each from the backend’s own tracking_entry() (a backend that does not override it contributes nothing). A backend that raises, or whose entry is not strings under a named backend, is skipped with a warning. run_training hands this to CheckpointManager.set_tracking_source, so every checkpoint’s metadata.json names the runs it was trained under (Checkpointing). It reads fields the backends already hold, never waits on a server, and returns an empty list before any run has started and after close().

Registering checkpoints¶

def checkpoint_fields(self, step: int) -> dict[str, dict[str, str]]:
    ...

def register_checkpoint(self, step: int, directory: Path, record: TrackingRecord) -> None:
    ...

run_training hands the tracker to CheckpointManager.set_checkpoint_registrar, which calls both on rank 0. checkpoint_fields runs inside save(): each backend with a live run names the fields it adds to its own entry of that save’s tracking list, keyed by backend name, from what it already holds. register_checkpoint runs once the checkpoint is committed and must not block: each backend with a live run gets record bound to its own name and calls it, then or later from another thread, to add fields to its entry; the manager applies each call to that checkpoint’s metadata.json atomically. Only MLflowBackend does either, when mlflow_register_checkpoints is on (MLflow). A backend that raises costs only its own fields or registration, and neither method does anything after close().

Logger: `get_logger`, `format_metrics`¶

Anywhere in the framework, call:

from kempnerforge.metrics import get_logger
logger = get_logger(__name__)
logger.info("starting distributed init")

Returns a logging.Logger named kempnerforge.<name>. On first call, configures the root kempnerforge logger with a _RankFormatter (adds the local time, [rank N] and the colored level, and prints any traceback) and a _RankFilter (drops records from non-zero ranks by default).

Override rank filtering with get_logger(__name__, rank_zero_only=False) when you genuinely want every rank to log — e.g. for distributed-init debugging. Once the root logger is configured it’s cached, so this flag only takes effect on the very first call.

A single record can pass the filter instead: log it with extra={"every_rank": True} and every rank emits it. Use this for an event that only the logging rank sees, such as a problem with a sample that rank read.