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

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"],

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

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",