471 lines
30 KiB
Python
471 lines
30 KiB
Python
from dataclasses import dataclass, field
|
||
from typing import List, Literal, Optional, Union
|
||
|
||
from swift.utils import json_parse_to_dict
|
||
|
||
|
||
@dataclass
|
||
class VllmArguments:
|
||
"""VllmArguments is a dataclass that holds the configuration for vllm.
|
||
|
||
Args:
|
||
vllm_gpu_memory_utilization (float): GPU memory utilization. Default is 0.9.
|
||
vllm_tensor_parallel_size (int): Tensor parallelism size. Default is 1.
|
||
vllm_pipeline_parallel_size (int): Pipeline parallelism size. Default is 1.
|
||
vllm_enable_expert_parallel (bool): Flag to enable expert parallelism for MoE models. Default is False.
|
||
vllm_max_num_seqs (int): Maximum number of sequences. Default is 256.
|
||
vllm_max_model_len (Optional[int]): Maximum model length. Default is None.
|
||
vllm_disable_custom_all_reduce (bool): Flag to disable custom all-reduce. Default is True.
|
||
vllm_enforce_eager (bool): Flag to enforce eager execution. Default is False.
|
||
vllm_limit_mm_per_prompt (Optional[str]): Limit multimedia per prompt. Default is None.
|
||
vllm_max_lora_rank (int): Maximum LoRA rank. Default is 16.
|
||
vllm_enable_prefix_caching (Optional[bool]): Flag to enable automatic prefix caching. Default is None.
|
||
vllm_use_async_engine (Optional[bool]): Whether to use async engine for vLLM. Default is None,
|
||
which will be set to True for encode tasks (embedding, seq_cls, reranker, generative_reranker),
|
||
deployment scenarios (swift deploy) and False otherwise.
|
||
vllm_quantization (Optional[str]): The quantization method for vLLM. Default is None.
|
||
vllm_reasoning_parser (Optional[str]): The reasoning parser for vLLM. Default is None.
|
||
vllm_disable_cascade_attn (bool): Flag to disable cascade attention. Default is False.
|
||
vllm_mm_processor_cache_gb (Optional[float]): MM processor cache size in GB. Default is None.
|
||
vllm_speculative_config (Optional[Union[dict, str]]): Speculative decoding configuration, passed in as a JSON
|
||
string. Defaults to None.
|
||
vllm_engine_kwargs (Optional[Union[dict, str]]): Additional parameters for vllm, formatted as a JSON string.
|
||
Defaults to None.
|
||
vllm_data_parallel_size (int): Data parallelism size for vLLM rollout. Default is 1.
|
||
"""
|
||
# vllm
|
||
vllm_gpu_memory_utilization: float = 0.9
|
||
vllm_tensor_parallel_size: int = 1
|
||
vllm_pipeline_parallel_size: int = 1
|
||
vllm_enable_expert_parallel: bool = False
|
||
vllm_max_num_seqs: int = 256
|
||
vllm_max_model_len: Optional[int] = None
|
||
vllm_disable_custom_all_reduce: bool = True
|
||
vllm_enforce_eager: bool = False
|
||
vllm_limit_mm_per_prompt: Optional[Union[dict, str]] = None # '{"image": 5, "video": 2}'
|
||
vllm_max_lora_rank: int = 16
|
||
vllm_enable_prefix_caching: Optional[bool] = None
|
||
vllm_use_async_engine: Optional[bool] = None
|
||
vllm_quantization: Optional[str] = None
|
||
vllm_reasoning_parser: Optional[str] = None
|
||
vllm_disable_cascade_attn: bool = False
|
||
vllm_mm_processor_cache_gb: Optional[float] = None
|
||
vllm_speculative_config: Optional[Union[dict, str]] = None
|
||
vllm_engine_kwargs: Optional[Union[dict, str]] = None
|
||
# rollout
|
||
vllm_data_parallel_size: int = 1
|
||
|
||
def __post_init__(self):
|
||
if self.vllm_limit_mm_per_prompt is not None:
|
||
self.vllm_limit_mm_per_prompt = json_parse_to_dict(self.vllm_limit_mm_per_prompt)
|
||
if self.vllm_speculative_config is not None:
|
||
self.vllm_speculative_config = json_parse_to_dict(self.vllm_speculative_config)
|
||
self.vllm_engine_kwargs = json_parse_to_dict(self.vllm_engine_kwargs)
|
||
|
||
def get_vllm_engine_kwargs(self):
|
||
# Some parameters are provided by BaseArguments
|
||
adapters = self.adapters
|
||
if hasattr(self, 'adapter_mapping'):
|
||
adapters = adapters + list(self.adapter_mapping.values())
|
||
kwargs = {
|
||
'gpu_memory_utilization': self.vllm_gpu_memory_utilization,
|
||
'tensor_parallel_size': self.vllm_tensor_parallel_size,
|
||
'pipeline_parallel_size': self.vllm_pipeline_parallel_size,
|
||
'enable_expert_parallel': self.vllm_enable_expert_parallel,
|
||
'max_num_seqs': self.vllm_max_num_seqs,
|
||
'max_model_len': self.vllm_max_model_len,
|
||
'disable_custom_all_reduce': self.vllm_disable_custom_all_reduce,
|
||
'enforce_eager': self.vllm_enforce_eager,
|
||
'limit_mm_per_prompt': self.vllm_limit_mm_per_prompt,
|
||
'max_lora_rank': self.vllm_max_lora_rank,
|
||
'enable_lora': len(adapters) > 0,
|
||
'max_loras': max(len(adapters), 1),
|
||
'enable_prefix_caching': self.vllm_enable_prefix_caching,
|
||
'use_async_engine': self.vllm_use_async_engine,
|
||
'quantization': self.vllm_quantization,
|
||
'reasoning_parser': self.vllm_reasoning_parser,
|
||
'disable_cascade_attn': self.vllm_disable_cascade_attn,
|
||
'mm_processor_cache_gb': self.vllm_mm_processor_cache_gb,
|
||
'speculative_config': self.vllm_speculative_config,
|
||
'num_labels': self.num_labels,
|
||
'engine_kwargs': self.vllm_engine_kwargs,
|
||
}
|
||
if self.task_type in ('embedding', 'seq_cls') or 'reranker' in self.task_type:
|
||
kwargs['task_type'] = self.task_type
|
||
|
||
return kwargs
|
||
|
||
|
||
@dataclass
|
||
class RolloutTrainerArgumentsMixin(VllmArguments):
|
||
"""A dataclass for configuring parameters required for rollout-based training (GRPO, GKD, etc.).
|
||
|
||
This mixin provides arguments for controlling the generation process during rollout, especially when using vLLM as
|
||
the inference backend for generation.
|
||
|
||
Args:
|
||
top_k (int): The number of highest probability vocabulary tokens to keep for top-k-filtering. -1 means
|
||
no filtering. Defaults to -1.
|
||
top_p (float): If set to a float < 1, only the smallest set of most probable tokens with probabilities that
|
||
add up to top_p or higher are kept for generation. Defaults to 1.0.
|
||
repetition_penalty (float): The parameter for repetition penalty. 1.0 means no penalty. Defaults to 1.0.
|
||
stop_words (List[str]): A list of strings that will stop the generation when they are generated. Defaults to an
|
||
empty list.
|
||
use_vllm (bool): Whether to use vLLM as the inference backend for generation. Defaults to False.
|
||
vllm_mode (Literal['server', 'colocate']): The vLLM integration mode. 'server' mode uses a vLLM server launched
|
||
by swift rollout for sampling. 'colocate' mode deploys vLLM within the same process. For full-parameter
|
||
training in 'server' mode, the environment variable `SWIFT_UPDATE_WEIGHTS_BUCKET_SIZE` can be set to
|
||
control the bucket size for weight synchronization (in MB, defaults to 512). Defaults to 'colocate'.
|
||
vllm_enable_prefix_caching (bool): A pass-through parameter for vLLM. Enables prefix caching. Used in
|
||
'colocate' mode. Defaults to True.
|
||
vllm_enable_lora (bool): Enables the vLLM Engine to load LoRA adapters. This is used to accelerate weight
|
||
synchronization during LoRA training. See documentation for details. Used in 'colocate' mode. Defaults
|
||
to False.
|
||
lora_rank (int): The rank for the LoRA adapter loaded by the vLLM engine. When using `vllm_enable_lora`, this
|
||
should be greater than or equal to (ideally equal to) the rank of the LoRA model being trained. Used in
|
||
'colocate' mode. Defaults to 8.
|
||
vllm_server_base_url (Optional[List[str]]): The base URL of the vLLM server (e.g., "http://localhost:8000").
|
||
If set, `vllm_server_host` and `vllm_server_port` are ignored. Used in 'server' mode. Defaults to None.
|
||
vllm_server_host (Optional[List[str]]): The host address(es) of the vLLM
|
||
server. Used in 'server' mode. Defaults to None.
|
||
vllm_server_port (List[int]): The port(s) of the vLLM server. Used in 'server' mode. Defaults to `[8000]`.
|
||
vllm_server_timeout (float): Timeout in seconds for connecting to the vLLM server. Used in 'server' mode.
|
||
Defaults to 240.0.
|
||
vllm_client: Internal client instance for vLLM server communication. Not intended to be set by the user.
|
||
async_generate (bool): Whether to perform asynchronous rollout to improve training speed. Note: When enabled,
|
||
sampling uses the model from the previous update step and is not compatible with multi-turn scenarios.
|
||
Defaults to False.
|
||
sleep_level (int): Specifies the level of GPU memory release for vLLM during training steps. Options: 0
|
||
(no release), 1, 2. A higher level releases more memory but may incur overhead. Defaults to 0.
|
||
move_model_batches (Optional[int]): The number of batches after which the model is moved back to the GPU if it
|
||
was offloaded. Used for memory management during training. Defaults to None.
|
||
offload_optimizer (bool): Whether to offload optimizer states to CPU/RAM when performing inference with vLLM to
|
||
save GPU memory. Defaults to False.
|
||
offload_model (bool): Whether to offload the model weights to CPU/RAM when performing inference with vLLM.
|
||
Defaults to False.
|
||
wandb_log_unique_prompts (Optional[bool]): Whether to log unique prompts to Weights & Biases for analysis
|
||
during training. Defaults to None.
|
||
structured_outputs_regex (Optional[str]): A regular expression pattern for structured outputs (guided
|
||
decoding). When set, the model's generation is constrained to match the specified regex pattern. This is
|
||
useful for tasks requiring structured outputs like reasoning chains. Defaults to None (disabled).
|
||
Only effective when using vLLM backend (`use_vllm=True`).
|
||
generation_batch_size (Optional[int]): The total number of samples generated in one vLLM inference call.
|
||
It should be a multiple of `num_processes * per_device_train_batch_size`. Defaults to
|
||
`per_device_train_batch_size * gradient_accumulation_steps * num_processes`.
|
||
steps_per_generation (Optional[int]): The number of training steps (micro-batches) that reuse the same
|
||
set of generated samples. Only one of `steps_per_generation` and `generation_batch_size` should be set.
|
||
Defaults to `gradient_accumulation_steps`.
|
||
"""
|
||
# generation args
|
||
top_k: int = -1
|
||
top_p: float = 1.0
|
||
repetition_penalty: float = 1.
|
||
stop_words: List[str] = field(default_factory=list)
|
||
|
||
# vllm
|
||
use_vllm: bool = False
|
||
vllm_mode: Optional[Literal['server', 'colocate']] = None
|
||
# internal vllm (colocate)
|
||
vllm_max_num_seqs: Optional[int] = None
|
||
vllm_enable_prefix_caching: bool = True # overwrite
|
||
vllm_enable_lora: bool = False
|
||
lora_rank: int = 8 # for vllm lora adapter
|
||
# external vllm (server)
|
||
vllm_server_base_url: Optional[List[str]] = None
|
||
vllm_server_host: Optional[List[str]] = None
|
||
vllm_server_port: List[int] = field(default_factory=lambda: [8000])
|
||
vllm_server_timeout: float = 240.0
|
||
vllm_client = None # Not required to set, used for client instantiation
|
||
vllm_server_group_port: Optional[List[int]] = None
|
||
enable_flattened_weight_sync: bool = True
|
||
async_generate: bool = False
|
||
# # structured outputs (guided decoding), only effective for vllm backend
|
||
structured_outputs_regex: Optional[str] = None
|
||
|
||
sleep_level: int = 0
|
||
move_model_batches: Optional[int] = None
|
||
offload_optimizer: bool = False
|
||
offload_model: bool = False
|
||
|
||
wandb_log_unique_prompts: Optional[bool] = None
|
||
|
||
generation_batch_size: Optional[int] = None
|
||
steps_per_generation: Optional[int] = None
|
||
|
||
teacher_tag_key: str = 'dataset'
|
||
|
||
def _init_generation_batch_params(self):
|
||
num_generations = getattr(self, 'num_generations', 1)
|
||
num_processes = self.world_size
|
||
global_batch_size = self.per_device_train_batch_size * num_processes
|
||
|
||
if self.generation_batch_size is None and self.steps_per_generation is None:
|
||
self.steps_per_generation = self.gradient_accumulation_steps
|
||
self.generation_batch_size = global_batch_size * self.steps_per_generation
|
||
elif self.generation_batch_size is not None and self.steps_per_generation is None:
|
||
if self.generation_batch_size % global_batch_size != 0:
|
||
raise ValueError(f'generation_batch_size ({self.generation_batch_size}) must be divisible by '
|
||
f'the global batch size ({global_batch_size}).')
|
||
self.steps_per_generation = self.generation_batch_size // global_batch_size
|
||
elif self.generation_batch_size is None and self.steps_per_generation is not None:
|
||
self.generation_batch_size = global_batch_size * self.steps_per_generation
|
||
else:
|
||
expected = global_batch_size * self.steps_per_generation
|
||
if self.generation_batch_size != expected:
|
||
raise ValueError(f'generation_batch_size ({self.generation_batch_size}) must equal '
|
||
f'per_device_train_batch_size * world_size * steps_per_generation = {expected}.')
|
||
|
||
if self.steps_per_generation <= 0:
|
||
raise ValueError(f'steps_per_generation must be > 0, got {self.steps_per_generation}.')
|
||
|
||
if num_generations > 1:
|
||
if self.generation_batch_size % num_generations != 0:
|
||
possible_values = [
|
||
n for n in range(2, self.generation_batch_size + 1) if self.generation_batch_size % n == 0
|
||
]
|
||
raise ValueError(f'generation_batch_size ({self.generation_batch_size}) must be evenly divisible by '
|
||
f'num_generations ({num_generations}). Valid values: {possible_values}.')
|
||
|
||
if hasattr(self, 'eval_strategy') and self.eval_strategy != 'no':
|
||
num_generations_eval = getattr(self, 'num_generations_eval', None) or num_generations
|
||
global_eval_batch_size = self.per_device_eval_batch_size * num_processes
|
||
if global_eval_batch_size % num_generations_eval != 0:
|
||
possible_values = [
|
||
n for n in range(1, global_eval_batch_size + 1) if global_eval_batch_size % n == 0
|
||
]
|
||
raise ValueError(
|
||
f'The global eval batch size ({num_processes} x {self.per_device_eval_batch_size}) must be '
|
||
f'evenly divisible by num_generations_eval ({num_generations_eval}). '
|
||
f'Valid values: {possible_values}.')
|
||
|
||
|
||
@dataclass
|
||
class GRPOArgumentsMixin(RolloutTrainerArgumentsMixin):
|
||
"""A dataclass for configuring parameters for algorithms like DAPO, Dr.GRPO, GSPO, RLOO, and REINFORCE++.
|
||
|
||
Args:
|
||
epsilon (float): The clipping coefficient. Defaults to 0.2.
|
||
epsilon_high (Optional[float]): The upper clipping coefficient. If set, it forms a clipping range of
|
||
`[epsilon, epsilon_high]` with epsilon. Defaults to None.
|
||
delta (Optional[float]): The upper clipping value for two-sided GRPO from the INTELLECT-2 tech
|
||
report. If set, it is recommended to be greater than `1 + epsilon`. Defaults to None.
|
||
cosine_min_len_value_wrong (float): The reward for wrong answers with zero completion length
|
||
(r^w_0 in the paper). Defaults to -0.5.
|
||
cosine_max_len_value_wrong (float): The reward for wrong answers with maximum completion length
|
||
(r^w_L in the paper). Defaults to 0.0.
|
||
cosine_min_len_value_correct (float): The reward for correct answers with zero completion length
|
||
(r^c_0 in the paper). Defaults to 1.0.
|
||
cosine_max_len_value_correct (float): The reward for correct answers with maximum completion length
|
||
(r^c_L in the paper). Defaults to 0.5.
|
||
cosine_max_len (Optional[int]): The maximum length for generated text (Lmax in the paper). Defaults
|
||
to `max_completion_length`.
|
||
repetition_n_grams (int): The n-gram size for repetition detection. Defaults to 3.
|
||
repetition_max_penalty (float): The maximum penalty value, used to control the strength of the
|
||
penalty. Defaults to -1.0.
|
||
reward_model (Optional[List[str]]): The reward model(s) to use. Defaults to None.
|
||
reward_model_plugin (Optional[List[str]]): The plugin logic for the reward model. Defaults to 'orm'
|
||
logic. See custom reward models for details. Defaults to None.
|
||
chord_sft_dataset (List[str]): The SFT dataset(s) used to provide expert data for the CHORD algorithm. Defaults
|
||
to `[]`.
|
||
chord_sft_per_device_train_batch_size (Optional[int]): The SFT mini-batch size per device for the CHORD
|
||
algorithm. Defaults to None.
|
||
chord_enable_phi_function (bool): Whether to enable the token-wise weighting function phi (φ) in the CHORD
|
||
algorithm. Defaults to False.
|
||
chord_mu_warmup_steps (Optional[int]): The number of training steps for the mu (μ) value to warm up to its peak
|
||
value. Defaults to None.
|
||
chord_mu_decay_steps (Optional[int]): The number of training steps for the mu (μ) value to decay from its peak
|
||
to its valley value. Defaults to None.
|
||
chord_mu_peak (Optional[float]): The peak value for mu (μ) during its schedule. Defaults to None.
|
||
chord_mu_valley (Optional[float]): The final (valley) value for mu (μ) after decay. Defaults to None.
|
||
sync_ref_model (bool): Whether to periodically synchronize the `ref_model`. Defaults to False.
|
||
ref_model_sync_steps (int): The synchronization frequency. Defaults to 512.
|
||
ref_model_mixup_alpha (float): Controls the mixup between the current model and the previous
|
||
`ref_model` during updates. Defaults to 0.6.
|
||
multi_turn_scheduler (Optional[str]): Parameter for multi-turn GRPO. Pass the corresponding plugin
|
||
name. The implementation should be added in `plugin/multi_turn.py`. Defaults to None.
|
||
max_turns (Optional[int]): The maximum number of turns for multi-turn GRPO. If None, there is no
|
||
limit. Defaults to None.
|
||
completion_length_limit_scope (Literal['total', 'per_round']): The scope of the
|
||
`max_completion_length` limit in multi-turn dialogue. 'total' limits the total output length across all
|
||
turns, while 'per_round' limits the output length for each turn. Defaults to 'per_round'.
|
||
vllm_server_pass_dataset (bool): Pass extra dataset information to the vLLM server, used for
|
||
multi-turn training. Defaults to False.
|
||
use_gym_env (Optional[bool]): If set, the trainer treats `rollout_infos['total_reward']` produced
|
||
by a gym-style multi-turn scheduler as the reward (no reward function needed). Works in both
|
||
`server` and `colocate` modes, and on the Megatron trainer. When `None` (default), it auto-defaults
|
||
to `True` if `gym_env` is set; otherwise it is auto-detected from the connected vLLM server in
|
||
`server` mode and `False` otherwise. An explicit value here is authoritative — it is never
|
||
overridden by the value reported by the rollout server.
|
||
gym_env (Optional[str]): Default gym environment name used by the `gym_scheduler`. Equivalent to
|
||
`--gym_env` on `swift rollout` but for the trainer-side colocate path; per-row `env_config.name`
|
||
still wins over this default. Defaults to None.
|
||
dynamic_sample (bool): If True, filters out data with a reward standard deviation of 0 within a group
|
||
and samples new data. Defaults to False.
|
||
max_resample_times (int): When `dynamic_sample` is enabled, this limits the number of resampling
|
||
attempts. Defaults to 3.
|
||
overlong_filter (bool): If True, skips samples that are truncated due to being too long, so they are
|
||
not included in the loss calculation. Defaults to False.
|
||
soft_max_length (Optional[int]): The maximum generation length of the model (L_max in the paper).
|
||
Defaults to `max_completion_length`.
|
||
soft_cache_length (Optional[int]): Controls the length penalty interval (L_cache in the paper).
|
||
The interval is `[soft_max_length - soft_cache_length, soft_max_length]`. Defaults to None.
|
||
scale_rewards (Optional[Literal['group', 'batch', 'none', 'gdpo']]): Specifies the reward scaling strategy.
|
||
Options are 'group' (scale by intra-group standard deviation), 'batch' (scale by the entire batch's
|
||
standard deviation), 'none' (no scaling), or 'gdpo' (normalize each reward function separately before
|
||
weighted aggregation; reward_weights serve as weights for each reward's advantage term). The default
|
||
value is tied to `advantage_estimator`: 'group' for 'grpo', 'none' for 'rloo', and 'batch' for
|
||
'reinforce_plus_plus'. In ms-swift < 3.10, this was a boolean where `True` corresponded to 'group'
|
||
and `False` to 'none'. Note: GDPO mode does not support kl_in_reward=True.
|
||
log_entropy (bool): Log the dynamics of entropy values during training. See documentation for
|
||
details. Defaults to False.
|
||
top_entropy_quantile (float): Only tokens with entropy in the top specified quantile participate in
|
||
the loss calculation. A value of 1.0 means no tokens are filtered. See documentation for details.
|
||
Defaults to 1.0.
|
||
importance_sampling_level (Literal['token', 'sequence', 'sequence_token']): Controls the importance
|
||
sampling ratio calculation. 'token' mode retains the original log probability ratio for each token.
|
||
'sequence' mode averages the log probability ratios of all valid tokens in the sequence. The GSPO paper
|
||
uses 'sequence' level for stable training. Defaults to 'token'.
|
||
tau_pos (float): The temperature parameter for positive dominance in the SAPO algorithm, controlling the
|
||
sharpness of the soft gating function. Larger values result in sharper gating (approaching hard
|
||
clipping), while smaller values result in smoother gating. The default value is 1.0.
|
||
tau_neg (float): The temperature parameter for negative dominance in the SAPO algorithm, controlling the
|
||
sharpness of the soft gating function. Typically, `tau_neg` is set > `tau_pos` to impose stronger
|
||
constraints on negative dominance. The default value is 1.05.
|
||
real_tau (float): The temperature parameter. REAL induces monotonic and bounded gradient weighting with
|
||
magnitude upper-bounded by 1/tau. The default value is 0.5.
|
||
fipo_decay_rate (float): Half-life used to derive `fipo_gamma`. Defaults to 32.0.
|
||
fipo_clip_range (Optional[float]): Clip range for the FIPO influence weight. `0.2` clips to
|
||
`[0.8, 1.2]`; `None` or `0` disables clipping. Defaults to 0.2.
|
||
fipo_clip_high_only (bool): If `True`, clips the FIPO influence weight to `[1, 1 + fipo_clip_range]`.
|
||
Defaults to True.
|
||
fipo_safety_threshold (Optional[float]): Safety threshold for negative advantages. Tokens with
|
||
`advantage < 0` and importance ratio above this value have their FIPO influence weight capped to
|
||
`[0.8, 1.0]` to avoid over-penalization. Defaults to 4.0.
|
||
advantage_estimator (Literal['grpo', 'rloo', 'reinforce_plus_plus']): The advantage estimation
|
||
function to use. 'grpo' calculates the relative advantage within a group. Options are 'grpo', 'rloo',
|
||
'reinforce_plus_plus'. Defaults to 'grpo'.
|
||
kl_in_reward (Optional[bool]): Controls how the KL divergence regularization term is handled. If
|
||
`False`, it's an independent term in the loss function. If `True`, KL is directly incorporated into the
|
||
reward (subtracted from it). The default is tied to `advantage_estimator`: `False` for 'grpo', `True` for
|
||
'rloo' and 'reinforce_plus_plus'.
|
||
num_generations_eval (Optional[int]): Number of generations to sample during evaluation. This allows
|
||
using fewer generations during evaluation to save computation. If `None`, uses the value of
|
||
`num_generations`. Defaults to None.
|
||
dataset_shuffle (Optional[bool]): Whether to shuffle the dataset. Defaults to True.
|
||
rollout_importance_sampling_mode (Optional[Literal['token_truncate', 'token_mask', 'sequence_truncate',
|
||
'sequence_mask']]): The training-pull inconsistency correction mode. Options are `token_truncate`,
|
||
`token_mask`, `sequence_truncate`, and `sequence_mask`. Defaults to None, disabling correction.
|
||
See the documentation for details.
|
||
rollout_importance_sampling_threshold (float): The threshold for importance sampling weights, used to truncate
|
||
or mask extreme weights. Defaults to 2.0.
|
||
log_rollout_offpolicy_metrics (bool): Whether to log rollout off-policy diagnostic metrics (KL, PPL, chi2, etc.)
|
||
when `rollout_importance_sampling_mode` is not set. When `rollout_importance_sampling_mode` is set,
|
||
metrics are always logged regardless of this setting. Defaults to False.
|
||
"""
|
||
epsilon: float = 0.2
|
||
epsilon_high: Optional[float] = None
|
||
delta: Optional[float] = None
|
||
|
||
# reward function args, see details in swift/rewards/orm.py
|
||
# cosine reward, https://arxiv.org/abs/2502.03373
|
||
cosine_min_len_value_wrong: float = -0.5 # r^w_0 in paper, Reward for wrong answers with zero completion length.
|
||
cosine_max_len_value_wrong: float = 0.0 # r^w_L in paper, Reward for wrong answers with max completion length.
|
||
cosine_min_len_value_correct: float = 1.0 # r^c_0 in paper, Reward for correct answers with zero completion length.
|
||
cosine_max_len_value_correct: float = 0.5 # r^c_L in paper, Reward for correct answers with max completion length.
|
||
cosine_max_len: Optional[int] = None # Lmax in paper, default equal to max_completion_length
|
||
# repetition penalty, https://arxiv.org/abs/2502.03373
|
||
repetition_n_grams: int = 3
|
||
repetition_max_penalty: float = -1.0
|
||
|
||
reward_model: Optional[List[str]] = None
|
||
reward_model_plugin: Optional[List[str]] = None
|
||
|
||
# chord
|
||
chord_sft_dataset: List[str] = field(default_factory=list)
|
||
chord_sft_per_device_train_batch_size: Optional[int] = None
|
||
|
||
chord_enable_phi_function: bool = False
|
||
chord_mu_warmup_steps: Optional[int] = None
|
||
chord_mu_decay_steps: Optional[int] = None
|
||
chord_mu_peak: Optional[float] = None
|
||
chord_mu_valley: Optional[float] = None
|
||
|
||
# sync ref model
|
||
sync_ref_model: bool = False
|
||
ref_model_sync_steps: int = 512
|
||
ref_model_mixup_alpha: float = 0.6
|
||
|
||
# multi turn
|
||
multi_turn_scheduler: Optional[str] = None
|
||
max_turns: Optional[int] = None
|
||
completion_length_limit_scope: Literal['total', 'per_round'] = 'per_round'
|
||
vllm_server_pass_dataset: bool = False
|
||
use_gym_env: Optional[bool] = None
|
||
gym_env: Optional[str] = None
|
||
|
||
# DAPO, https://arxiv.org/abs/2503.14476
|
||
dynamic_sample: bool = False
|
||
max_resample_times: int = 3
|
||
overlong_filter: bool = False
|
||
soft_max_length: Optional[int] = None
|
||
soft_cache_length: Optional[int] = None
|
||
|
||
# Dr. GRPO, https://arxiv.org/abs/2503.20783
|
||
# GDPO: normalize each reward function separately, https://arxiv.org/abs/2601.05242
|
||
scale_rewards: Optional[Literal['group', 'batch', 'none', 'gdpo']] = None
|
||
|
||
# entropy
|
||
log_entropy: bool = False
|
||
# Beyond the 80/20 Rule, https://arxiv.org/abs/2506.01939
|
||
top_entropy_quantile: float = 1.0
|
||
|
||
# GSPO https://arxiv.org/abs/2507.18071
|
||
importance_sampling_level: Literal['token', 'sequence', 'sequence_token'] = 'token'
|
||
|
||
# SAPO https://arxiv.org/abs/2511.20347
|
||
# Temperature parameters for soft adaptive gate
|
||
tau_pos: float = 1.0
|
||
tau_neg: float = 1.05
|
||
|
||
# RLOO, REINFORCE++
|
||
advantage_estimator: Literal['grpo', 'rloo', 'reinforce_plus_plus'] = 'grpo'
|
||
# If false, add KL into loss, otherwise add into reward
|
||
kl_in_reward: Optional[bool] = None # rloo/reinforce_plus_plus: true, grpo: false (default)
|
||
|
||
# OPD-RL (On-Policy Distillation as RL)
|
||
# enabled when a teacher (teacher_model / teacher_model_server) is set on a GRPO run.
|
||
teacher_kl_coef: float = 1.0
|
||
|
||
# REAL https://arxiv.org/abs/2602.05630
|
||
real_tau: float = 0.5
|
||
|
||
# FIPO https://arxiv.org/abs/2603.19835
|
||
fipo_decay_rate: float = 32.0
|
||
fipo_clip_range: Optional[float] = 0.2
|
||
fipo_clip_high_only: bool = True
|
||
fipo_safety_threshold: Optional[float] = 4.0
|
||
|
||
num_generations_eval: Optional[int] = None
|
||
|
||
# dataset
|
||
dataset_shuffle: Optional[bool] = True
|
||
|
||
# Rollout Importance Sampling Correction (off-policy correction)
|
||
# Set to None to disable, or choose from: 'token_truncate', 'token_mask', 'sequence_truncate', 'sequence_mask'
|
||
rollout_importance_sampling_mode: Optional[Literal['token_truncate', 'token_mask', 'sequence_truncate',
|
||
'sequence_mask']] = None
|
||
rollout_importance_sampling_threshold: float = 2.0 # Threshold for truncation/masking (C in paper)
|
||
log_rollout_offpolicy_metrics: bool = False # Log off-policy metrics even when IS correction is disabled
|
||
# Off-Policy Sequence Masking: mask out sequences that deviate too much from rollout policy
|
||
# If set, compute mean(rollout_per_token_logps - per_token_logps) per sequence,
|
||
# and mask sequences where this delta > threshold AND advantage < 0
|
||
# Falls back to old_per_token_logps if rollout_per_token_logps is not available
|
||
off_policy_sequence_mask_delta: Optional[float] = None
|
||
|
||
def __post_init__(self):
|
||
super().__post_init__()
|
||
# gym_env implies use_gym_env unless the user said otherwise; mirrors deploy_args behavior so the
|
||
# default propagates to GRPOConfig too (not just RLHFArguments).
|
||
if self.use_gym_env is None and self.gym_env is not None:
|
||
self.use_gym_env = True
|