215 lines
9.5 KiB
Python
215 lines
9.5 KiB
Python
from dataclasses import field
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
import torch.nn as nn
|
|
from pydantic import ConfigDict, field_validator, model_validator
|
|
from pydantic.dataclasses import dataclass
|
|
from torch.optim import Optimizer
|
|
from torch.optim.lr_scheduler import LRScheduler
|
|
from torch.utils.data import Dataset
|
|
|
|
from astrai.config.base import BaseConfig
|
|
from astrai.model.components.lora import LoRAConfig
|
|
|
|
_TRAIN_TYPES = frozenset({"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"})
|
|
_PARALLEL_MODES = frozenset({"none", "ddp", "fsdp"})
|
|
_BACKENDS = frozenset({"nccl", "gloo"})
|
|
_START_METHODS = frozenset({"spawn", "fork", "forkserver"})
|
|
_COMPILE_MODES = frozenset({"default", "reduce-overhead", "max-autotune"})
|
|
|
|
|
|
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
|
|
class TrainConfig(BaseConfig):
|
|
"""Training configuration.
|
|
|
|
Combines hyperparameters with runtime objects (model_fn, dataset, etc.).
|
|
Only JSON-serializable fields are written to checkpoint meta via to_dict().
|
|
|
|
Args:
|
|
model_fn (Callable[[], nn.Module]): Model factory for training.
|
|
strategy (str): Training strategy (seq, sft, dpo, grpo, online_*).
|
|
dataset (Dataset): Dataset for training.
|
|
optimizer_fn (Callable[[nn.Module], Optimizer]): Optimizer factory for training.
|
|
optimizer_name (Optional[str]): Serializable built-in optimizer identifier. Defaults to None.
|
|
optimizer_hyperparameters (Dict[str, Any]): Serializable optimizer settings. Defaults to {}.
|
|
scheduler_fn (Callable[[Optimizer], LRScheduler]): Scheduler factory for training.
|
|
n_epoch (int): Number of epochs for training. Defaults to 1.
|
|
batch_per_device (int): Batch size per device. Defaults to 4.
|
|
grad_accum_steps (int): Number of iterations between optimizer steps. Defaults to 1.
|
|
max_grad_norm (Optional[float]): Maximum gradient norm. None disables clipping. Defaults to 1.0.
|
|
gradient_checkpointing_modules (List[type]): Module types to enable activation checkpointing for. Defaults to [].
|
|
compile_mode (Optional[str]): torch.compile mode: 'default', 'reduce-overhead', 'max-autotune', or None. Defaults to None.
|
|
start_epoch (int): Start epoch for training. Defaults to 0.
|
|
start_samples (int): Start samples count (per rank). Superseded by checkpoint consumed_samples. Defaults to 0.
|
|
ckpt_dir (str): Checkpoint directory. Defaults to "./checkpoint".
|
|
ckpt_interval (int): Number of optimizer steps between checkpoints. Defaults to 5000.
|
|
lora (Optional[LoRAConfig]): LoRA config. None means full fine-tuning. Defaults to None.
|
|
metrics (List[str]): Metrics to record during training. Defaults to ["loss", "lr", "grad_norm"].
|
|
random_seed (int): Random seed. Defaults to 3407.
|
|
num_workers (int): Number of workers for dataloader. Defaults to 0.
|
|
prefetch_factor (Optional[int]): Prefetch factor for dataloader. Defaults to None.
|
|
pin_memory (bool): Pin memory for dataloader. Defaults to False.
|
|
collate_fn (Optional[Callable[[List[Any]], Any]]): Collate function for dataloader (e.g. dpo_collate_fn). Defaults to None.
|
|
nprocs (int): Number of processes for distributed training. Defaults to 1.
|
|
backend (str): Distributed training backend. Defaults to "nccl".
|
|
master_addr (str): Master address for distributed training. Defaults to "localhost".
|
|
master_port (str): Master port for distributed training. Defaults to "29500".
|
|
parallel_mode (str): Parallel strategy: none, ddp, fsdp. Defaults to "none".
|
|
start_method (str): Multiprocessing start method: spawn/fork/forkserver. Defaults to "spawn".
|
|
device_type (str): Device type for distributed training. Defaults to "cuda".
|
|
val_dataset (Optional[Dataset]): Dataset for validation. Defaults to None.
|
|
val_split (Optional[float]): Ratio to split from training dataset for validation, e.g. 0.05. Defaults to None.
|
|
val_step (int): Number of optimizer steps between validation runs. Defaults to 1000.
|
|
neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
|
|
rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512.
|
|
rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7.
|
|
rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0.
|
|
rollout_top_p (float): Top-p (nucleus) filtering for online rollout. Defaults to 0.9.
|
|
rollout_max_tokens (int): Maximum generated tokens per response in rollout. Defaults to 1024.
|
|
reward_model_fn (Optional[Callable]): Factory for reward model, required for online RL strategies. Defaults to None.
|
|
executor_kwargs (Dict[str, Any]): Extra kwargs passed to ExecutorFactory.create(). Defaults to {}.
|
|
extra_kwargs (Dict[str, Any]): Other arguments. Defaults to {}.
|
|
"""
|
|
|
|
model_fn: Callable[[], nn.Module]
|
|
strategy: str
|
|
dataset: Dataset
|
|
optimizer_fn: Callable[[nn.Module], Optimizer]
|
|
scheduler_fn: Callable[[Optimizer], LRScheduler]
|
|
optimizer_name: Optional[str] = None
|
|
optimizer_hyperparameters: Dict[str, Any] = field(default_factory=dict)
|
|
n_epoch: int = 1
|
|
batch_per_device: int = 4
|
|
grad_accum_steps: int = 1
|
|
max_grad_norm: Optional[float] = 1.0
|
|
gradient_checkpointing_modules: List[type] = field(default_factory=list)
|
|
compile_mode: Optional[str] = None
|
|
|
|
start_epoch: int = 0
|
|
start_samples: int = 0
|
|
ckpt_dir: str = "./checkpoint"
|
|
ckpt_interval: int = 5000
|
|
|
|
lora: Optional[LoRAConfig] = None
|
|
|
|
metrics: List[str] = field(default_factory=lambda: ["loss", "lr", "grad_norm"])
|
|
|
|
random_seed: int = 3407
|
|
num_workers: int = 0
|
|
prefetch_factor: Optional[int] = None
|
|
pin_memory: bool = False
|
|
collate_fn: Optional[Callable[[List[Any]], Any]] = None
|
|
|
|
nprocs: int = 1
|
|
backend: str = "nccl"
|
|
master_addr: str = "localhost"
|
|
master_port: str = "29500"
|
|
parallel_mode: str = "none"
|
|
start_method: str = "spawn"
|
|
|
|
device_type: str = "cuda"
|
|
val_dataset: Optional[Dataset] = None
|
|
val_split: Optional[float] = None
|
|
val_step: int = 1000
|
|
neftune_alpha: float = 0.0
|
|
|
|
rollout_interval: int = 512
|
|
rollout_temperature: float = 0.7
|
|
rollout_top_k: int = 0
|
|
rollout_top_p: float = 0.9
|
|
rollout_max_tokens: int = 1024
|
|
reward_model_fn: Optional[Callable] = None
|
|
|
|
executor_kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
extra_kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
@field_validator("strategy")
|
|
def _validate_strategy(cls, v: str) -> str:
|
|
if v not in _TRAIN_TYPES:
|
|
raise ValueError(
|
|
f"strategy must be one of {sorted(_TRAIN_TYPES)}, got {v!r}"
|
|
)
|
|
return v
|
|
|
|
@field_validator("parallel_mode")
|
|
def _validate_parallel_mode(cls, v: str) -> str:
|
|
if v not in _PARALLEL_MODES:
|
|
raise ValueError(
|
|
f"parallel_mode must be one of {sorted(_PARALLEL_MODES)}, got {v!r}"
|
|
)
|
|
return v
|
|
|
|
@field_validator("backend")
|
|
def _validate_backend(cls, v: str) -> str:
|
|
if v not in _BACKENDS:
|
|
raise ValueError(f"backend must be one of {sorted(_BACKENDS)}, got {v!r}")
|
|
return v
|
|
|
|
@field_validator("start_method")
|
|
def _validate_start_method(cls, v: str) -> str:
|
|
if v not in _START_METHODS:
|
|
raise ValueError(
|
|
f"start_method must be one of {sorted(_START_METHODS)}, got {v!r}"
|
|
)
|
|
return v
|
|
|
|
@field_validator("compile_mode")
|
|
def _validate_compile_mode(cls, v: Optional[str]) -> Optional[str]:
|
|
if v is not None and v not in _COMPILE_MODES:
|
|
raise ValueError(
|
|
f"compile_mode must be one of {sorted(_COMPILE_MODES)} or None, got {v!r}"
|
|
)
|
|
return v
|
|
|
|
@field_validator(
|
|
"n_epoch",
|
|
"batch_per_device",
|
|
"grad_accum_steps",
|
|
"ckpt_interval",
|
|
"val_step",
|
|
"rollout_interval",
|
|
"rollout_max_tokens",
|
|
)
|
|
def _validate_positive_int(cls, v: int) -> int:
|
|
if v <= 0:
|
|
raise ValueError(f"must be positive, got {v}")
|
|
return v
|
|
|
|
@field_validator("rollout_temperature")
|
|
def _validate_positive_float(cls, v: float) -> float:
|
|
if v <= 0:
|
|
raise ValueError(f"must be positive, got {v}")
|
|
return v
|
|
|
|
@field_validator("rollout_top_p")
|
|
def _validate_top_p(cls, v: float) -> float:
|
|
if not 0 < v <= 1:
|
|
raise ValueError(f"rollout_top_p must be in (0, 1], got {v}")
|
|
return v
|
|
|
|
@field_validator("rollout_top_k", "num_workers", "neftune_alpha")
|
|
def _validate_non_negative(cls, v):
|
|
if v < 0:
|
|
raise ValueError(f"must be non-negative, got {v}")
|
|
return v
|
|
|
|
@field_validator("max_grad_norm")
|
|
def _validate_max_grad_norm(cls, v: Optional[float]) -> Optional[float]:
|
|
if v is not None and v <= 0:
|
|
raise ValueError(f"max_grad_norm must be positive or None, got {v}")
|
|
return v
|
|
|
|
@field_validator("val_split")
|
|
def _validate_val_split(cls, v: Optional[float]) -> Optional[float]:
|
|
if v is not None and not 0 < v < 1:
|
|
raise ValueError(f"val_split must be in (0, 1) or None, got {v}")
|
|
return v
|
|
|
|
@model_validator(mode="after")
|
|
def _validate_online_strategy(self) -> "TrainConfig":
|
|
if self.strategy.startswith("online_") and self.reward_model_fn is None:
|
|
raise ValueError(
|
|
f"reward_model_fn is required for online RL strategy {self.strategy!r}"
|
|
)
|
|
return self
|