From 22cf798d819623982aed2e1db1dbf364f699c166 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 30 Jul 2026 08:40:41 +0800 Subject: [PATCH] feat: add field and model validators to config classes - TrainConfig: enum validators (strategy, parallel_mode, backend, start_method, compile_mode), positive/non-negative/range validators, model_validator requiring reward_model_fn for online RL strategies - AutoRegressiveLMConfig/EncoderConfig: attn_type, ffn_type enum validators - ProcessingConfig: packing_strategy, truncation_mode enums, positive int validators - OutputConfig: storage_format, position_ids_mode enum validators --- astrai/config/model_config.py | 32 ++++++++++ astrai/config/preprocess_config.py | 50 +++++++++++++++ astrai/config/train_config.py | 98 +++++++++++++++++++++++++++++- 3 files changed, 179 insertions(+), 1 deletion(-) diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py index cc6f414..0aefffe 100644 --- a/astrai/config/model_config.py +++ b/astrai/config/model_config.py @@ -1,10 +1,14 @@ from typing import Any, Dict, Optional +from pydantic import field_validator from pydantic.dataclasses import dataclass from astrai.config.base import BaseConfig from astrai.factory import BaseFactory +_ATTN_TYPES = frozenset({"gqa", "mla"}) +_FFN_TYPES = frozenset({"mlp", "moe"}) + class ConfigFactory(BaseFactory[BaseConfig]): """Factory that dispatches config classes by ``model_type``.""" @@ -84,6 +88,20 @@ class AutoRegressiveLMConfig(BaseModelConfig): n_activated_experts: Optional[int] = None topk_method: Optional[str] = None + @field_validator("attn_type") + def _validate_attn_type(cls, v: str) -> str: + if v not in _ATTN_TYPES: + raise ValueError( + f"attn_type must be one of {sorted(_ATTN_TYPES)}, got {v!r}" + ) + return v + + @field_validator("ffn_type") + def _validate_ffn_type(cls, v: str) -> str: + if v not in _FFN_TYPES: + raise ValueError(f"ffn_type must be one of {sorted(_FFN_TYPES)}, got {v!r}") + return v + @dataclass @ConfigFactory.register("embedding") @@ -127,3 +145,17 @@ class EncoderConfig(BaseModelConfig): ffn_type: str = "mlp" pooling_type: Optional[str] = None normalize_embeddings: Optional[bool] = None + + @field_validator("attn_type") + def _validate_attn_type(cls, v: str) -> str: + if v not in _ATTN_TYPES: + raise ValueError( + f"attn_type must be one of {sorted(_ATTN_TYPES)}, got {v!r}" + ) + return v + + @field_validator("ffn_type") + def _validate_ffn_type(cls, v: str) -> str: + if v not in _FFN_TYPES: + raise ValueError(f"ffn_type must be one of {sorted(_FFN_TYPES)}, got {v!r}") + return v diff --git a/astrai/config/preprocess_config.py b/astrai/config/preprocess_config.py index 28abcbe..48bba4d 100644 --- a/astrai/config/preprocess_config.py +++ b/astrai/config/preprocess_config.py @@ -8,10 +8,16 @@ modes, both driven declaratively through ``input.sections`` or from dataclasses import field from typing import Dict, List, Optional +from pydantic import field_validator from pydantic.dataclasses import dataclass from astrai.config.base import BaseConfig +_PACKING_STRATEGIES = frozenset({"simple", "bfd", "bfd_split"}) +_TRUNCATION_MODES = frozenset({"keep_start", "keep_end"}) +_STORAGE_FORMATS = frozenset({"bin", "jsonl"}) +_POSITION_IDS_MODES = frozenset({"none", "doc_reset", "continuous"}) + @dataclass class InputConfig(BaseConfig): @@ -61,6 +67,34 @@ class ProcessingConfig(BaseConfig): max_packed_len: int = 8192 truncation_mode: str = "keep_start" + @field_validator("packing_strategy") + def _validate_packing_strategy(cls, v: str) -> str: + if v not in _PACKING_STRATEGIES: + raise ValueError( + f"packing_strategy must be one of {sorted(_PACKING_STRATEGIES)}, got {v!r}" + ) + return v + + @field_validator("truncation_mode") + def _validate_truncation_mode(cls, v: str) -> str: + if v not in _TRUNCATION_MODES: + raise ValueError( + f"truncation_mode must be one of {sorted(_TRUNCATION_MODES)}, got {v!r}" + ) + return v + + @field_validator("max_seq_len", "batch_size", "max_packed_len") + def _validate_positive_int(cls, v: int) -> int: + if v <= 0: + raise ValueError(f"must be positive, got {v}") + return v + + @field_validator("min_chars") + def _validate_non_negative(cls, v: int) -> int: + if v < 0: + raise ValueError(f"min_chars must be non-negative, got {v}") + return v + @dataclass class OutputConfig(BaseConfig): @@ -80,6 +114,22 @@ class OutputConfig(BaseConfig): dtype: Dict[str, str] = field(default_factory=dict) position_ids_mode: str = "doc_reset" + @field_validator("storage_format") + def _validate_storage_format(cls, v: str) -> str: + if v not in _STORAGE_FORMATS: + raise ValueError( + f"storage_format must be one of {sorted(_STORAGE_FORMATS)}, got {v!r}" + ) + return v + + @field_validator("position_ids_mode") + def _validate_position_ids_mode(cls, v: str) -> str: + if v not in _POSITION_IDS_MODES: + raise ValueError( + f"position_ids_mode must be one of {sorted(_POSITION_IDS_MODES)}, got {v!r}" + ) + return v + @dataclass class PipelineConfig(BaseConfig): diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index d0e2e45..fadc973 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -2,7 +2,7 @@ from dataclasses import field from typing import Any, Callable, Dict, List, Optional import torch.nn as nn -from pydantic import ConfigDict +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 @@ -11,6 +11,12 @@ 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): @@ -112,3 +118,93 @@ class TrainConfig(BaseConfig): 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