From 74b9308883e5acfed9519f7973504eee24e8f5e8 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 20 Jul 2026 01:31:40 +0800 Subject: [PATCH] refactor: pass model_fn/optimizer_fn to executor.prepare - BaseExecutor.prepare now takes factories and instantiates model via model_fn(), runs before_wrap hook, wraps DDP/FSDP, then builds optimizer/scheduler on the wrapped model - optimizer/scheduler creation moved into executor.prepare, eliminating the old 'create-then-wrap' hack reliance on use_orig_params=True - FSDPExecutor/BaseExecutor accept **_extra kwargs to tolerate DDP-only keys (broadcast_buffers, gradient_as_bucket_view) being forwarded via executor_kwargs - dataloader builds stay external; executor only handles model/optimizer/scheduler - train_context.py rewritten to load checkpoint state_dict before prepare via a before_wrap closure --- astrai/parallel/executor.py | 34 +++++++----- astrai/trainer/train_context.py | 98 ++++++++++++++++----------------- 2 files changed, 69 insertions(+), 63 deletions(-) diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index 7b4c112..58b90d8 100644 --- a/astrai/parallel/executor.py +++ b/astrai/parallel/executor.py @@ -4,7 +4,7 @@ import contextlib import logging import os from contextlib import contextmanager -from typing import Optional, Tuple +from typing import Callable, Optional, Tuple import torch import torch.distributed as dist @@ -14,7 +14,6 @@ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler -from torch.utils.data import DataLoader from astrai.factory import BaseFactory from astrai.parallel.setup import get_rank, get_world_size @@ -81,24 +80,30 @@ class AccumScheduler: class BaseExecutor: - def __init__(self, grad_accum_steps: int = 1): + def __init__(self, grad_accum_steps: int = 1, **_extra): self.gradient_state = GradientState(grad_accum_steps) def prepare( self, - model: nn.Module, - optimizer: Optional[Optimizer] = None, - dataloader: Optional[DataLoader] = None, - scheduler: Optional[LRScheduler] = None, - ) -> Tuple[ - nn.Module, Optional[Optimizer], Optional[DataLoader], Optional[LRScheduler] - ]: + model_fn: Callable[[], nn.Module], + optimizer_fn: Optional[Callable[[nn.Module], Optimizer]] = None, + scheduler_fn: Optional[Callable[[Optimizer], LRScheduler]] = None, + before_wrap: Optional[Callable[[nn.Module], nn.Module]] = None, + ) -> Tuple[nn.Module, Optional[Optimizer], Optional[LRScheduler]]: + model = model_fn() + if before_wrap is not None: + model = before_wrap(model) model = self._prepare_model(model) - if optimizer is not None: + optimizer = None + scheduler = None + if optimizer_fn is not None: + optimizer = optimizer_fn(model) + if scheduler_fn is not None: + scheduler = scheduler_fn(optimizer) optimizer = AccumOptimizer(optimizer, self.gradient_state) - if scheduler is not None: - scheduler = AccumScheduler(scheduler, self.gradient_state) - return model, optimizer, dataloader, scheduler + if scheduler is not None: + scheduler = AccumScheduler(scheduler, self.gradient_state) + return model, optimizer, scheduler def _prepare_model(self, model: nn.Module) -> nn.Module: return model @@ -243,6 +248,7 @@ class FSDPExecutor(BaseExecutor): limit_all_gathers: bool = True, ignored_states=None, device_mesh=None, + **_ddp_only_kwargs, ): super().__init__(grad_accum_steps=grad_accum_steps) self._fsdp_kwargs = { diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index dc7d411..6c7eacd 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -72,61 +72,70 @@ class TrainContextBuilder: **cfg.executor_kwargs, ) - model = cfg.model_fn() - model = model.to(device=device) - model_config = {} if self._param_path: config_path = Path(self._param_path) / "config.json" if config_path.exists(): model_config = load_json(config_path) - if not model_config and hasattr(model, "config"): - model_config = model.config.to_dict() + preloaded_state_dict = None + preloaded_epoch = cfg.start_epoch + preloaded_consumed = cfg.start_samples * get_world_size() + preloaded_checkpoint = None + if self._param_path: + checkpoint = Checkpoint.load_any(self._param_path) + if checkpoint is not None: + preloaded_state_dict = checkpoint.state_dict + if checkpoint.config: + model_config = checkpoint.config + if self._resume: + preloaded_epoch = checkpoint.epoch or cfg.start_epoch + if checkpoint.consumed_samples > 0: + per_step = ( + cfg.batch_per_device + * get_world_size() + * cfg.grad_accum_steps + ) + preloaded_consumed = ( + checkpoint.consumed_samples // per_step + ) * per_step + else: + preloaded_consumed = cfg.start_samples * get_world_size() + preloaded_checkpoint = checkpoint + + if not model_config and hasattr(cfg.model_fn(), "config"): + model_config = cfg.model_fn().config.to_dict() + + def _before_wrap(m): + m = m.to(device=device) + if preloaded_state_dict is not None: + m.load_state_dict(preloaded_state_dict, strict=False) + if cfg.lora is not None: + inject_lora( + m, + r=cfg.lora.r, + alpha=cfg.lora.alpha, + target_modules=set(cfg.lora.target_modules), + ) + return m context = TrainContext( - model=model, world_size=get_world_size(), rank=get_rank(), config=cfg, model_config=model_config, executor=executor, + epoch=preloaded_epoch, + consumed_samples=preloaded_consumed, + checkpoint=preloaded_checkpoint, ) - if self._param_path: - checkpoint = Checkpoint.load_any(self._param_path) - if checkpoint is not None: - model.load_state_dict(checkpoint.state_dict, strict=False) - if checkpoint.config: - context.model_config = checkpoint.config - - if self._resume: - context.epoch = checkpoint.epoch or cfg.start_epoch - if checkpoint.consumed_samples > 0: - per_step = ( - cfg.batch_per_device - * context.world_size - * cfg.grad_accum_steps - ) - context.consumed_samples = ( - checkpoint.consumed_samples // per_step - ) * per_step - else: - context.consumed_samples = ( - cfg.start_samples * context.world_size - ) - context.checkpoint = checkpoint - - if cfg.lora is not None: - inject_lora( - model, - r=cfg.lora.r, - alpha=cfg.lora.alpha, - target_modules=set(cfg.lora.target_modules), - ) - - context.optimizer = cfg.optimizer_fn(model) - context.scheduler = cfg.scheduler_fn(context.optimizer) + context.model, context.optimizer, context.scheduler = executor.prepare( + cfg.model_fn, + cfg.optimizer_fn, + cfg.scheduler_fn, + before_wrap=_before_wrap, + ) train_dataset = cfg.dataset val_dataset = cfg.val_dataset @@ -175,15 +184,6 @@ class TrainContextBuilder: collate_fn=cfg.collate_fn, ) - context.model, context.optimizer, context.dataloader, context.scheduler = ( - executor.prepare( - model, - context.optimizer, - context.dataloader, - context.scheduler, - ) - ) - if context.checkpoint and context.checkpoint.extra: extra = context.checkpoint.extra for name in ("optimizer", "scheduler"):