From b8747edf27deee27018ef96153ac8921145ca052 Mon Sep 17 00:00:00 2001 From: xhrxgr <1749567727@qq.com> Date: Sun, 10 May 2026 18:26:09 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E9=87=8F=E5=8C=96=E9=85=8D=E7=BD=AE?= =?UTF-8?q?=E6=9C=AA=E4=BC=A0=E9=80=92=E7=BB=99=20LLaMA-Factory=20?= =?UTF-8?q?=E5=AF=BC=E8=87=B4=E9=80=80=E5=8C=96=E4=B8=BA=E5=85=A8=E7=B2=BE?= =?UTF-8?q?=E5=BA=A6=E5=8A=A0=E8=BD=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 QuantizationArgs 配置类，与 LLaMA-Factory QuantizationArguments 对齐 - 移除 load_in_4bit/load_in_8bit（由 quantization_bit 内部处理） - 新增 quantization_method 支持多种量化后端（bnb, gptq, awq 等） - 修复 dict_to_argv 跳过 None 值，避免 HfArgumentParser 报错 - 使用 exclude_none=True 防止空字段传递给 run_exp() - 训练和推理时自动展平嵌套的量化配置 --- weclone/train/train_sft.py | 6 +++- weclone/utils/config.py | 29 ++++++++++++++++++-- weclone/utils/config_models.py | 50 ++++++++++++++++++++++++++++++++-- weclone/utils/tools.py | 5 ++-- 4 files changed, 81 insertions(+), 9 deletions(-) diff --git a/weclone/train/train_sft.py b/weclone/train/train_sft.py index 2c498ea..5a9732f 100644 --- a/weclone/train/train_sft.py +++ b/weclone/train/train_sft.py @@ -40,7 +40,11 @@ def main(): formatted_config = json.dumps(train_config.model_dump(mode="json"), indent=4, ensure_ascii=False) logger.info(f"Fine-tuning configuration:\n{formatted_config}") - run_exp(train_config.model_dump(mode="json")) + # Build config dict and remove nested 'quantization' key (its fields are already flattened at top level) + config_dict = train_config.model_dump(mode="json", exclude_none=True) + config_dict.pop("quantization", None) + + run_exp(config_dict) if __name__ == "__main__": diff --git a/weclone/utils/config.py b/weclone/utils/config.py index e578447..53825a9 100644 --- a/weclone/utils/config.py +++ b/weclone/utils/config.py @@ -47,6 +47,17 @@ def load_base_config() -> WcConfig: return wc_config +def _flatten_quantization_args(train_sft_args) -> dict: + """Extract quantization sub-config and return as a flat dict with non-None values. + + LLaMA-Factory expects quantization params at the top level of the argument + namespace (e.g. ``--quantization_bit 4``), not nested under a ``quantization`` + prefix. This helper flattens the nested ``QuantizationArgs`` model so it can + be merged directly into the config dict passed to HfArgumentParser. + """ + return train_sft_args.quantization.get_non_none_dict() + + def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel: """Create corresponding configuration object based on argument type, merge common_config""" if arg_type == "cli_args": @@ -55,7 +66,13 @@ def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel: common_config = wc_config.common_args.model_dump() if arg_type == "web_demo" or arg_type == "api_service": - config_dict = {**common_config, **wc_config.infer_args.model_dump()} + # Inherit quantization settings from train_sft_args for inference + quant_dict = _flatten_quantization_args(wc_config.train_sft_args) + config_dict = { + **common_config, + **wc_config.infer_args.model_dump(), + **quant_dict, + } return WCInferConfig(**config_dict) elif arg_type == "vllm": @@ -66,12 +83,18 @@ def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel: elif arg_type == "train_sft": common_config["include_type"] = wc_config.make_dataset_args.include_type - config_dict = {**common_config, **wc_config.train_sft_args.model_dump()} + + # Merge training params; flatten the nested quantization sub-config + train_dict = wc_config.train_sft_args.model_dump() + # Remove the nested "quantization" dict — its fields are flattened below + train_dict.pop("quantization", None) + train_dict.update(_flatten_quantization_args(wc_config.train_sft_args)) + + config_dict = {**common_config, **train_dict} return WCTrainSftConfig(**config_dict) elif arg_type == "make_dataset": make_dataset_config = wc_config.make_dataset_args.model_dump() - # TODO: Should the following three parameters be moved to common? train_sft_args = wc_config.train_sft_args extra_values = { "dataset": train_sft_args.dataset, diff --git a/weclone/utils/config_models.py b/weclone/utils/config_models.py index 085402c..6d94218 100644 --- a/weclone/utils/config_models.py +++ b/weclone/utils/config_models.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, List, Literal, Optional from loguru import logger from pydantic import BaseModel, Field, model_validator @@ -135,7 +135,7 @@ class MakeDatasetArgs(BaseConfigModel): max_image_num: int = Field(2, description="Maximum number of images per single data entry") blocked_words: List[str] = Field([], description="List of blocked words") add_time: bool = Field(False, description="Whether to add time to the dataset") - add_relation: bool = Field(False, description="Whether to add chat member relationship to the dataset") + add_relation: bool = Field(False, description="Whether to add chat relation to the dataset") single_combine_strategy: CombineStrategy = Field( CombineStrategy.TIME_WINDOW, description="Strategy for combining single person's messages into a single sentence", @@ -163,6 +163,40 @@ class MakeDatasetArgs(BaseConfigModel): ) clean_batch_size: int = Field(10, description="Batch size for data cleaning") vision_api: VisionApiConfig = Field(VisionApiConfig()) + pure_text: bool = Field(False, description="Whether to use pure text mode (disable vision encoder)") + + +class QuantizationArgs(BaseConfigModel): + """Quantization arguments aligned with LLaMA-Factory QuantizationArguments. + + These parameters are passed directly to LLaMA-Factory's HfArgumentParser + for both training and inference. LLaMA-Factory internally maps + ``quantization_bit`` to ``BitsAndBytesConfig(load_in_4bit/load_in_8bit)``, + so there is no need to expose ``load_in_4bit`` / ``load_in_8bit`` directly. + + Reference: LLaMA-Factory src/llamafactory/hparams/model_args.py QuantizationArguments + """ + + quantization_method: Optional[str] = Field( + None, + description="Quantization method: bnb, gptq, awq, aqlm, quanto, eetq, hqq, mxfp4, fp8", + ) + quantization_bit: Optional[int] = Field( + None, + description="Number of bits for on-the-fly quantization (e.g. 4 or 8)", + ) + quantization_type: Optional[Literal["nf4", "fp4"]] = Field( + None, + description="Quantization data type for bitsandbytes int4 training: nf4 or fp4", + ) + double_quantization: Optional[bool] = Field( + None, + description="Whether to use double quantization in bitsandbytes int4 training", + ) + + def get_non_none_dict(self) -> dict: + """Return only the non-None fields as a dict, for merging into other configs.""" + return {k: v for k, v in self.model_dump().items() if v is not None} class TrainSftArgs(BaseConfigModel): @@ -190,6 +224,10 @@ class TrainSftArgs(BaseConfigModel): plot_loss: bool = Field(True, description="Whether to plot loss curve") fp16: bool = Field(True, description="Whether to use fp16") flash_attn: str = Field("fa2", description="Flash Attention type") + quantization: QuantizationArgs = Field( + default_factory=QuantizationArgs, + description="Quantization settings for on-the-fly quantization (QLoRA, etc.)", + ) preprocessing_num_workers: int = Field(16, description="Number of preprocessing worker processes") dataloader_num_workers: int = Field(4, description="Number of dataloader worker processes") deepspeed: Optional[str] = Field( @@ -207,6 +245,12 @@ class InferArgs(BaseConfigModel): class VllmArgs(BaseConfigModel): gpu_memory_utilization: float = Field(default=0.9, description="vllm GPU memory utilization") + quantization: Optional[str] = Field( + default=None, description="Quantization method for vLLM, e.g. 'awq', 'gptq'" + ) + load_format: Optional[str] = Field( + default=None, description="Format for loading weights, e.g. 'awq', 'gptq'" + ) class TestModelArgs(BaseConfigModel): @@ -235,7 +279,7 @@ class WcConfig(BaseModel): class WCInferConfig(CommonArgs, InferArgs): - """Final configuration model for Web Demo""" + """Final configuration model for Web Demo / API Service (based on LLaMA-Factory ChatModel)""" pass diff --git a/weclone/utils/tools.py b/weclone/utils/tools.py index 3c2fb11..9b0d86e 100644 --- a/weclone/utils/tools.py +++ b/weclone/utils/tools.py @@ -1,7 +1,8 @@ def dict_to_argv(d): argv = [] for k, v in d.items(): + if v is None: + continue argv.append("--" + k) - if v is not None: - argv.append(str(v)) + argv.append(str(v)) return argv