fix: 修复训练循环 step/backward 顺序，重构为三重循环嵌套

- 训练循环改用 itertools.batched 实现 epoch→step→batch 三重嵌套
- on_step_begin 包裹 batch 循环，on_step_end 后接 optimizer.step/scheduler.step
- 修复首次 iteration=0 时 optimizer.step() 在 backward 之前触发的 bug
- GradientClippingCallback 改为 on_step_end（梯度已累积，step 前裁剪）
- SchedulerCallback 移除，schduler.step 由 trainer 在 optimizer.step 后直接调用
- metric_util 提取 _grad_stat 公共 helper，if param.grad: 修正为 is not None

astrai/trainer/metric_util.py

@@ -1,75 +1,42 @@
-from typing import Dict
+from typing import Any, Callable, Dict
 
+import torch
 import torch.nn as nn
 
 
-def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+def _grad_stat(
-    """Compute gradient norm for each parameter in the model."""
+    model: nn.Module, fn: Callable[[torch.Tensor], Any], default: Any
-    norms = {}
+) -> dict:
+    results = {}
     for name, param in model.named_parameters():
-        norms[name] = 0.0
+        results[name] = default
-        if param.grad:
+        if param.grad is not None:
-            norm = param.grad.data.norm(norm_type).item()
+            results[name] = fn(param.grad.data)
-            norms[name] = norm
+    return results
-    return norms
+
+
+def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+    return _grad_stat(model, lambda g: g.norm(norm_type).item(), 0.0)
 
 
 def grad_std(model: nn.Module) -> Dict[str, float]:
-    """Compute standard deviation of gradients for each parameter."""
+    return _grad_stat(model, lambda g: g.std().item(), 0.0)
-    stds = {}
-    for name, param in model.named_parameters():
-        stds[name] = 0.0
-        if param.grad:
-            std = param.grad.data.std().item()
-            stds[name] = std
-    return stds
 
 
 def grad_max(model: nn.Module) -> Dict[str, float]:
-    """Find the maximum absolute gradient value for each parameter."""
+    return _grad_stat(model, lambda g: g.max().item(), -float("inf"))
-    max_vals = {}
-    for name, param in model.named_parameters():
-        max_vals[name] = -float("inf")
-        if param.grad:
-            max_val = param.grad.data.max().item()
-            max_vals[name] = max_val
-
-    return max_vals
 
 
 def grad_min(model: nn.Module) -> Dict[str, float]:
-    """Find the minimum absolute gradient value for each parameter."""
+    return _grad_stat(model, lambda g: g.min().item(), float("inf"))
-    min_vals = {}
-    for name, param in model.named_parameters():
-        min_vals[name] = float("inf")
-        if param.grad:
-            min_val = param.grad.data.min().item()
-            min_vals[name] = min_val
-
-    return min_vals
 
 
 def grad_mean(model: nn.Module) -> Dict[str, float]:
-    """Compute mean of gradients for each parameter."""
+    return _grad_stat(model, lambda g: g.mean().item(), 0.0)
-    means = {}
-    for name, param in model.named_parameters():
-        means[name] = 0.0
-        if param.grad:
-            mean = param.grad.data.mean().item()
-            means[name] = mean
-
-    return means
 
 
 def grad_nan_num(model: nn.Module) -> Dict[str, int]:
-    """Count the number of NaNs in gradients for each parameter."""
+    return _grad_stat(model, lambda g: g.isnan().sum().item(), 0)
-    nan_nums = {}
-    for name, param in model.named_parameters():
-        nan_nums[name] = 0
-        if param.grad:
-            nan_num = param.grad.isnan().sum().item()
-            nan_nums[name] = nan_num
-    return nan_nums
 
 
 def ctx_get_loss(ctx):

astrai/trainer/train_callback.py

@@ -79,30 +79,11 @@ class GradientClippingCallback(TrainCallback):
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
 
-    def on_step_begin(self, context: TrainContext):
+    def on_step_end(self, context: TrainContext):
         _ = context
         clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
 
 
-@CallbackFactory.register("scheduler")
-class SchedulerCallback(TrainCallback):
-    """
-    Scheduler callback for trainer.
-    """
-
-    def __init__(self):
-        pass
-
-    def on_train_begin(self, context: TrainContext):
-        for group in context.optimizer.param_groups:
-            if "initial_lr" not in group:
-                group["initial_lr"] = group["lr"]
-
-    def on_batch_end(self, context: TrainContext):
-        if context.scheduler:
-            context.scheduler.step()
-
-
 @CallbackFactory.register("checkpoint")
 class CheckpointCallback(TrainCallback):
     """

astrai/trainer/trainer.py

@@ -1,4 +1,5 @@
 import logging
+from itertools import batched
 from typing import List, Optional
 
 from astrai.config import TrainConfig
@@ -30,7 +31,6 @@ class Trainer:
             CallbackFactory.create("checkpoint", cfg.ckpt_dir, cfg.ckpt_interval),
             CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
             CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
-            CallbackFactory.create("scheduler"),
         ]
 
     def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
@@ -62,32 +62,33 @@ class Trainer:
 
         try:
             context.model.train()
-            # 1.epoch
+            accumulation_steps = max(self.train_config.accumulation_steps, 1)
+
             for epoch in range(context.epoch, self.train_config.n_epoch):
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
-                accumulation_steps = max(self.train_config.accumulation_steps, 1)
+                for steps in batched(context.dataloader, accumulation_steps):
-                for batch in context.dataloader:
-                    if context.iteration % accumulation_steps == 0:
-                        # 2. step
                     self._call_callbacks("on_step_begin", context)
-                        context.optimizer.step()
-                        context.optimizer.zero_grad()
-                        self._call_callbacks("on_step_end", context)
 
-                    # 3. batch
+                    step_batch_nums = len(steps)
+                    for batch in steps:
                         self._call_callbacks("on_batch_begin", context)
                         loss = context.strategy(batch)
                         context.loss = loss.item()
                         context.iteration += 1
 
-                    # to make the loss normalized by accumulation steps
+                        stand_loss = loss / step_batch_nums
-                    stand_loss = loss / accumulation_steps
                         stand_loss.backward()
-
                         self._call_callbacks("on_batch_end", context)
 
+                    self._call_callbacks("on_step_end", context)
+                    context.optimizer.step()
+                    context.optimizer.zero_grad()
+
+                    if context.scheduler:
+                        context.scheduler.step()
+
                 self._call_callbacks("on_epoch_end", context)
 
         except Exception as e: