fix bf16 option in AutoScheme (#1079)

Author: wenhuach21

auto_round/__main__.py

@@ -78,7 +78,9 @@ def __init__(self, *args, **kwargs):
             help="The batch size for tuning/calibration."
             "Larger batch sizes may improve stability but require more memory.",
         )
-        basic.add_argument("--avg_bits", default=None, type=float, help="for auto scheme, number of avg weight bits")
+        basic.add_argument(
+            "--avg_bits", "--target_bits", default=None, type=float, help="for auto scheme, number of avg weight bits"
+        )
         basic.add_argument(
             "--options", default=None, type=str, help="for auto scheme, options for auto scheme, e.g. 'W4A16,W8A16'"
         )

auto_round/compressors/base.py

@@ -427,10 +427,6 @@ def _gen_auto_scheme(
 
         all_dtypes = []
         for option in scheme.options:
-            # Skip pure BF16 option
-            if option == "BF16":
-                continue
-
             # Resolve the quantization scheme or data type
             dtype = "int"
             if isinstance(option, str):

test/test_cpu/test_autoround.py

@@ -713,14 +713,22 @@ def test_invalid_layer_config(self):
 
     def test_quant_lm_head(self):
         model_name = "/tf_dataset/auto_round/models/Qwen/Qwen3-8B"
-        ar = AutoRound(model_name, quant_lm_head=True, iters=0, disable_opt_rtn=True)
+        ar = AutoRound(model_name, quant_lm_head=True, iters=0, seqlen=8, nsamples=1, disable_opt_rtn=True)
         ar.quantize_and_save(output_dir=self.save_folder, format="auto_round")
         model = AutoModelForCausalLM.from_pretrained(self.save_folder, device_map="cpu")
         assert "lm_head" in model.config.quantization_config.extra_config
         assert model.config.quantization_config.extra_config["lm_head"]["bits"] == 4
 
         layer_config = {"lm_head": {"bits": 4}}
-        ar = AutoRound(model_name, quant_lm_head=False, iters=0, disable_opt_rtn=True, layer_config=layer_config)
+        ar = AutoRound(
+            model_name,
+            quant_lm_head=False,
+            iters=0,
+            seqlen=8,
+            nsamples=1,
+            disable_opt_rtn=True,
+            layer_config=layer_config,
+        )
         ar.quantize_and_save(output_dir=self.save_folder, format="auto_round")
         model = AutoModelForCausalLM.from_pretrained(self.save_folder, device_map="cpu")
         assert "lm_head" in model.config.quantization_config.extra_config
@@ -729,7 +737,15 @@ def test_quant_lm_head(self):
     def test_quant_lm_head_layer_config(self):
         model_name = "/tf_dataset/auto_round/models/Qwen/Qwen3-8B"
         layer_config = {"lm_head": {"bits": 4}}
-        ar = AutoRound(model_name, quant_lm_head=True, iters=0, disable_opt_rtn=True, layer_config=layer_config)
+        ar = AutoRound(
+            model_name,
+            quant_lm_head=True,
+            iters=0,
+            seqlen=8,
+            nsamples=1,
+            disable_opt_rtn=True,
+            layer_config=layer_config,
+        )
         ar.quantize_and_save(output_dir=self.save_folder, format="auto_round")
         model = AutoModelForCausalLM.from_pretrained(self.save_folder, device_map="cpu")
         assert "lm_head" in model.config.quantization_config.extra_config