kempnerforge.metrics

Metrics, MFU computation, memory tracking, and logging for KempnerForge.

class kempnerforge.metrics.DeviceMemoryMonitor

Bases: object

Tracks GPU memory usage across training steps.

Resets peak memory stats at each reporting interval so that the peak reflects per-interval usage rather than all-time peak.

Supports memory snapshot capture at a configurable step for debugging OOM and memory fragmentation with pytorch.org/memory_viz.

Parameters:

• device – CUDA device index, or None for the device that is current at each call.

• snapshot_step – Step at which to capture a memory snapshot. None to disable.

• snapshot_dir – Directory to save snapshots.

__init__(device=None, snapshot_step=None, snapshot_dir='memory_snapshots')

Parameters:

• device (int | None)

• snapshot_step (int | None)

• snapshot_dir (str)

Return type:

None

report(step)

Report memory stats for the current interval and reset peak.

Parameters:

step (int) – Current training step.

Returns:

Dict with memory stats.

Return type:

dict[str, float]

capture_snapshot(step)

Capture a CUDA memory snapshot and save as pickle.

The snapshot can be visualized at https://pytorch.org/memory_viz

Parameters:

step (int) – Current step (used in filename).

Returns:

Path to the saved snapshot, or None if capture failed.

Return type:

str | None

class kempnerforge.metrics.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.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

kempnerforge.metrics.compute_mfu(config, tokens_per_sec, num_gpus=1, gpu_peak_tflops=None, seq_len=None)

Compute Model FLOPs Utilization.

Parameters:

• config (ModelConfig) – Model configuration.

• tokens_per_sec (float) – Global throughput (tokens/sec across all GPUs).

• num_gpus (int) – Number of GPUs.

• gpu_peak_tflops (float | None) – Peak bf16 TFLOPS per GPU. Auto-detected if None.

• seq_len (int | None) – Actual training sequence length for attention FLOPS. Falls back to config.max_seq_len if not provided.

Returns:

MFU as a fraction (0.0 to 1.0).

Return type:

float

kempnerforge.metrics.estimate_model_flops_per_token(config, seq_len=None)

Estimate FLOPS per token for forward + backward pass.

Uses the PaLM paper approximation: 6*P + 12*L*D*S

For MoE: uses active params, moe_top_k routed plus moe_shared_experts shared experts per MoE layer (not all experts), each at the expert FFN width computed_expert_ffn_hidden_dim; dense layers keep the dense FFN width. Excludes embedding (table lookup, not matmul). Includes output projection. The 12*L*D*S attention term does not discount GQA — FlashAttention expands GQA internally, so the hardware performs full attention compute. Router FLOPS (dim × num_experts) are intentionally omitted — negligible.

Parameters:

• config (ModelConfig) – Model configuration.

• seq_len (int | None) – Actual training sequence length. Falls back to config.max_seq_len if not provided.

Returns:

Estimated FLOPS per token.

Return type:

int

kempnerforge.metrics.format_memory_stats(device=None)

Format memory stats of device (None: the current device) as a human-readable string.

Parameters:

device (int | None)

Return type:

str

kempnerforge.metrics.format_metrics(step, metrics)

Format a metrics dict into a compact, color-coded log line.

Example output:

[step 1000] loss=2.34 | lr=3.00e-04 | grad_norm=1.2 | tok/s=125k | mfu=52.3% | mem=71.2/80GB

Parameters:

• step (int)

• metrics (dict[str, float | int | str])

Return type:

str

kempnerforge.metrics.get_gpu_peak_tflops(device=0)

Auto-detect GPU peak bf16 TFLOPS.

Tries to match the GPU name against known models. Falls back to a conservative estimate based on compute capability.

Parameters:

device (int) – CUDA device index.

Returns:

Peak bf16 TFLOPS for this GPU.

Return type:

float

kempnerforge.metrics.get_logger(name, rank_zero_only=True)

Get a logger for the given module name.

Parameters:

• name (str) – Logger name (typically __name__).

• rank_zero_only (bool) – If True, only rank 0 emits logs, except records logged with extra={"every_rank": True}, which every rank emits.

Returns:

A configured logging.Logger instance.

Return type:

Logger

kempnerforge.metrics.get_memory_stats(device=None)

Get current GPU memory statistics in GB.

Parameters:

device (int | None) – CUDA device index, or None for the current device.

Returns:

Dict with allocated, peak, reserved, and total memory in GB. Peak is the most allocated since the process started or the device’s last reset_peak_memory.

Return type:

dict[str, float]

kempnerforge.metrics.get_memory_utilization(device=None)

Get peak memory utilization as a fraction of total GPU memory.

Parameters:

device (int | None) – CUDA device index, or None for the current device.

Returns:

Utilization between 0.0 and 1.0.

Return type:

float

kempnerforge.metrics.reset_peak_memory(device=None)

Reset the peak memory counter of device (None: the current device).

Parameters:

device (int | None)

Return type:

None

Modules

logger	Rank-aware logging utilities for KempnerForge.
memory	GPU memory tracking and reporting.
mfu	Model FLOPs Utilization (MFU) computation.
mlflow	MLflow experiment tracking: the training backend, checkpoint models and the eval publisher.
tracker	Metrics collection, accumulation, and reporting.