project-llm-trainer 0.12.3__py3-none-any.whl

llm_trainer/sft_trainer.py

@@ -0,0 +1,97 @@
+from typing import Optional, Tuple, List
+from torch.utils.data import Dataset
+
+from llm_model import (
+    VLMConfig,
+    LlmModel,
+    VlmModel
+)
+
+from .base_trainer import BaseTrainer
+from .train_configs import TrainConfig
+from .dataset import SFTDataset
+from .utils import get_sft_collate_fn
+from .tools import TrainerTools
+
+
+class SFTTrainer(BaseTrainer):
+    def __init__(
+            self,
+            *,
+            train_config: TrainConfig,
+            eval_prompts: List[str],
+            eval_image_tags: Optional[List[str]] = None
+    ):
+        self.sft_config = train_config.sft_config
+        self.pixel_values_provider = self.sft_config.pixel_values_provider
+        self.eval_image_tags = eval_image_tags
+
+        super().__init__(
+            train_config=train_config,
+            eval_prompts=eval_prompts,
+            kd_config=self.sft_config.kd_config,
+            gradient_accumulation_steps=self.sft_config.gradient_accumulation_steps
+        )
+
+        if isinstance(train_config.model_config, VLMConfig):
+            self.pixel_values_provider = self.sft_config.pixel_values_provider
+        else:
+            self.pixel_values_provider = None
+
+    def _new_model(self, train_config: TrainConfig):
+        if isinstance(train_config.model_config, VLMConfig):
+            return VlmModel(train_config.model_config)
+        else:
+            return LlmModel(train_config.model_config)
+
+    def _check_freeze_llm_model(self, model):
+        # freeze llm model for vlm training
+        if self.sft_config.freeze_llm_model:
+            for name, param in model.named_parameters():
+                if not any(sub_module in name for sub_module in ['multi_modal_projector']):
+                    param.requires_grad = False
+
+            # model.embed_tokens.eval()
+            # model.layers.eval()
+            # model.head_norm.eval()
+            # model.lm_head.eval()
+
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
+        sft_collate_fn = get_sft_collate_fn(self.sft_config.mask_prompt)
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs = super()._convert_train_args()
+        data_loader_kwargs.update({"collate_fn": sft_collate_fn})
+
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
+
+    def _get_pixel_values(self, batch_data):
+        if self.pixel_values_provider and 'image_tags' in batch_data:
+            image_tags = batch_data['image_tags']
+            return self.pixel_values_provider(image_tags).to(TrainerTools().parallel.device)
+
+        return None
+
+    def _get_eval_pixel_values_and_tokens_count(self, eval_idx):
+        if not self.eval_image_tags:
+            return None, None
+
+        eval_image_tag = self.eval_image_tags[eval_idx]
+        if isinstance(self.train_config.model_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
+            return self.pixel_values_provider([eval_image_tag]), self.train_config.model_config.tokens_per_image
+
+        return None, None
+
+    def _create_dataset(self, file_idx) -> Tuple[Dataset, str]:
+        file_path = self.train_config.file_dataset[file_idx]
+        max_seq_len = self.train_config.max_seq_len
+
+        image_tag_file_path = None
+        tokens_per_image = -1
+
+        if isinstance(self.train_config.model_config, VLMConfig):
+            if self.sft_config.image_tags_file_dataset:
+                image_tag_file_path = self.sft_config.image_tags_file_dataset[file_idx]
+
+            if self.train_config.model_config.tokens_per_image:
+                tokens_per_image = self.train_config.model_config.tokens_per_image
+
+        return SFTDataset(file_path, max_seq_len, image_tag_file_path, tokens_per_image), file_path

llm_trainer/tokenizer.py

@@ -0,0 +1,162 @@
+import os
+import warnings
+from typing import List, Dict, Union
+from transformers import AutoTokenizer
+import torch
+
+
+class Tokenizer:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(os.environ['TOKEN_DIR'])
+
+        self.text_end = '</s>'
+
+        self.text_pad = '<pad>'
+        self.text_unk = '<unk>'
+
+        self.text_user = '<user>'
+        self.text_assistant = '<assistant>'
+
+        self.text_think_start = '<think>'
+        self.text_think_end = '</think>'
+
+        self.text_answer_start = '<answer>'
+        self.text_answer_end = '</answer>'
+
+        self.text_system = '<system>'
+
+        self.text_image = '<image>'
+
+        self.end = self.tokenizer.convert_tokens_to_ids(self.text_end)
+
+        self.pad = self.tokenizer.convert_tokens_to_ids(self.text_pad)
+        self.unk = self.tokenizer.convert_tokens_to_ids(self.text_unk)
+
+        self.user = self.tokenizer.convert_tokens_to_ids(self.text_user)
+        self.assistant = self.tokenizer.convert_tokens_to_ids(self.text_assistant)
+
+        self.think_start = self.tokenizer.convert_tokens_to_ids(self.text_think_start)
+        self.think_end = self.tokenizer.convert_tokens_to_ids(self.text_think_end)
+
+        self.answer_start = self.tokenizer.convert_tokens_to_ids(self.text_answer_start)
+        self.answer_end = self.tokenizer.convert_tokens_to_ids(self.text_answer_end)
+
+        self.system = self.tokenizer.convert_tokens_to_ids(self.text_system)
+        self.image = self.tokenizer.convert_tokens_to_ids(self.text_image)
+
+        self.vocab_size = len(self.tokenizer)
+
+    def encode(
+            self,
+            text: str,
+            unsqueeze: bool = False,
+            covert_tensor: bool = False
+    ) -> Union[torch.Tensor, List[int]]:
+        # [x,x,x]
+        encoded = self.tokenizer.encode(text, add_special_tokens=False)
+
+        if unsqueeze:
+            # tensor: [[x,x,x]]
+            return torch.tensor(encoded, dtype=torch.long).unsqueeze(0)
+        else:
+            # tensor: # [x,x,x]
+            if covert_tensor:
+                return torch.tensor(encoded, dtype=torch.long)
+
+            return encoded
+
+    def batch_encode(
+            self,
+            text: List[str],
+            padding = False,
+            truncation = False,
+            covert_tensor: bool = False,
+            return_attention_mask: bool = False
+    ) -> Union[torch.Tensor, List[List[int]]]:
+        encoded = self.tokenizer(
+            text,
+            padding=padding,
+            truncation=truncation,
+            return_attention_mask=return_attention_mask
+        )['input_ids']
+
+        if covert_tensor:
+            encoded = torch.tensor(encoded, dtype=torch.long)
+
+        return encoded
+
+    def decode(
+            self,
+            token: Union[torch.Tensor, List[int]],
+            skip_special_tokens: bool = False
+    ) -> str:
+        return self.tokenizer.decode(token, skip_special_tokens=skip_special_tokens)
+
+    def batch_decode(
+            self,
+            tokens: Union[torch.Tensor, List[int], List[List[int]]],
+            skip_special_tokens: bool = False
+    ) -> List[str]:
+        return self.tokenizer.batch_decode(tokens, skip_special_tokens=skip_special_tokens)
+
+    def encode_to_token(self, text: str, unsqueeze=True, covert_tensor=True):
+        warnings.warn('encode_to_token is deprecated. Please use `encode` instead.')
+        return self.encode(text, unsqueeze, covert_tensor)
+
+    def decode_to_text(self, token: torch.Tensor, skip_special_tokens: bool = False) -> str:
+        warnings.warn('decode_to_text is deprecated. Please use `decode` instead.')
+        return self.decode(token.squeeze(0), skip_special_tokens)
+
+    def apply_chat_template(
+            self,
+            conversations: List[Dict[str, str]],
+            tokenizer: bool = True,
+            add_answer_tag_for_assistant: bool = True,
+            unsqueeze=False,
+            covert_tensor=False
+    ):
+        """
+            [
+                {"role":"system", "content":"system prompt"},
+                {"role":"user", "content":"hello?"},
+                {"role":"assistant", "content":"hello"},
+                {"role":"user", "content":"hello hello?"},
+                {"role":"assistant", "think":"thinking", "content":"hello hello"},
+            ]
+            <system>{system_prompt}</s><user>hello?</s><assistant>hello</s><user>hello hello?</s><assistant><think>thinking</think><answer>hello hello</answer></s>
+        """
+
+        chat_template = ''
+        support_roles = {'system': self.text_system, 'user': self.text_user, 'assistant': self.text_assistant}
+        for conversation in conversations:
+            role = conversation['role']
+            if role in support_roles:
+                content = conversation['content']
+                if add_answer_tag_for_assistant and role == 'assistant':
+                    content = f"{self.text_answer_start}{content}{self.text_answer_end}"
+
+                if 'think' in conversation:
+                    content = f"{self.text_think_start}{conversation['think']}{self.text_think_end}{content}"
+
+                chat_template = f"{chat_template}{support_roles[role]}{content}{self.text_end}"
+
+        if tokenizer:
+            return self.encode(chat_template, unsqueeze, covert_tensor)
+
+        return chat_template
+
+    def get_special_tokens_dict(self):
+        return {
+            self.text_end: self.end,
+            self.text_pad: self.pad,
+            self.text_unk: self.unk,
+            self.text_user: self.user,
+            self.text_assistant: self.assistant,
+            self.text_think_start: self.think_start,
+            self.text_think_end: self.think_end,
+            self.text_answer_start: self.answer_start,
+            self.text_answer_end: self.answer_end,
+            self.text_system: self.system,
+            self.text_image: self.image,
+        }
+

llm_trainer/tools.py

@@ -0,0 +1,116 @@
+import os
+from abc import ABC, abstractmethod
+import torch
+from .tokenizer import Tokenizer
+from .parallel import DsParallel, DdpParallel, NoneParallel
+from .log import Logger
+
+
+parallel_types = {
+    'ds': DsParallel,
+    'ddp': DdpParallel,
+    'none': NoneParallel
+}
+
+dtypes = {
+    'float': torch.float,
+    'float16': torch.float16,
+    'float32': torch.float32,
+    'float64': torch.float64
+}
+
+class TrainerTools:
+    def __init__(self):
+        if not hasattr(TrainerTools, "_first_init"):
+            TrainerTools._first_init = True
+
+            self.parallel = self._new_parallel()
+
+            self.tokenizer = Tokenizer()
+            self.use_amp = 'cuda' in self.parallel.device and not isinstance(self.parallel, DsParallel)
+
+            Logger.std_log(f'word_size={self.parallel.world_size}, use_amp={self.use_amp}')
+
+    def _new_parallel(self):
+        parallel_type = os.environ.get('PARALLEL_TYPE', 'none')
+        Logger.std_log(f'parallel_type={parallel_type}')
+        return parallel_types[parallel_type]()
+
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(TrainerTools, "_instance"):
+            TrainerTools._instance = object.__new__(cls)
+
+        return TrainerTools._instance
+
+
+class FileDataset(ABC):
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+    @abstractmethod
+    def __getitem__(self, idx) -> str: ...
+
+
+def estimate_data_size(
+        file_dataset: FileDataset,
+        max_seq_len: int,
+        type: str
+) -> int:
+    """
+    估计数据集大小
+    """
+    data_size = 0
+    files_count = len(file_dataset)
+
+    if type == 'sft':
+        from .dataset import SFTDataset
+        for idx in range(files_count):
+            dataset = SFTDataset(file_dataset[idx], max_seq_len)
+            data_size += len(dataset)
+    elif type == 'dpo':
+        from .dataset import DPODataset
+        for idx in range(files_count):
+            dataset = DPODataset(file_dataset[idx], max_seq_len)
+            data_size += len(dataset)
+    elif type == 'grpo' or type == 'ppo':
+        from .dataset import RLDataset
+        for idx in range(files_count):
+            dataset = RLDataset(file_dataset[idx])
+            data_size += len(dataset)
+    else:
+        from .dataset import PretrainDataset
+        for idx in range(files_count):
+            dataset = PretrainDataset(
+                file_dataset[idx],
+                max_seq_len,
+                max_seq_len
+            )
+            data_size += len(dataset)
+
+    return data_size
+
+
+def extract_policy_weights_from_ppo(model_config, ppo_weights):
+    from llm_model import LlmModel
+    from .ppo_trainer import PolicyAndValueModelWrapper, ValueModel
+
+    policy_model = LlmModel(model_config)
+    value_model = ValueModel(LlmModel(model_config))
+
+    wrapper = PolicyAndValueModelWrapper(policy_model, value_model)
+    wrapper.load_state_dict(ppo_weights)
+
+    return wrapper.policy_model.state_dict()
+
+
+def extract_value_weights_from_ppo(model_config, ppo_weights):
+    from llm_model import LlmModel
+    from .ppo_trainer import PolicyAndValueModelWrapper, ValueModel
+
+    policy_model = LlmModel(model_config)
+    value_model = ValueModel(LlmModel(model_config))
+
+    wrapper = PolicyAndValueModelWrapper(policy_model, value_model)
+    wrapper.load_state_dict(ppo_weights)
+
+    return wrapper.value_model.state_dict()

llm_trainer/train_configs.py

@@ -0,0 +1,324 @@
+from typing import Optional, Union, Callable, List, Mapping, Any, Tuple
+from dataclasses import dataclass, field
+
+import torch
+from llm_model import ModelConfig, VLMConfig
+from .tools import FileDataset
+
+
+@dataclass(kw_only=True)
+class DsOffloadConfig:
+    device: str = 'cpu'
+    pin_memory: bool = True
+
+
+@dataclass(kw_only=True)
+class DsActivationCheckpointingConfig:
+    partition_activations: bool = True
+    cpu_checkpointing: bool = False
+    contiguous_memory_optimization: bool = True
+    number_checkpoints: Optional[int] = None
+    synchronize_checkpoint_boundary: bool = False
+    profile: bool = False
+
+
+@dataclass(kw_only=True)
+class DsZeROConfig:
+    stage: int
+    allgather_partitions: Optional[bool] = True
+    allgather_bucket_size: Optional[int] = 5e8
+    overlap_comm: Optional[bool] = True
+    reduce_scatter: Optional[bool] = True
+    reduce_bucket_size: Optional[Union[str, int]] = 5e8
+    contiguous_gradients: Optional[bool] = True
+
+@dataclass(kw_only=True)
+class DsZero0Config(DsZeROConfig):
+    stage: int = field(default=0, init=False)
+
+
+@dataclass(kw_only=True)
+class DsZero1Config(DsZeROConfig):
+    stage: int = field(default=1, init=False)
+
+
+@dataclass(kw_only=True)
+class DsZero2Config(DsZeROConfig):
+    stage: int = field(default=2, init=False)
+    offload_optimizer: Optional[DsOffloadConfig] = None
+    offload_param: Optional[DsOffloadConfig] = None
+
+
+@dataclass(kw_only=True)
+class DsZero3Config(DsZeROConfig):
+    stage: int = field(default=3, init=False)
+    sub_group_size: Optional[int] = 1e9
+    stage3_prefetch_bucket_size: Optional[Union[str, int]] = 'auto'
+    stage3_param_persistence_threshold: Optional[Union[str, int]] = 'auto'
+    stage3_max_live_parameters: Optional[int] = 1e9
+    stage3_max_reuse_distance: Optional[int] = 1e9
+    stage3_gather_16bit_weights_on_model_save: Optional[bool] = True
+    offload_optimizer: Optional[DsOffloadConfig] = None
+    offload_param: Optional[DsOffloadConfig] = None
+
+
+@dataclass(kw_only=True)
+class DsFp16Config:
+    enabled: Union[str, bool] = 'auto'
+    loss_scale: int = 0
+    loss_scale_window: int = 1000
+    initial_scale_power: int = 16
+    hysteresis: int = 2
+    min_loss_scale: int = 1
+    fp16_opt_level: Optional[str] = 'O2'
+
+
+@dataclass(kw_only=True)
+class DsBf16Config:
+    enabled: bool = True
+
+
+@dataclass(kw_only=True)
+class DsConfig:
+    zero_config: Optional[DsZeROConfig] = field(default_factory=DsZero3Config)
+    fp16_config: Optional[DsFp16Config] = field(default_factory=DsFp16Config)
+    bf16_config: Optional[DsBf16Config] = field(default_factory=DsBf16Config)
+    gradient_clipping: Optional[float] = 1.0
+    activation_checkpointing: Optional[DsActivationCheckpointingConfig] = None
+
+
+@dataclass(kw_only=True)
+class DataLoaderConfig:
+    """
+        data loader配置项
+        Args:
+            data_loader_pin_memory (`bool`, *optional*, default is None):
+                data_loader pin_memory config
+            data_loader_num_workers (`int`, *optional*, default is 0):
+                data_loader num_workers config
+            data_loader_shuffle (`bool`, *optional*, default is False):
+                是否需要shuffle数据
+            data_loader_drop_last (`bool`, default is False):
+                最后一个batch不满足batch_size时，是否丢弃
+    """
+    data_loader_pin_memory: bool = False
+    data_loader_num_workers: int = 0
+    data_loader_shuffle: bool = False
+    data_loader_drop_last: bool = True
+
+
+@dataclass(kw_only=True)
+class OptimConfig:
+    optim_type: str = 'adam' # or 'lion'
+    enable_lr_scheduler: bool = False
+    initial_lr: float
+    weight_decay: Optional[float] = None
+    betas: Optional[Tuple[float, float]] = None
+    warmup_iters: Optional[int] = None
+    max_lr: Optional[float] = None
+    min_lr: Optional[float] = None
+    cosine_annealing_period: Optional[int] = None
+    cosine_annealing_period_mul: int = 0
+
+
+@dataclass(kw_only=True)
+class LossConfig:
+    critical_tokens: Optional[List[int]] = None
+    critical_alpha: float = 1.0
+    aux_loss_coef: Optional[float] = 0.001
+
+
+@dataclass(kw_only=True)
+class KDConfig:
+    """
+        知识蒸馏模式配置项
+
+        Args:
+            teacher_logits_provider (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):
+                知识蒸馏教师模型logits的提供者
+            kd_coef (`float`, *optional*, default is 0.4):
+                蒸馏loss的占比，loss = kd_coef * kd_loss + (1 - kd_coef) * lm_loss
+    """
+    teacher_logits_provider: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+    kd_coef: float = 0.4
+
+
+@dataclass(kw_only=True)
+class EvalConfig:
+    """
+    训练参数配置项
+
+    Args:
+        eval_batch_interval (`int`, default is 100):
+            每隔多少个batch进行模型eval
+    """
+    max_new_tokens: int
+    eval_batch_interval: int = 100
+    temperature: float = 1.0
+    top_p: float = 0.95
+    top_k: Optional[float] = None
+
+
+@dataclass(kw_only=True)
+class PretrainConfig:
+    """
+    训练参数配置项
+
+    Args:
+        gradient_accumulation_steps (`int`, *Optional*, default is 1):
+            梯度累积步数，为0时不使用梯度累积
+            目前仅适用于pretrain\sft\dpo，不适用于ppo\grpo\gspo
+        kd_config: (`KDConfig`, *Optional*, default is None):
+            知识蒸馏配置项，为None时不使用知识蒸馏
+    """
+    gradient_accumulation_steps: int = 1
+    kd_config: Optional[KDConfig] = None
+
+
+@dataclass(kw_only=True)
+class SFTConfig:
+    """
+    训练参数配置项
+
+    Args:
+        mask_prompt (`bool`)
+            指定是否mask prompt部分的token
+        gradient_accumulation_steps (`int`, *Optional*, default is 1):
+            梯度累积步数，为0时不使用梯度累积
+            目前仅适用于pretrain\sft\dpo，不适用于ppo\grpo\gspo
+        kd_config: (`KDConfig`, *Optional*, default is None):
+            知识蒸馏配置项，为None时不使用知识蒸馏
+        pixel_values_provider: (`Callable[[list[str]], torch.Tensor]`, *Optional*, default is None):
+            训练vlm时根据image_tag提供pixel_values信息
+        freeze_llm_model:
+            是否冻结llm部分model参数，用于训练vlm
+    """
+    mask_prompt: bool = True
+    gradient_accumulation_steps: int = 1
+    kd_config: Optional[KDConfig] = None
+    image_tags_file_dataset: Optional[FileDataset] = None
+    pixel_values_provider: Optional[Callable[[list[str]], torch.Tensor]] = None
+    freeze_llm_model: bool = False
+
+
+@dataclass(kw_only=True)
+class DPOConfig:
+    """
+    训练参数配置项
+
+    Args:
+        mask_prompt (`bool`)
+            指定是否mask prompt部分的token
+        gradient_accumulation_steps (`int`, *Optional*, default is 1):
+            梯度累积步数，为0时不使用梯度累积
+            目前仅适用于pretrain\sft\dpo，不适用于ppo\grpo\gspo
+    """
+    ref_model_checkpoint: Mapping[str, Any]
+    mask_prompt: bool = True
+    gradient_accumulation_steps: int = 1
+    loss_beta: float
+    loss_label_smoothing: float = 0.0
+    loss_ipo: bool = False
+    nll_loss_coef: Optional[float] = None
+
+
+@dataclass(kw_only=True)
+class PPOConfig:
+    ppo_epochs: int
+    ppo_batch_size: int
+    ref_model_checkpoint: Mapping[str, Any]
+    value_model_checkpoint: Optional[Mapping[str, Any]] = None
+    gradient_accumulation_steps: int = 1
+    gamma: float = 1.0
+    lam: float = 0.95
+    clip_eps: float = 0.1
+    vf_coef: float = 0.5
+    kl_beta: float = 0.02
+    kl_estimator: str = 'k1' # or k3
+    missing_eos_penalty: Optional[float] = None
+    normalize_rewards: bool = False
+    whiten_rewards: bool = False
+    gen_max_new_tokens: int
+    gen_temperature: Optional[float] = None
+    gen_k: Optional[int] = None
+    gen_p: Optional[float] = None
+    gen_suppress_tokens: Optional[list[int]] = None
+
+
+@dataclass(kw_only=True)
+class GRPOConfig:
+    grpo_steps: int = 1
+    group_size: int = 12
+    mixup_alpha: float = 1.0
+    loss_beta: float = 0.0 # or 0.04 for grpo
+    loss_clip_eps: float = 3e-4
+    loss_clip_eps_high: Optional[float] = 4e-4
+    loss_delta: Optional[float] = None
+    loss_importance_sampling_level: str = 'seq' # token or seq
+    loss_type: str = 'grpo' # grpo or bnpo or dr_grpo
+    gen_max_new_tokens: int
+    gen_temperature: Optional[float] = None
+    gen_k: Optional[int] = None
+    gen_p: Optional[float] = None
+    gen_suppress_tokens: Optional[list[int]] = None
+
+
+@dataclass(kw_only=True)
+class TrainConfig:
+    """
+    训练参数配置项
+
+    Args:
+        n_epochs (`int`):
+            训练epochs
+        batch_size (`int`):
+            每个batch的大小
+        model_config (`ModelConfig`):
+            模型的配置
+        init_state_dict:
+            初始化检查点
+        file_dataset (`FileDataset`):
+            训练文件dataset
+        max_seq_len (`int`, default is None)
+            训练序列最大长度，为None时取model的max_position_embedding
+        data_loader_config: (`DataLoaderConfig`):
+            data loader配置项
+        loss_config:
+            配置loss
+        ds_config:
+            配置deepspeed
+        eval_config:
+            配置eval
+        optim_config (`OptimConfig`):
+            optim配置项
+        pretrain_config:
+            预训练配置项，仅适用于使用Trainer
+        sft_config:
+            sft配置项，仅适用于使用SFTTrainer
+        dpo_config:
+            dpo配置项，仅适用于使用DPOTrainer
+        ppo_config:
+            ppo配置项，仅适用于使用PPOTrainer
+        grpo_config:
+            grpo配置项，仅适用于使用GRPOTrainer
+    """
+    n_epochs: int
+    batch_size: int
+    model_config: Union[ModelConfig, VLMConfig]
+    init_state_dict: Optional[Mapping[str, Any]] = None
+
+    file_dataset: FileDataset
+    max_seq_len: int
+    data_loader_config: DataLoaderConfig = field(default_factory=DataLoaderConfig)
+
+    loss_config: LossConfig = field(default_factory=LossConfig)
+    optim_config: OptimConfig = field(default_factory=OptimConfig)
+    ds_config: DsConfig = field(default_factory=DsConfig)
+
+    eval_config: EvalConfig = field(default_factory=EvalConfig)
+
+    pretrain_config: Optional[PretrainConfig] = None
+    sft_config: Optional[SFTConfig] = None
+    dpo_config: Optional[DPOConfig] = None
+    ppo_config: Optional[PPOConfig] = None
+    grpo_config: Optional[GRPOConfig] = None