| |
| import os |
| from dataclasses import dataclass |
| from typing import Literal, Optional |
|
|
| import torch |
| import torch.distributed as dist |
|
|
| from swift.utils import get_logger, init_process_group, set_default_ddp_config |
| from .base_args import BaseArguments, to_abspath |
| from .merge_args import MergeArguments |
|
|
| logger = get_logger() |
|
|
|
|
| @dataclass |
| class ExportArguments(MergeArguments, BaseArguments): |
| """ |
| ExportArguments is a dataclass that inherits from BaseArguments and MergeArguments. |
| |
| Args: |
| output_dir (Optional[str]): Directory where the output will be saved. |
| quant_n_samples (int): Number of samples for quantization. |
| max_length (int): Sequence length for quantization. |
| quant_batch_size (int): Batch size for quantization. |
| to_ollama (bool): Flag to indicate export model to ollama format. |
| push_to_hub (bool): Flag to indicate if the output should be pushed to the model hub. |
| hub_model_id (Optional[str]): Model ID for the hub. |
| hub_private_repo (bool): Flag to indicate if the hub repository is private. |
| commit_message (str): Commit message for pushing to the hub. |
| to_peft_format (bool): Flag to indicate if the output should be in PEFT format. |
| This argument is useless for now. |
| """ |
| output_dir: Optional[str] = None |
|
|
| |
| quant_method: Literal['awq', 'gptq', 'bnb'] = None |
| quant_n_samples: int = 256 |
| max_length: int = 2048 |
| quant_batch_size: int = 1 |
| group_size: int = 128 |
|
|
| |
| to_ollama: bool = False |
|
|
| |
| to_mcore: bool = False |
| to_hf: bool = False |
| mcore_model: Optional[str] = None |
| thread_count: Optional[int] = None |
| test_convert_precision: bool = False |
|
|
| |
| push_to_hub: bool = False |
| |
| hub_model_id: Optional[str] = None |
| hub_private_repo: bool = False |
| commit_message: str = 'update files' |
| |
| to_peft_format: bool = False |
| exist_ok: bool = False |
|
|
| def _init_output_dir(self): |
| if self.output_dir is None: |
| ckpt_dir = self.ckpt_dir or f'./{self.model_suffix}' |
| ckpt_dir, ckpt_name = os.path.split(ckpt_dir) |
| if self.to_peft_format: |
| suffix = 'peft' |
| elif self.quant_method: |
| suffix = f'{self.quant_method}-int{self.quant_bits}' |
| elif self.to_ollama: |
| suffix = 'ollama' |
| elif self.merge_lora: |
| suffix = 'merged' |
| elif self.to_mcore: |
| suffix = 'mcore' |
| elif self.to_hf: |
| suffix = 'hf' |
| else: |
| return |
|
|
| self.output_dir = os.path.join(ckpt_dir, f'{ckpt_name}-{suffix}') |
|
|
| self.output_dir = to_abspath(self.output_dir) |
| if not self.exist_ok and os.path.exists(self.output_dir): |
| raise FileExistsError(f'args.output_dir: `{self.output_dir}` already exists.') |
| logger.info(f'args.output_dir: `{self.output_dir}`') |
|
|
| def __post_init__(self): |
| if self.quant_batch_size == -1: |
| self.quant_batch_size = None |
| if self.quant_bits and self.quant_method is None: |
| raise ValueError('Please specify the quantization method using `--quant_method awq/gptq/bnb`.') |
| if self.quant_method and self.quant_bits is None: |
| raise ValueError('Please specify `--quant_bits`.') |
| if self.quant_method in {'gptq', 'awq'} and self.torch_dtype is None: |
| self.torch_dtype = torch.float16 |
| if self.to_mcore or self.to_hf: |
| self.mcore_model = to_abspath(self.mcore_model, check_path_exist=True) |
| if not dist.is_initialized(): |
| set_default_ddp_config() |
| init_process_group() |
|
|
| BaseArguments.__post_init__(self) |
| self._init_output_dir() |
| if self.quant_method in {'gptq', 'awq'} and len(self.dataset) == 0: |
| raise ValueError(f'self.dataset: {self.dataset}, Please input the quant dataset.') |
|
|
