refactor: 重构训练后端为 Executor 模式

- backend.py → executor.py，BaseTrainingBackend → BaseExecutor
- 新增 NoneExecutor（单卡）和 DDPExecutor（DDP，world_size=1 自动降级）
- 新增 GradientState 分离梯度同步状态，AccumOptimizer/AccumScheduler 包裹拦截
- 新增 astrai/protocols.py：OptimizerProtocol/SchedulerProtocol 结构子类型
- TrainContext.backend → executor，TrainConfig 移除 parallel_wrapper/state_dict_fn，新增 parallel_mode/executor_kwargs
- 训练循环用 accumulate() 包裹，on_optimizer_step 命名约定=gate
- scripts/tools/train.py 移除 ddp_wrap/prepare_checkpoint，新增 --parallel_mode

astrai/trainer/train_context.py

@@ -2,13 +2,13 @@ from dataclasses import dataclass, field
 from typing import Optional, Self
 
 import torch.nn as nn
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import ResumableDistributedSampler
+from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
@@ -18,10 +18,11 @@ class TrainContext:
     model: nn.Module = field(default=None)
     strategy: BaseStrategy = field(default=None)
     dataloader: DataLoader = field(default=None)
-    optimizer: Optimizer = field(default=None)
-    scheduler: LRScheduler = field(default=None)
+    optimizer: OptimizerProtocol = field(default=None)
+    scheduler: SchedulerProtocol = field(default=None)
     checkpoint: Checkpoint = field(default=None)
     config: TrainConfig = field(default=None)
+    executor: BaseExecutor = field(default=None)
 
     epoch: int = field(default=0)
     iteration: int = field(default=0)
@@ -47,22 +48,28 @@ class TrainContextBuilder:
         return self
 
     def build(self) -> TrainContext:
+        cfg = self.config
+        device = get_current_device()
+
+        executor = ExecutorFactory.create(
+            cfg.parallel_mode,
+            grad_accum_steps=cfg.grad_accum_steps,
+            **cfg.executor_kwargs,
+        )
+
         context = TrainContext(
-            model=self.config.model,
+            model=cfg.model,
             world_size=get_world_size(),
             rank=get_rank(),
-            config=self.config,
+            config=cfg,
+            executor=executor,
         )
 
-        device = get_current_device()
         context.model = context.model.to(device=device)
 
-        if self.config.nprocs > 1 and self.config.parallel_wrapper:
-            context.model = self.config.parallel_wrapper(context.model)
-
         if self._checkpoint is not None:
-            context.epoch = max(self._checkpoint.epoch, self.config.start_epoch)
-            context.iteration = max(self._checkpoint.iteration, self.config.start_batch)
+            context.epoch = max(self._checkpoint.epoch, cfg.start_epoch)
+            context.iteration = max(self._checkpoint.iteration, cfg.start_batch)
             context.model.load_state_dict(self._checkpoint.state_dict)
             context.checkpoint = self._checkpoint
         else:
@@ -70,10 +77,9 @@ class TrainContextBuilder:
                 state_dict=context.model.state_dict(),
             )
 
-        context.optimizer = self.config.optimizer_fn(context.model)
-        context.scheduler = self.config.scheduler_fn(context.optimizer)
+        context.optimizer = cfg.optimizer_fn(context.model)
+        context.scheduler = cfg.scheduler_fn(context.optimizer)
 
-        cfg = self.config
         sampler_offset = context.iteration * cfg.batch_per_device
         sampler = ResumableDistributedSampler(
             data_source=cfg.dataset,
@@ -107,11 +113,20 @@ class TrainContextBuilder:
                 prefetch_factor=cfg.prefetch_factor,
             )
 
+        context.model, context.optimizer, context.dataloader, context.scheduler = (
+            executor.prepare(
+                context.model,
+                context.optimizer,
+                context.dataloader,
+                context.scheduler,
+            )
+        )
+
         context.strategy = StrategyFactory.create(
             model=context.model,
-            train_type=self.config.strategy,
+            train_type=cfg.strategy,
             device=device,
-            **self.config.extra_kwargs,
+            **cfg.extra_kwargs,
         )
 
         return context