112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
|
"""Training metric buffering and wandb emission helpers."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections import deque
|
||
|
|
from typing import Any, Deque, Dict, List, Optional, Tuple
|
||
|
|
|
||
|
|
|
||
|
|
class MetricsLogger:
|
||
|
|
"""Buffer per-step stats and emit averaged batches to wandb.
|
||
|
|
|
||
|
|
Also formats the per-step console line via ``format_training_line``,
|
||
|
|
which consumes an ``(stat_key, pretty_label, fmt_spec)`` triple list
|
||
|
|
from ``adapter.console_fields()``.
|
||
|
|
|
||
|
|
``smooth_window`` controls a separate rolling-mean buffer used purely
|
||
|
|
for display smoothing (console + tqdm). It is independent of the
|
||
|
|
``log_interval`` wandb-flush buffer; the smoothed dict is returned by
|
||
|
|
:meth:`smooth` for the caller to assign to ``_current_step_stats``.
|
||
|
|
1 (default) preserves historical per-step values.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, *, wandb_run, log_interval: int, smooth_window: int = 1):
|
||
|
|
self._wandb_run = wandb_run
|
||
|
|
self._log_interval = log_interval
|
||
|
|
self._buffer: List[Dict[str, Any]] = []
|
||
|
|
self._smooth_window = max(1, int(smooth_window))
|
||
|
|
self._smooth_buffers: Dict[str, Deque[float]] = {}
|
||
|
|
|
||
|
|
def record_step(self, step_stats: Dict[str, Any], *, is_main: bool) -> None:
|
||
|
|
"""Push a per-step stats dict into the rank-0 buffer (no-op off rank-0)."""
|
||
|
|
if is_main:
|
||
|
|
self._buffer.append(step_stats)
|
||
|
|
|
||
|
|
def flush_if_due(self, global_step: int) -> Optional[Dict[str, Any]]:
|
||
|
|
"""At log_interval boundaries, average + emit; return avg dict or None."""
|
||
|
|
if global_step % self._log_interval != 0 or not self._buffer:
|
||
|
|
return None
|
||
|
|
avg = self._average(self._buffer)
|
||
|
|
if self._wandb_run is not None and hasattr(self._wandb_run, "log"):
|
||
|
|
self._wandb_run.log(avg, step=global_step)
|
||
|
|
self._buffer = []
|
||
|
|
return avg
|
||
|
|
|
||
|
|
def smooth(self, stats: Dict[str, Any]) -> Dict[str, Any]:
|
||
|
|
"""Return *stats* with numeric entries replaced by their rolling mean.
|
||
|
|
|
||
|
|
Uses one ``deque(maxlen=smooth_window)`` per metric key. Non-numeric
|
||
|
|
entries pass through unchanged. When ``smooth_window <= 1`` this is
|
||
|
|
the identity (historical behavior).
|
||
|
|
"""
|
||
|
|
if self._smooth_window <= 1:
|
||
|
|
return stats
|
||
|
|
out: Dict[str, Any] = {}
|
||
|
|
for key, val in stats.items():
|
||
|
|
if not isinstance(val, (int, float)):
|
||
|
|
out[key] = val
|
||
|
|
continue
|
||
|
|
buf = self._smooth_buffers.get(key)
|
||
|
|
if buf is None or buf.maxlen != self._smooth_window:
|
||
|
|
buf = deque(maxlen=self._smooth_window)
|
||
|
|
self._smooth_buffers[key] = buf
|
||
|
|
buf.append(float(val))
|
||
|
|
out[key] = sum(buf) / len(buf)
|
||
|
|
return out
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _average(buf: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||
|
|
all_keys: set = set()
|
||
|
|
for stats in buf:
|
||
|
|
all_keys.update(stats.keys())
|
||
|
|
avg: Dict[str, Any] = {}
|
||
|
|
for key in all_keys:
|
||
|
|
values = [s[key] for s in buf if key in s]
|
||
|
|
if values:
|
||
|
|
avg[key] = sum(values) / len(values)
|
||
|
|
return avg
|
||
|
|
|
||
|
|
def format_training_line(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
epoch: int,
|
||
|
|
total_epoch: int,
|
||
|
|
current_iter: int,
|
||
|
|
total_iter: int,
|
||
|
|
loss: float,
|
||
|
|
lr: float,
|
||
|
|
time_per_step: float,
|
||
|
|
stats: Dict[str, Any],
|
||
|
|
fields: List[Tuple[str, str, str]],
|
||
|
|
mfu_info: Optional[Dict[str, Any]] = None,
|
||
|
|
) -> str:
|
||
|
|
"""Render a line matching pre-refactor training_log output exactly.
|
||
|
|
|
||
|
|
Layout: " epoch E/T | iter I/T | loss L | [adapter fields] | lr L |
|
||
|
|
time_current_backward_step Ts | [MFU P% |]"
|
||
|
|
"""
|
||
|
|
parts = [
|
||
|
|
" epoch {:3d}/{:3d} |".format(epoch, total_epoch),
|
||
|
|
" iter {:6d}/{:6d} |".format(current_iter, total_iter),
|
||
|
|
" loss {:.6f} |".format(loss),
|
||
|
|
]
|
||
|
|
for key, label, fmt in fields:
|
||
|
|
value = stats.get(key)
|
||
|
|
if value is not None:
|
||
|
|
parts.append(" {} {:{fmt}} |".format(label, value, fmt=fmt))
|
||
|
|
parts.append(" lr {:.6f} |".format(lr))
|
||
|
|
parts.append(" time_current_backward_step {:.6f}s |".format(time_per_step))
|
||
|
|
if mfu_info is not None:
|
||
|
|
parts.append(" MFU {:.2f}% |".format(mfu_info["mfu"] * 100))
|
||
|
|
return "".join(parts)
|