From 44579ea6dc0a8b95b84a5e90b952fdf57a71eab5 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Tue, 30 Jun 2026 15:12:31 +0800 Subject: [PATCH] =?UTF-8?q?refactor=20:=20metric=20=E6=97=A5=E5=BF=97?= =?UTF-8?q?=E6=94=B9=E4=B8=BA=E4=BB=A5=20optimizer=20step=20=E4=B8=BA?= =?UTF-8?q?=E5=8D=95=E4=BD=8D=EF=BC=8C=E9=BB=98=E8=AE=A4=E6=AF=8F=E6=AD=A5?= =?UTF-8?q?=E8=AE=B0=E5=BD=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - log_interval 默认 100 -> 1，语义从 batch iteration 改为 optimizer step - step 指标从 on_batch_end 移到 on_optimizer_step，不受梯度累积影响 - JSONL 条目新增 step 字段，保留 iter - flush 落盘仍在 on_batch_end --- astrai/config/train_config.py | 4 ++-- astrai/trainer/train_callback.py | 14 +++++++++----- scripts/tools/train.py | 4 ++-- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 846c67b..8575301 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -68,8 +68,8 @@ class TrainConfig(BaseConfig): default="./checkpoint/logs", metadata={"help": "Directory for metric logs."} ) log_interval: int = field( - default=100, - metadata={"help": "Number of batch iterations between metric logs."}, + default=1, + metadata={"help": "Number of optimizer steps between metric logs."}, ) metrics: List[str] = field( default_factory=lambda: ["loss", "lr", "grad_norm"], diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index a052122..ab22421 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -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 diff --git a/scripts/tools/train.py b/scripts/tools/train.py index e1a5291..10ad134 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -162,8 +162,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--log_interval", type=int, - default=100, - help="Number of batch iterations between metric logs.", + default=1, + help="Number of optimizer steps between metric logs.", ) parser.add_argument( "--grpo_sync_interval",