fix extra config (#847)

Author: n1ck-guo

auto_round/__main__.py

@@ -436,16 +436,17 @@ def tune(args):
     extra_config = ExtraConfig()
     tuning_config = TuningExtraConfig(
         amp=not args.disable_amp,
-        lr=args.lr,
-        minmax_lr=args.minmax_lr,
-        enable_quanted_input=not args.disable_quanted_input,
-        nblocks=args.nblocks,
+        disable_opt_rtn=args.disable_opt_rtn,
+        enable_alg_ext=args.enable_alg_ext,
         enable_minmax_tuning=not args.disable_minmax_tuning,
         enable_norm_bias_tuning=args.enable_norm_bias_tuning,
+        enable_quanted_input=not args.disable_quanted_input,
         enable_deterministic_algorithms=args.enable_deterministic_algorithms,
+        lr=args.lr,
+        minmax_lr=args.minmax_lr,
+        mem_per_param_scale=args.mem_per_param_scale,
+        nblocks=args.nblocks,
         to_quant_block_names=args.to_quant_block_names,
-        disable_opt_rtn=args.disable_opt_rtn,
-        enable_alg_ext=args.enable_alg_ext,
         scale_dtype=args.scale_dtype,
     )
     scheme_config = SchemeExtraConfig(
@@ -459,6 +460,8 @@ def tune(args):
         act_dynamic=act_dynamic,
         super_bits=args.super_bits,
         super_group_size=args.super_group_size,
+        quant_lm_head=args.quant_lm_head,
+        fp_layers=args.fp_layers,
     )
     mllm_config = MLLMExtraConfig(
         quant_nontext_module=args.quant_nontext_module, extra_data_dir=args.extra_data_dir, template=args.template
@@ -480,7 +483,6 @@ def tune(args):
         device_map=args.device_map,
         enable_torch_compile=enable_torch_compile,
         seed=args.seed,
-        fp_layers=args.fp_layers,
         not_use_best_mse=args.not_use_best_mse,
         enable_adam=args.adam,
         extra_config=extra_config,

auto_round/compressors/config.py

@@ -42,7 +42,6 @@ def __init__(
         minmax_lr: float = None,
         mem_per_param_scale: int = None,
         nblocks: int = 1,
-        quant_lm_head: bool = False,
         to_quant_block_names: Union[str, list, None] = None,
         scale_dtype: str = "fp16",
         # scheme
@@ -58,6 +57,8 @@ def __init__(
         super_bits: int = None,
         super_group_size: int = None,
         static_kv_dtype: Union[str, torch.dtype] = None,
+        quant_lm_head: bool = False,
+        fp_layers: str = None,
         # mllm
         processor: Callable = None,
         image_processor: Callable = None,
@@ -116,7 +117,6 @@ def __init__(
             minmax_lr=minmax_lr,
             mem_per_param_scale=mem_per_param_scale,
             nblocks=nblocks,
-            quant_lm_head=quant_lm_head,
             to_quant_block_names=to_quant_block_names,
             scale_dtype=scale_dtype,
         )
@@ -133,6 +133,8 @@ def __init__(
             super_bits=super_bits,
             super_group_size=super_group_size,
             static_kv_dtype=static_kv_dtype,
+            quant_lm_head=quant_lm_head,
+            fp_layers=fp_layers,
         )
         self.mllm_config = MLLMExtraConfig(
             processor=processor,
@@ -232,7 +234,6 @@ class TuningExtraConfig(BaseExtraConfig):
     minmax_lr: float = None
     mem_per_param_scale: int = None
     nblocks: int = 1
-    quant_lm_head: bool = False
     to_quant_block_names: Union[str, list, None] = None
     scale_dtype: str = "fp16"
 
@@ -251,6 +252,8 @@ class SchemeExtraConfig(BaseExtraConfig):
     super_bits: int = None
     super_group_size: int = None
     static_kv_dtype: Union[str, torch.dtype] = None
+    quant_lm_head: bool = False
+    fp_layers: str = None
 
 
 @dataclass