arguments.py

from dataclasses import dataclass, field
from typing import List, Optional

from transformers.training_args import TrainingArguments


@dataclass
class CustomTrainingArguments(TrainingArguments):
    use_wandb: bool = field(
        default=False,
        metadata={
            "help": "whether to use wandb to log the training process.",
            "choices": [True, False],
        },
    )

    adapter_path: str = field(
        default="",
        metadata={"help": "the path to load the adapter."},
    )

    # tokenizer params
    padding_side: str = field(
        default="right",
        metadata={"help": "the direction for tokenizer to add padding tokens."},
    )

    truncation_side: str = field(
        default="left",
        metadata={"help": "the direction for tokenizer to add padding tokens."},
    )

    add_sep_token: bool = field(
        default=False,
        metadata={"help": "whether add a <sep> token between query and response."},
    )

    # model params
    model_type: str = field(
        default="llama",
        metadata={
            "help": "the base model type for reward model, selected from [llama, bert]."
        },
    )

    model_prefix: str = field(
        default="llama",
        metadata={
            "help": "the base model type for reward model, selected from [llama, bert]."
        },
    )

    pooling_type: str = field(
        default="average",
        metadata={
            "help": "the pooling method for reward model, selected from [average, max, last]."
        },
    )

    model_name_or_path: str = field(
        default="/data/models/Llama-3.2-3B-Instruct",
        metadata={"help": "the path to load pretrained model."},
    )

    ref_model_name_or_path: str = field(
        default="",
        metadata={"help": "the path to load pretrained model."},
    )

    critic_model_name_or_path: str = field(
        default="FacebookAI/roberta-base",
        metadata={"help": "the path to load pretrained critic model."},
    )

    # data params
    game_name: str = field(
        default="competitive_taboo",
        metadata={
            "help": "the name of the task.",
            "choices": [
                "competitive_taboo",
                "cooperative_rsa",
                "strategic",
                "20q",
                "guessmycity",
            ],
        },
    )

    game_max_turn: int = field(
        default=6,
        metadata={"help": "the max_turn to play the adversarial taboo game."},
    )

    data_dir: str = field(
        default="path/to/cleaned_data", metadata={"help": "the directory to load data."}
    )

    data_type: str = field(default="no_type", metadata={"help": "the type of data."})

    data_path: str = field(
        default="yahma/alpaca-cleaned", metadata={"help": "the path to load data."}
    )

    train_data_path: List[str] = field(
        default_factory=lambda: [], metadata={"help": "train datasets paths."}
    )

    eval_data_path: List[str] = field(
        default_factory=lambda: [], metadata={"help": "evaluation datasets paths."}
    )

    data_prefix: str = field(
        default="yahma/alpaca-cleaned",
        metadata={"help": "the prefix to load train and test data."},
    )

    data_suffix: str = field(
        default="yahma/alpaca-cleaned",
        metadata={"help": "the suffix to save inference data."},
    )

    # training hyperparams
    task_type: str = field(default="training", metadata={"help": "the task type"})

    train_method: str = field(
        default="SFTWeightedWithKL",
        metadata={
            "help": "the LLM training method name.",
            "choices": [
                "SFTWeightedWithKL",
                "SelfPlayPPO",
                "SelfPlayFoPO",
                "SelfPlayGRPO",
                "SelfPlayGRFoPO",
                "SelfPlayOnlinePPO",
                "SelfPlayOnlineGRPO",
                "SelfPlayOnlineFoPO",
                "SelfPlayOnlineGRFoPO",
            ],
        },
    )

    use_lora: bool = field(
        default=False,
        metadata={"help": "whether use lora to train the model."},
    )

    debug_mode: bool = field(
        default=False, metadata={"help": "whether use the debug mode."}
    )

    cache_dir: Optional[str] = field(default=None)

    clip_range: float = field(
        default=0.2,
        metadata={
            "help": "the range to clip the importance reweighting ratio for policy optimization."
        },
    )

    length_penalty: float = field(
        default=1.0, metadata={"help": "the penalty for seq length."}
    )

    lm_sft_coeff: float = field(
        default=0.0,
        metadata={"help": "the coefficient for SFT data language modeling loss."},
    )

    lm_kl_coeff: float = field(
        default=0.0, metadata={"help": "the coefficient of kl regularizer."}
    )

    max_length: int = field(
        default=256, metadata={"help": "the max sentence sequence length."}
    )

    valid_data_size: int = field(
        default=0, metadata={"help": "the data size for validation data"}
    )

    resume_from_checkpoint: Optional[str] = field(
        default=None, metadata={"help": "either training checkpoint or final adapter"}
    )

    rollout_size: int = field(
        default=128,  # the original value is 256
        metadata={"help": "the number of rollouts for PPO training, default is 64."},
    )

    replay_buffer_size: int = field(
        default=10000,
        metadata={
            "help": "the size of the replay buffer for PPO training, default is 10000."
        },
    )

    replay_batch_size: int = field(
        default=16,  # the original value is 256
        metadata={
            "help": "the batch size for sampling from the replay buffer, default is 16."
        },
    )

    critic_learning_rate: float = field(
        default=2e-5, metadata={"help": "the learning rate for critic model."}
    )

    gamma: float = field(default=0.99, metadata={"help": "the discount factor for RL."})

    tau: float = field(
        default=0.95, metadata={"help": "the decay factor for critic model."}
    )

    # generation parameters:
    max_new_tokens: int = field(
        default=128, metadata={"help": "the max sentence sequence length."}
    )

    temperature: float = field(
        default=0.9, metadata={"help": "the temperature for sampling."}
    )

    top_p: float = field(default=0.95, metadata={"help": "the top_p for sampling."})

    player_one_model_name_or_path: str = field(
        default="",
        metadata={"help": "the path to load the attacker or speaker model."},
    )

    player_two_model_name_or_path: str = field(
        default="",
        metadata={"help": "the path to load the defender or listener model."},
    )

    # GRPO 特有参数
    group_size: int = field(
        default=4, metadata={"help": "Number of responses per group for GRPO"}
    )

    group_sampling_strategy: str = field(
        default="random",
        metadata={"help": "Strategy for sampling groups: random, diverse, etc."},
    )