refactor : metric 日志改为以 optimizer step 为单位,默认每步记录

- log_interval 默认 100 -> 1,语义从 batch iteration 改为 optimizer step
- step 指标从 on_batch_end 移到 on_optimizer_step,不受梯度累积影响
- JSONL 条目新增 step 字段,保留 iter
- flush 落盘仍在 on_batch_end
This commit is contained in:
2026-06-30 15:12:31 +08:00
parent 0f1fcb079f
commit 44579ea6dc
3 changed files with 13 additions and 9 deletions
+9 -5
View File
@@ -229,13 +229,14 @@ class MetricLoggerCallback(TrainCallback):
self,
log_dir: str,
save_interval: int,
log_interval: int = 10,
log_interval: int = 1,
metrics: List[str] = None,
):
self.last_log_iter = 0
self._last_val_loss = None
self._last_log_step = 0
self.save_interval = save_interval
self.log_interval = log_interval
self.log_interval = max(log_interval, 1)
self.metrics = metrics or ["loss", "lr"]
self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
@@ -263,6 +264,7 @@ class MetricLoggerCallback(TrainCallback):
"type": event_type,
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
"epoch": context.epoch,
"step": context.iteration // context.config.grad_accum_steps,
"iter": context.iteration,
**extra,
}
@@ -277,14 +279,16 @@ class MetricLoggerCallback(TrainCallback):
f.write(json.dumps(log) + "\n")
def on_batch_end(self, context):
if context.iteration % self.log_interval == 0:
step_metrics = [m for m in self.metrics if m != "val_loss"]
self._append("step", context, **self._metrics(context, step_metrics))
if context.iteration - self.last_log_iter >= self.save_interval:
self._flush(context.epoch, context.iteration)
self.last_log_iter = context.iteration
def on_optimizer_step(self, context):
step = context.iteration // context.config.grad_accum_steps
if step - self._last_log_step >= self.log_interval:
step_metrics = [m for m in self.metrics if m != "val_loss"]
self._append("step", context, **self._metrics(context, step_metrics))
self._last_log_step = step
if context.val_loss is not None and context.val_loss != self._last_val_loss:
self._append("validation", context, val_loss=context.val_loss)
self._last_val_loss = context.val_loss