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
This commit is contained in:
2026-07-30 08:41:14 +08:00
parent 164be9708b
commit 22cf798d81
3 changed files with 179 additions and 1 deletions
+32
View File
@@ -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