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
This commit is contained in:
+18
-12
@@ -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
|
||||
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 = {
|
||||
|
||||
@@ -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.model, context.optimizer, context.scheduler = executor.prepare(
|
||||
cfg.model_fn,
|
||||
cfg.optimizer_fn,
|
||||
cfg.scheduler_fn,
|
||||
before_wrap=_before_wrap,
|
||||
)
|
||||
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"):
|
||||
|
||||
Reference in New Issue
Block a user