ViperEkura
8150ab6c32
- Add --compile flag (default/reduce-overhead/max-autotune) - Apply torch.compile in _before_wrap before DDP/FSDP wrapping - Profiling shows MFU 85.5% -> 88.5% (+3%), time -3.2%, memory -7.9%
289 lines
10 KiB
Python
289 lines
10 KiB
Python
import logging
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Self
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader, random_split
|
|
|
|
from astrai.config.train_config import TrainConfig
|
|
from astrai.dataset import RDSampler
|
|
from astrai.inference.core.scheduler import InferenceScheduler
|
|
from astrai.model.components.lora import inject_lora
|
|
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, load_json
|
|
from astrai.tokenize import AutoTokenizer
|
|
from astrai.trainer.rollout import RolloutGenerator, RolloutRunner
|
|
from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TrainContext:
|
|
model: nn.Module = field(default=None)
|
|
strategy: BaseStrategy = field(default=None)
|
|
dataloader: DataLoader = field(default=None)
|
|
optimizer: OptimizerProtocol = field(default=None)
|
|
scheduler: SchedulerProtocol = field(default=None)
|
|
checkpoint: Checkpoint = field(default=None)
|
|
config: TrainConfig = field(default=None)
|
|
model_config: dict = field(default_factory=dict)
|
|
executor: BaseExecutor = field(default=None)
|
|
epoch: int = field(default=0)
|
|
consumed_samples: int = field(default=0)
|
|
loss: float = field(default=0.0)
|
|
grad_norm: Optional[float] = field(default=None)
|
|
val_dataloader: Optional[DataLoader] = field(default=None)
|
|
val_loss: Optional[float] = field(default=None)
|
|
|
|
world_size: int = field(default=1)
|
|
rank: int = field(default=0)
|
|
kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
_stop_event: threading.Event = field(default_factory=threading.Event)
|
|
|
|
@property
|
|
def stop_requested(self) -> bool:
|
|
return self._stop_event.is_set()
|
|
|
|
def request_stop(self) -> None:
|
|
self._stop_event.set()
|
|
|
|
@property
|
|
def optimizer_step(self) -> int:
|
|
return self.consumed_samples // (
|
|
self.config.batch_per_device
|
|
* self.world_size
|
|
* self.config.grad_accum_steps
|
|
)
|
|
|
|
|
|
class TrainContextBuilder:
|
|
def __init__(
|
|
self,
|
|
config: TrainConfig,
|
|
):
|
|
self.config = config
|
|
self._param_path: Optional[str] = None
|
|
self._resume: bool = False
|
|
|
|
def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Self:
|
|
self._param_path = param_path
|
|
self._resume = resume
|
|
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,
|
|
)
|
|
|
|
model_config = {}
|
|
if self._param_path:
|
|
config_path = Path(self._param_path) / "config.json"
|
|
if config_path.exists():
|
|
model_config = load_json(config_path)
|
|
|
|
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
|
|
per_step = (
|
|
cfg.batch_per_device * get_world_size() * cfg.grad_accum_steps
|
|
)
|
|
preloaded_consumed = (
|
|
checkpoint.consumed_samples // per_step
|
|
) * per_step
|
|
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 cfg.lora is not None:
|
|
inject_lora(
|
|
m,
|
|
r=cfg.lora.r,
|
|
alpha=cfg.lora.alpha,
|
|
target_modules=set(cfg.lora.target_modules),
|
|
)
|
|
if preloaded_state_dict is not None:
|
|
m.load_state_dict(preloaded_state_dict, strict=False)
|
|
if cfg.compile_mode is not None:
|
|
logger.info("torch.compile enabled (mode=%s)", cfg.compile_mode)
|
|
m = torch.compile(m, mode=cfg.compile_mode)
|
|
return m
|
|
|
|
context = TrainContext(
|
|
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,
|
|
)
|
|
|
|
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
|
|
|
|
if val_dataset is None and cfg.val_split is not None:
|
|
n_total = len(cfg.dataset)
|
|
n_val = max(1, int(n_total * cfg.val_split))
|
|
n_train = n_total - n_val
|
|
generator = torch.Generator().manual_seed(cfg.random_seed)
|
|
train_dataset, val_dataset = random_split(
|
|
cfg.dataset, [n_train, n_val], generator=generator
|
|
)
|
|
|
|
sampler_offset = context.consumed_samples // context.world_size
|
|
|
|
if self._resume and sampler_offset > 0:
|
|
offset = context.world_size - 1
|
|
num_samples_per_replica = (
|
|
len(train_dataset) + offset
|
|
) // context.world_size
|
|
if num_samples_per_replica > 0:
|
|
context.epoch = sampler_offset // num_samples_per_replica
|
|
|
|
sampler = RDSampler(
|
|
data_source=train_dataset,
|
|
start_epoch=context.epoch,
|
|
start_iter=sampler_offset,
|
|
seed=cfg.random_seed,
|
|
)
|
|
context.dataloader = DataLoader(
|
|
train_dataset,
|
|
batch_size=cfg.batch_per_device,
|
|
sampler=sampler,
|
|
num_workers=cfg.num_workers,
|
|
pin_memory=cfg.pin_memory,
|
|
prefetch_factor=cfg.prefetch_factor,
|
|
collate_fn=cfg.collate_fn,
|
|
)
|
|
|
|
if val_dataset is not None:
|
|
val_sampler = RDSampler(
|
|
data_source=val_dataset,
|
|
start_epoch=0,
|
|
start_iter=0,
|
|
seed=cfg.random_seed,
|
|
shuffle=False,
|
|
)
|
|
context.val_dataloader = DataLoader(
|
|
val_dataset,
|
|
batch_size=cfg.batch_per_device,
|
|
sampler=val_sampler,
|
|
num_workers=cfg.num_workers,
|
|
pin_memory=cfg.pin_memory,
|
|
prefetch_factor=cfg.prefetch_factor,
|
|
collate_fn=cfg.collate_fn,
|
|
)
|
|
|
|
if context.checkpoint and context.checkpoint.extra:
|
|
extra = context.checkpoint.extra
|
|
for name in ("optimizer", "scheduler"):
|
|
if name in extra:
|
|
obj = getattr(context, name, None)
|
|
if obj is not None:
|
|
obj.load_state_dict(extra[name])
|
|
|
|
strategy_kwargs = dict(cfg.extra_kwargs)
|
|
|
|
needs_ref = cfg.strategy in (
|
|
"dpo",
|
|
"grpo",
|
|
"online_grpo",
|
|
"online_dpo",
|
|
)
|
|
needs_old = cfg.strategy in ("grpo", "online_grpo")
|
|
|
|
if needs_ref:
|
|
ref_model = create_ref_model(
|
|
cfg.model_fn, executor.unwrap_model(context.model)
|
|
).to(device=device)
|
|
strategy_kwargs["ref_model"] = ref_model
|
|
|
|
old_model = None
|
|
if needs_old:
|
|
old_model = create_ref_model(
|
|
cfg.model_fn, executor.unwrap_model(context.model)
|
|
).to(device=device)
|
|
strategy_kwargs["old_model"] = old_model
|
|
|
|
context.strategy = StrategyFactory.create(
|
|
cfg.strategy,
|
|
model=context.model,
|
|
device=device,
|
|
executor=executor,
|
|
**strategy_kwargs,
|
|
)
|
|
|
|
# Enable online rollout when the train_type is an ``online_*`` variant.
|
|
is_online = cfg.strategy.startswith("online_")
|
|
if is_online:
|
|
if not context.strategy.supports_online():
|
|
raise ValueError(
|
|
f"Strategy '{cfg.strategy}' does not support online rollout"
|
|
)
|
|
if cfg.reward_model_fn is None:
|
|
raise ValueError("reward_model_fn is required for online RL strategies")
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(self._param_path)
|
|
reward_model = cfg.reward_model_fn()
|
|
|
|
group_size = strategy_kwargs.get("group_size", 1)
|
|
rollout_batch_size = group_size * max(1, cfg.batch_per_device)
|
|
max_seq_len = getattr(context.model.config, "max_position_embeddings", None)
|
|
|
|
scheduler = InferenceScheduler(
|
|
model=context.model,
|
|
tokenizer=tokenizer,
|
|
max_batch_size=rollout_batch_size,
|
|
max_seq_len=max_seq_len,
|
|
)
|
|
|
|
generator = RolloutGenerator(
|
|
scheduler=scheduler,
|
|
tokenizer=tokenizer,
|
|
max_tokens=cfg.rollout_max_tokens,
|
|
group_size=group_size,
|
|
temperature=cfg.rollout_temperature,
|
|
top_k=cfg.rollout_top_k,
|
|
top_p=cfg.rollout_top_p,
|
|
)
|
|
runner = RolloutRunner(
|
|
generator=generator,
|
|
reward_model=reward_model,
|
|
rollout_interval=cfg.rollout_interval,
|
|
)
|
|
context.strategy.set_rollout_runner(runner)
|
|
|
|
return context
|