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