wehub-resource-sync
205 lines
14 KiB
Python
205 lines
14 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from typing import Literal, Optional
|
|
|
|
from swift.template import TEMPLATE_MAPPING, get_template_meta
|
|
from swift.utils import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class TemplateArguments:
|
|
"""TemplateArguments class holds various arguments for template configuration.
|
|
|
|
This dataclass manages settings related to how data is formatted and processed using templates, including
|
|
tokenization, truncation, loss calculation, and special handling for multimodal and agent-based models.
|
|
|
|
Args:
|
|
template (Optional[str]): The dialogue template type. Defaults to None, which automatically selects the
|
|
template corresponding to the model type. Refer to the list of supported models for mappings.
|
|
system (Optional[str]): Custom system prompt. Can be a string or a path to a .txt file. Defaults to None,
|
|
which uses the default system from the registered template.
|
|
Note: The priority for the system prompt is as follows:
|
|
1. System prompt from the dataset.
|
|
2. The `--system` command-line argument.
|
|
3. The `default_system` set when the template was registered.
|
|
max_length (Optional[int]): The maximum number of tokens for a single sample after tokenization. Samples
|
|
exceeding this length are handled according to `truncation_strategy` to prevent OOM errors. Defaults to
|
|
None, which uses the model's maximum supported length (`max_model_len`). In PPO, GRPO, and inference
|
|
scenarios, this argument specifies the `max_prompt_length`.
|
|
truncation_strategy (Literal['delete', 'left', 'right', 'split']): Strategy for handling samples exceeding
|
|
`max_length`. Options are 'delete', 'left' (truncate from the left), 'right' (truncate from the right),
|
|
and 'split' (split into multiple samples). Defaults to 'delete'.
|
|
Note: The 'split' strategy is only supported during pre-training (e.g., `swift/megatron pt`),
|
|
and is incompatible with `cached_dataset`. It splits long samples to avoid wasting tokens.
|
|
Note: For multimodal models, setting this to 'left' or 'right' preserves all image tokens, which may lead
|
|
to OOM errors.
|
|
max_pixels (Optional[int]): The maximum number of pixels (H*W) for an input image in a multimodal model.
|
|
Images exceeding this limit will be scaled down to prevent OOM errors. Defaults to None, meaning no limit.
|
|
Note: This parameter applies to all multimodal models. The model-specific `MAX_PIXELS` parameter for
|
|
Qwen2.5-VL is separate and only applies to that model.
|
|
agent_template (Optional[str]): The Agent template to use. This determines how the 'tools' list is converted
|
|
into a 'system' prompt, how tool calls are extracted from the model's response during inference, and the
|
|
format for tool call messages. Options include "react_en", "hermes", "glm4", "qwen_en", "toolbench", etc.
|
|
Defaults to None, which auto-selects based on the model type. Refer to the Agent documentation for more
|
|
details.
|
|
norm_bbox (Optional[Literal['norm1000', 'none']]): Controls how bounding box coordinates (from the "bbox"
|
|
field in the dataset) are scaled. 'norm1000' scales coordinates to a 1000x1000 grid, while 'none' performs
|
|
no scaling. Defaults to None, which auto-selects based on the model. This handles cases where images are
|
|
resized during training (e.g., due to `max_pixels`).
|
|
use_chat_template (bool): Whether to use the chat template or the generation template. The generation template
|
|
is typically used for pre-training. Defaults to True.
|
|
Note: Defaults to False for `swift pt`, which uses the generation template. This parameter is compatible
|
|
with multimodal models.
|
|
padding_side (Literal['left', 'right']): The side to pad on when `batch_size >= 2` during training.
|
|
Options are 'left' or 'right'. Defaults to 'right'. For inference with `batch_size >= 2`, padding is always
|
|
on the left.
|
|
Note: Defaults to 'left' for PPO and GKD.
|
|
padding_free (bool): If True, flattens the data within a batch to avoid padding, reducing memory usage and
|
|
speeding up training. Sequences within the batch remain causally isolated. Defaults to False. Supported for
|
|
CPT/SFT/DPO/GRPO/KTO/GKD.
|
|
Note: This requires `--attn_impl flash_attn` and `transformers>=4.44`. Compared to packing, padding_free
|
|
has no preprocessing overhead, but packing offers faster training speeds and more stable memory usage.
|
|
loss_scale (str): Loss weight configuration for training tokens. Default is `'default'`.
|
|
loss_scale includes 3 basic strategies: 'default', 'last_round', 'all', and other strategies:
|
|
'ignore_empty_think' and agent-specific ones: 'react', 'hermes', 'qwen', 'agentflan', 'alpha_umi', etc.
|
|
For available options, refer to
|
|
[loss_scale module](https://github.com/modelscope/ms-swift/blob/main/swift/loss_scale/mapping.py).
|
|
ms-swift supports mixing basic strategies with other strategies,
|
|
for example: `'default+ignore_empty_think'`, `'last_round+ignore_empty_think'`.
|
|
If no basic strategy is specified, it defaults to 'default',
|
|
for example: 'hermes' is equivalent to 'default+hermes'.
|
|
Multiple non-base strategies can be chained together
|
|
(each strategy processes the output segments of the previous one, with weights
|
|
multiplied accordingly). For example: `'last_round+hermes+ignore_empty_think'`, where
|
|
`'last_round'` is the base strategy, and `'hermes+ignore_empty_think'` represents a
|
|
chain of multiple non-base strategies that share the same base strategy.
|
|
- 'default': All responses (including history) are calculated with weight 1 for cross-entropy loss
|
|
(**system/user/multimodal tokens in messages and `tool_response` parts in Agent training are
|
|
not included in loss calculation**). (**Default value for SFT**)
|
|
- 'last_round': Only calculate loss for the last round response. The last round
|
|
means all content after the last "user". (**Default value for RLHF**)
|
|
- 'all': Calculate loss for all tokens. (**Default value for `swift pt`**)
|
|
- 'ignore_empty_think': Ignore loss computation for empty `'<think>\n\n</think>\n\n'`
|
|
(as long as it matches the regex `'<think>\\s*</think>\\s*'`).
|
|
- 'react', 'hermes', 'qwen': Adjust the loss weight of the `tool_call` part to 2.
|
|
sequence_parallel_size (int): The size of sequence parallelism. Defaults to 1. Currently supported for CPT,
|
|
SFT, DPO, and GRPO.
|
|
template_backend (Literal['swift', 'jinja']): The backend to use for templating. Options are 'swift' or
|
|
'jinja'. Defaults to 'swift'. If 'jinja' is used, it will leverage `transformers.apply_chat_template`.
|
|
Note: The 'jinja' backend is only supported for inference, not for training, as it cannot determine the
|
|
token range for loss calculation.
|
|
response_prefix (Optional[str]): A prefix string for the response, e.g., '<think>\\n' for Qwen-32B. This
|
|
parameter only affects inference. Defaults to None, which is auto-set based on the model.
|
|
enable_thinking (Optional[bool]): This parameter takes effect during inference,
|
|
indicating whether to enable thinking mode. Default is None, the default value is determined by the
|
|
template (model) type (True for thinking/hybrid thinking templates, False for non-thinking templates).
|
|
If enable_thinking is False, a non-thinking prefix is added, for example the Qwen3-8B hybrid thinking
|
|
model adds the prefix `'<think>\n\n</think>\n\n'`, while Qwen3-8B-Thinking does not add a prefix.
|
|
If enable_thinking is True, a thinking prefix is added, for example `'<think>\n'`.
|
|
Note: The priority of this parameter is lower than the response_prefix parameter.
|
|
preserve_thinking (Optional[bool]): Whether to preserve historical thinking content during inference and
|
|
training. When set to `True`, thinking content from all rounds is retained. When set to `False`,
|
|
only the thinking content from the last round is retained (i.e., the content following the last
|
|
user message). Defaults to `None`.
|
|
Default behavior: For thinking models (thinking/hybrid-thinking) or when `enable_thinking` is
|
|
explicitly enabled, this is set to `False` by default during inference and training, retaining
|
|
only the last round of thinking content. If the `loss_scale` base strategy during training is
|
|
not `'last_round'` (e.g., `'default'`), it defaults to `True`, and historical thinking content will
|
|
not be removed.
|
|
add_non_thinking_prefix (bool): This parameter only takes effect during training, indicating whether to
|
|
add a non-thinking prefix to data samples whose assistant part does not start with the thinking
|
|
marker `'<think>'` (typically hybrid thinking models contain a non-thinking prefix).
|
|
This feature allows swift's built-in datasets to train hybrid thinking models. Default value is True.
|
|
For example: the non-thinking prefix for the Qwen3-8B hybrid thinking model is
|
|
`'<think>\n\n</think>\n\n'`, while the non-thinking prefix for Qwen3-8B-Thinking/Instruct is `''`.
|
|
Note: During training, if the basic strategy of loss_scale is last_round, this modification is only
|
|
applied to the last round; otherwise, for example 'default' or 'all', this modification is applied to
|
|
every round of data. If set to False, no non-thinking prefix is added to data samples.
|
|
|
|
|
|
"""
|
|
template: Optional[str] = field(
|
|
default=None, metadata={'help': f'template choices: {list(TEMPLATE_MAPPING.keys())}'})
|
|
system: Optional[str] = None # Override the default_system in the template.
|
|
max_length: Optional[int] = None
|
|
|
|
truncation_strategy: Literal['delete', 'left', 'right', 'split', None] = None
|
|
max_pixels: Optional[int] = None
|
|
agent_template: Optional[str] = None
|
|
norm_bbox: Literal['norm1000', 'none', None] = None
|
|
use_chat_template: Optional[bool] = None
|
|
padding_side: Literal['left', 'right'] = 'right'
|
|
# train
|
|
padding_free: bool = False
|
|
loss_scale: str = 'default'
|
|
sequence_parallel_size: int = 1
|
|
is_binary_loss_scale: Optional[bool] = None
|
|
# infer/deploy
|
|
template_backend: Literal['swift', 'jinja'] = 'swift'
|
|
# thinking
|
|
response_prefix: Optional[str] = None
|
|
enable_thinking: Optional[bool] = None
|
|
preserve_thinking: Optional[bool] = None
|
|
add_non_thinking_prefix: bool = True
|
|
disable_ignore_empty_think: bool = False
|
|
|
|
def __post_init__(self):
|
|
if getattr(self, 'model_meta', None) is not None:
|
|
self.template_meta = get_template_meta(self.model_info, self.model_meta, template_type=self.template)
|
|
self.template = self.template_meta.template_type
|
|
if self.use_chat_template is None:
|
|
self.use_chat_template = True
|
|
if self.system is not None:
|
|
if self.system.endswith('.txt'):
|
|
assert os.path.isfile(self.system), f'self.system: {self.system}'
|
|
with open(self.system, 'r', encoding='utf-8') as f:
|
|
self.system = f.read()
|
|
else:
|
|
self.system = self.system.replace('\\n', '\n')
|
|
if self.response_prefix is not None:
|
|
self.response_prefix = self.response_prefix.replace('\\n', '\n')
|
|
if self.truncation_strategy is None:
|
|
self.truncation_strategy = 'delete'
|
|
self._set_loss_scale()
|
|
|
|
def _set_loss_scale(self):
|
|
"""For hybrid thinking models, automatically append '+ignore_empty_think' to loss_scale."""
|
|
if not self.disable_ignore_empty_think and getattr(self, 'template_meta', None) is not None:
|
|
template_meta = self.template_meta
|
|
if template_meta.is_thinking and template_meta.non_thinking_prefix:
|
|
# hybrid thinking model detected
|
|
if self.loss_scale and 'ignore_empty_think' not in self.loss_scale:
|
|
self.loss_scale = self.loss_scale + '+ignore_empty_think'
|
|
|
|
def get_template_kwargs(self):
|
|
truncation_strategy = self.truncation_strategy
|
|
if truncation_strategy == 'delete':
|
|
truncation_strategy = 'raise'
|
|
return {
|
|
'template_type': self.template,
|
|
'default_system': self.system,
|
|
'max_length': self.max_length,
|
|
'truncation_strategy': truncation_strategy,
|
|
'max_pixels': self.max_pixels,
|
|
'agent_template': self.agent_template,
|
|
'norm_bbox': self.norm_bbox,
|
|
'use_chat_template': self.use_chat_template,
|
|
'remove_unused_columns': self.remove_unused_columns, # from DataArguments
|
|
'padding_side': self.padding_side,
|
|
# train
|
|
'padding_free': self.padding_free,
|
|
'loss_scale': self.loss_scale,
|
|
'is_binary_loss_scale': self.is_binary_loss_scale,
|
|
'sequence_parallel_size': self.sequence_parallel_size,
|
|
# infer/deploy
|
|
'template_backend': self.template_backend,
|
|
# thinking
|
|
'response_prefix': self.response_prefix,
|
|
'enable_thinking': self.enable_thinking,
|
|
'preserve_thinking': self.preserve_thinking,
|
|
'add_non_thinking_prefix': self.add_non_thinking_prefix,
|
|
}
|