wehub-resource-sync
250 lines
14 KiB
Python
250 lines
14 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import ast
|
|
import math
|
|
import os
|
|
import torch
|
|
from dataclasses import dataclass, field
|
|
from transformers.utils import is_torch_mps_available
|
|
from typing import Any, Dict, List, Literal, Optional, Union
|
|
|
|
from swift.model import MODEL_MAPPING, get_model_info_meta, get_model_name
|
|
from swift.utils import HfConfigFactory, get_dist_setting, get_logger, json_parse_to_dict
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class ModelArguments:
|
|
"""A dataclass that holds various arguments related to model configuration and usage.
|
|
|
|
Args:
|
|
model (Optional[str]): The model ID from the Hub or a local path to the model. Defaults to None.
|
|
model_type (Optional[str]): The model type. In ms-swift, a 'model_type' groups models with the same
|
|
architecture, loading process, and template. Defaults to None, which enables auto-selection based on
|
|
the suffix of `--model` and the 'architectures' attribute in `config.json`. The `model_type` for a
|
|
corresponding model can be found in the list of supported models. Note: The concept of `model_type`
|
|
in ms-swift differs from the `model_type` in `config.json`. Custom models usually require registering
|
|
their own `model_type` and `template`.
|
|
model_revision (Optional[str]): The revision of the model. Defaults to None.
|
|
task_type (str): The task type. Can be 'causal_lm', 'seq_cls', 'embedding', 'reranker', or
|
|
'generative_reranker'. If set to 'seq_cls', you usually need to specify `--num_labels` and
|
|
`--problem_type`. Defaults to 'causal_lm'.
|
|
torch_dtype (Optional[str]): The data type of the model weights. Supports 'float16', 'bfloat16', 'float32'.
|
|
Defaults to None, in which case it's read from the 'config.json' file.
|
|
attn_impl (Optional[str]): The attention implementation to use. Options include 'sdpa', 'eager', 'flash_attn',
|
|
'flash_attention_2', 'flash_attention_3', 'flash_attention_4', etc.
|
|
Defaults to None, which means it will be read from 'config.json'.
|
|
Note: Support for these implementations depends on the model's transformers implementation.
|
|
If set to 'flash_attn' (for backward compatibility), 'flash_attention_2' will be used.
|
|
experts_impl (Optional[str]): Expert implementation type, options are 'grouped_mm', 'batched_mm', 'eager'.
|
|
Defaults to None. This feature requires "transformers>=5.0.0".
|
|
new_special_tokens (List[str]): Additional special tokens to be added to the tokenizer. Can also be a path to
|
|
a `.txt` file, where each line is a special token. Defaults to an empty list `[]`.
|
|
num_labels (Optional[int]): The number of labels for classification tasks (when `--task_type` is 'seq_cls').
|
|
Required for such tasks. Defaults to None.
|
|
problem_type (Optional[str]): The problem type for classification tasks (`--task_type` 'seq_cls'). Options are
|
|
'regression', 'single_label_classification', 'multi_label_classification'. Defaults to None, but is
|
|
automatically set to 'regression' if the model is a reward_model or `num_labels` is 1, and
|
|
'single_label_classification' otherwise.
|
|
rope_scaling (Optional[str]): The RoPE scaling type. You can pass a string like 'linear', 'dynamic', or
|
|
'yarn', and ms-swift will automatically set the corresponding `rope_scaling` and override the
|
|
'config.json' value. Alternatively, you can pass a JSON string (e.g., '{"factor":2.0, "type":"yarn"}'),
|
|
which will directly override the `rope_scaling` in 'config.json'. Defaults to None.
|
|
device_map (Optional[str]): The device map configuration for the model, e.g., 'auto', 'cpu', a JSON string,
|
|
or a path to a JSON file. This argument is passed directly to the `from_pretrained` method of transformers.
|
|
Defaults to None, and will be set automatically based on the device and distributed training settings.
|
|
max_memory (Optional[str]): The maximum memory allocation for each device when `device_map` is 'auto' or
|
|
'sequential'. Example: '{0: "20GB", 1: "20GB"}'. This argument is passed directly to the `from_pretrained`
|
|
method of transformers. Defaults to None.
|
|
max_model_len (Optional[int]): The maximum model length. This is used to calculate the RoPE scaling factor
|
|
when `rope_scaling` is specified as a string. If not None, it overrides the `max_position_embeddings`
|
|
value in 'config.json'. Defaults to None.
|
|
local_repo_path (Optional[str]): Path to a local repository for models that require a GitHub repo during
|
|
loading (e.g., deepseek-vl2). This avoids network issues during `git clone`. Defaults to None.
|
|
init_strategy (Optional[str]): The strategy to initialize all uninitialized parameters when loading a model
|
|
(especially for custom architectures). Options include 'zero', 'uniform', 'normal', 'xavier_uniform',
|
|
'xavier_normal', 'kaiming_uniform', 'kaiming_normal', 'orthogonal'. Defaults to None.
|
|
"""
|
|
model: Optional[str] = None # model id or model path
|
|
model_type: Optional[str] = field(
|
|
default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})
|
|
model_revision: Optional[str] = None
|
|
task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'reranker', 'generative_reranker'] = None
|
|
|
|
torch_dtype: Literal['bfloat16', 'float16', 'float32', None] = None
|
|
# flash_attn: It will automatically convert names based on the model.
|
|
# None: It will be automatically selected between sdpa and eager.
|
|
# 'flash_attn', 'sdpa', 'eager', 'flex_attention',
|
|
# 'flash_attention_2', 'flash_attention_3', 'flash_attention_4'
|
|
attn_impl: Optional[str] = None
|
|
experts_impl: Optional[str] = None
|
|
new_special_tokens: List[str] = field(default_factory=list)
|
|
|
|
num_labels: Optional[int] = None
|
|
problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None
|
|
rope_scaling: Optional[str] = None
|
|
device_map: Optional[Union[dict, str]] = None
|
|
max_memory: Optional[Union[dict, str]] = None
|
|
max_model_len: Optional[int] = None
|
|
# When some model code needs to be downloaded from GitHub,
|
|
# this parameter specifies the path to the locally downloaded repository.
|
|
local_repo_path: Optional[str] = None
|
|
init_strategy: Literal['zero', 'uniform', 'normal', 'xavier_uniform', 'xavier_normal', 'kaiming_uniform',
|
|
'kaiming_normal', 'orthogonal'] = None
|
|
|
|
def _init_device_map(self):
|
|
"""Prepare device map args"""
|
|
if self.device_map:
|
|
self.device_map: Union[str, Dict[str, Any], None] = json_parse_to_dict(self.device_map, strict=False)
|
|
# compat mp&ddp
|
|
_, local_rank, _, local_world_size = get_dist_setting()
|
|
if local_world_size > 1 and isinstance(self.device_map, dict) and local_rank > 0:
|
|
for k, v in self.device_map.items():
|
|
if isinstance(v, int):
|
|
self.device_map[k] += local_rank
|
|
|
|
def _init_max_memory(self):
|
|
if isinstance(self.max_memory, str):
|
|
try:
|
|
self.max_memory = ast.literal_eval(self.max_memory)
|
|
except Exception:
|
|
pass
|
|
self.max_memory = json_parse_to_dict(self.max_memory)
|
|
# compat mp&ddp
|
|
_, local_rank, _, local_world_size = get_dist_setting()
|
|
if local_world_size > 1 and isinstance(self.max_memory, dict) and local_rank > 0:
|
|
for k in list(self.max_memory.keys()):
|
|
if isinstance(k, int):
|
|
self.max_memory[k + local_rank] = self.max_memory.pop(k)
|
|
|
|
def _init_torch_dtype(self) -> None:
|
|
""""If torch_dtype is None, find a proper dtype by the config.json/GPU"""
|
|
from ..sft_args import SftArguments
|
|
|
|
self.torch_dtype: Optional[torch.dtype] = HfConfigFactory.to_torch_dtype(self.torch_dtype)
|
|
self.torch_dtype: torch.dtype = self._init_model_info()
|
|
# Mixed Precision Training
|
|
if isinstance(self, SftArguments):
|
|
self._init_mixed_precision()
|
|
|
|
def _init_mixed_precision(self):
|
|
if is_torch_mps_available():
|
|
fp16, bf16 = False, False
|
|
elif self.torch_dtype in {torch.float16, torch.float32}:
|
|
fp16, bf16 = True, False
|
|
elif self.torch_dtype == torch.bfloat16:
|
|
fp16, bf16 = False, True
|
|
else:
|
|
raise ValueError(f'args.torch_dtype: {self.torch_dtype}')
|
|
if self.fp16 is None:
|
|
self.fp16 = fp16
|
|
if self.bf16 is None:
|
|
self.bf16 = bf16
|
|
|
|
def _init_rope_scaling(self):
|
|
if self.rope_scaling:
|
|
rope_scaling: dict = json_parse_to_dict(self.rope_scaling, strict=False)
|
|
if isinstance(rope_scaling, str):
|
|
assert rope_scaling in ['linear', 'dynamic', 'yarn']
|
|
rope_scaling = {'type': rope_scaling}
|
|
else:
|
|
rope_scaling = self.model_info.rope_scaling
|
|
# reset the factor
|
|
rope_scaling.pop('factor', None)
|
|
|
|
rope_type = rope_scaling.get('rope_type', rope_scaling.get('type', 'default'))
|
|
if 'factor' not in rope_scaling and self.max_model_len is None and rope_type == 'default':
|
|
# fix megatron qwen2_5_vl
|
|
self.rope_scaling = rope_scaling
|
|
logger.info(f'Setting args.rope_scaling: {rope_scaling}')
|
|
return
|
|
|
|
# get origin_max_model_len
|
|
origin_max_model_len = None
|
|
if rope_scaling and rope_scaling.get('original_max_position_embeddings') is not None:
|
|
origin_max_model_len = rope_scaling['original_max_position_embeddings']
|
|
elif self.model_info.rope_scaling:
|
|
if self.model_info.rope_scaling.get('original_max_position_embeddings') is not None:
|
|
origin_max_model_len = self.model_info.rope_scaling['original_max_position_embeddings']
|
|
elif self.model_info.rope_scaling.get('factor') is not None:
|
|
origin_max_model_len = self.model_info.max_model_len // self.model_info.rope_scaling['factor']
|
|
if origin_max_model_len is None:
|
|
origin_max_model_len = self.model_info.max_model_len
|
|
assert origin_max_model_len is not None, '`origin_max_model_len` from model config is not set'
|
|
rope_scaling['original_max_position_embeddings'] = origin_max_model_len
|
|
|
|
if 'factor' not in rope_scaling:
|
|
assert self.max_model_len is not None, (
|
|
'max_model_len must be set if rope_scaling does not contain a "factor"')
|
|
rope_scaling['factor'] = max(float(math.ceil(self.max_model_len / origin_max_model_len)), 1.0)
|
|
rope_model_len = int(origin_max_model_len * rope_scaling['factor'])
|
|
if self.max_model_len is None:
|
|
self.max_model_len = rope_model_len
|
|
elif self.max_model_len > rope_model_len:
|
|
logger.warning(f'rope config ({rope_model_len} = {rope_scaling["factor"]} * '
|
|
f'{origin_max_model_len}) should be bigger than max_model_len '
|
|
f'from command line ({self.max_model_len})')
|
|
self.rope_scaling = rope_scaling
|
|
logger.info(f'Setting args.rope_scaling: {rope_scaling}')
|
|
logger.info(f'Setting args.max_model_len: {self.max_model_len}')
|
|
|
|
def _init_model_info(self) -> torch.dtype:
|
|
model_kwargs = self.get_model_kwargs()
|
|
if self.tuner_backend == 'unsloth':
|
|
model_kwargs['download_model'] = True
|
|
self.model_info, self.model_meta = get_model_info_meta(**model_kwargs)
|
|
self.task_type = self.model_info.task_type
|
|
self.num_labels = self.model_info.num_labels
|
|
|
|
self.model_dir = self.model_info.model_dir
|
|
self.model_type = self.model_info.model_type
|
|
if self.rope_scaling or self.model_info.rope_scaling and self.max_model_len is not None:
|
|
self._init_rope_scaling()
|
|
return self.model_info.torch_dtype
|
|
|
|
def _init_new_special_tokens(self):
|
|
if isinstance(self.new_special_tokens, str):
|
|
self.new_special_tokens = [self.new_special_tokens]
|
|
new_special_tokens = []
|
|
for token in self.new_special_tokens:
|
|
if token.endswith('.txt'):
|
|
assert os.path.isfile(token), f'special_tokens_path: {token}'
|
|
with open(token, 'r', encoding='utf-8') as f:
|
|
text = f.read()
|
|
new_special_tokens += text.split()
|
|
else:
|
|
new_special_tokens.append(token)
|
|
self.new_special_tokens = new_special_tokens
|
|
|
|
def __post_init__(self):
|
|
if self.model is None:
|
|
raise ValueError(f'Please set --model <model_id_or_path>`, model: {self.model}')
|
|
self._init_new_special_tokens()
|
|
self.model_suffix = get_model_name(self.model)
|
|
self._init_device_map()
|
|
self._init_max_memory()
|
|
self._init_torch_dtype()
|
|
|
|
def get_model_kwargs(self):
|
|
return {
|
|
'model_id_or_path': self.model,
|
|
'torch_dtype': self.torch_dtype,
|
|
'model_type': self.model_type,
|
|
'revision': self.model_revision,
|
|
'use_hf': self.use_hf,
|
|
'hub_token': self.hub_token,
|
|
'local_repo_path': self.local_repo_path,
|
|
'device_map': self.device_map,
|
|
'max_memory': self.max_memory,
|
|
'quantization_config': self.get_quantization_config(),
|
|
'attn_impl': self.attn_impl,
|
|
'experts_impl': self.experts_impl,
|
|
'new_special_tokens': self.new_special_tokens,
|
|
'rope_scaling': self.rope_scaling,
|
|
'max_model_len': self.max_model_len,
|
|
'task_type': self.task_type,
|
|
'num_labels': self.num_labels,
|
|
'problem_type': self.problem_type,
|
|
'init_strategy': self.init_strategy,
|
|
}
|