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

@@ -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,
+        model_fn: Callable[[], nn.Module],
-        optimizer: Optional[Optimizer] = None,
+        optimizer_fn: Optional[Callable[[nn.Module], Optimizer]] = None,
-        dataloader: Optional[DataLoader] = None,
+        scheduler_fn: Optional[Callable[[Optimizer], LRScheduler]] = None,
-        scheduler: Optional[LRScheduler] = None,
+        before_wrap: Optional[Callable[[nn.Module], nn.Module]] = None,
-    ) -> Tuple[
+    ) -> Tuple[nn.Module, Optional[Optimizer], Optional[LRScheduler]]:
-        nn.Module, Optional[Optimizer], Optional[DataLoader], 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
+        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 = {

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"):
+        preloaded_state_dict = None
-            model_config = model.config.to_dict()
+        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:
+        context.model, context.optimizer, context.scheduler = executor.prepare(
-            checkpoint = Checkpoint.load_any(self._param_path)
+            cfg.model_fn,
-            if checkpoint is not None:
+            cfg.optimizer_fn,
-                model.load_state_dict(checkpoint.state_dict, strict=False)
+            cfg.scheduler_fn,
-                if checkpoint.config:
+            before_wrap=_before_wrap,
-                    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)
 
         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"):