add target_bits but not good

Signed-off-by: xinhe3 <[email protected]>

Author: xinhe3

auto_round/utils.py

@@ -2686,20 +2686,7 @@ def _generate_recipe(
 def _generate_block_recipe(self, block, block_name, input_ids, q_input_ids, input_others):
     from itertools import combinations
 
-    # fetch mix-precision recipe configuration
-    sample_num = self.recipe_mp_config.get("sample_num", 8)
     quantizable_layers = [n for n, m in block.named_modules() if isinstance(m, SUPPORTED_LAYER_TYPES)]
-    target_bits = self.recipe_mp_config.get("target_bits", None)
-    if target_bits is None:
-        mp_ratio = self.recipe_mp_config.get("mp_ratio", 1 / 3)
-
-        # calculate the number of layers to use mix-precision
-        mp_ratio_list = [f"{i}/{len(quantizable_layers)}" for i in range(1, len(quantizable_layers))]
-        quantizable_num = int(mp_ratio * len(quantizable_layers))  # It's ceiling
-        logger.warning_once(
-            f"[Recipe Mode] {len(quantizable_layers)} layers are detected, so the available mp_ratio values are {mp_ratio_list}"
-        )
-        logger.warning_once(f"[Recipe Mode] {quantizable_num} layers of each block use the mixed precision.")
     # fetch raw low-bits dtype of block for recovering mix-precision block
     layer = get_module(block, quantizable_layers[0])
     raw_dtype = {
@@ -2722,6 +2709,35 @@ def _generate_block_recipe(self, block, block_name, input_ids, q_input_ids, inpu
         }
     )
 
+    # fetch mix-precision recipe configuration
+    sample_num = self.recipe_mp_config.get("sample_num", 8)
+    target_bits = self.recipe_mp_config.get("target_bits", None)
+    mp_ratio = self.recipe_mp_config.get("mp_ratio", None)
+    assert target_bits or mp_ratio, "Either target_bits or mp_ratio should be set in recipe_mp_config."
+    if target_bits and mp_ratio:
+        logger.warning_once("Both target_bits and mp_ratio are set in recipe_mp_config. target_bits will be used.")
+    if target_bits:
+        # get the average bits of all combinations
+        bits_of_combination = {}
+        for quantizable_num in range(len(quantizable_layers)):
+            for mp_layers in combinations(quantizable_layers, quantizable_num):
+                block = create_mp_block(block, mp_layers, self.recipe_mp_dtype)
+                # get average bits
+                avg_bits = get_avg_bits(block)
+                bits_of_combination[mp_layers] = avg_bits
+                block = recover_mp_block(block, mp_layers, raw_dtype)
+        acceptable_combination_set = {i for i in bits_of_combination if bits_of_combination[i] <= target_bits}
+    else:
+        mp_ratio = self.recipe_mp_config.get("mp_ratio", 1 / 3)
+        # calculate the number of layers to use mix-precision
+        mp_ratio_list = [f"{i}/{len(quantizable_layers)}" for i in range(1, len(quantizable_layers))]
+        quantizable_num = int(mp_ratio * len(quantizable_layers))  # It's ceiling
+        logger.warning_once(
+            f"[Recipe Mode] {len(quantizable_layers)} layers are detected, so the available mp_ratio values are {mp_ratio_list}"
+        )
+        logger.warning_once(f"[Recipe Mode] {quantizable_num} layers of each block use the mixed precision.")
+        acceptable_combination_set = combinations(quantizable_layers, quantizable_num)
+
     # generate reference output of sample input_ids
     def get_output(block, input_ids):
         output = self.get_block_outputs(
@@ -2756,19 +2772,31 @@ def get_loss(q_block, q_input_ids):
     combination_list = []
     avg_bits_list = []
     loss_list = []
-    for hp_layers in combinations(quantizable_layers, quantizable_num):
-        combination_list.append(hp_layers)
+    for mp_layers in acceptable_combination_set:
+        combination_list.append(mp_layers)
         # get loss
-        block = create_mp_block(block, hp_layers, self.recipe_mp_dtype)
+        block = create_mp_block(block, mp_layers, self.recipe_mp_dtype)
         # get average bits
         avg_bits = get_avg_bits(block)
         avg_bits_list.append(avg_bits)
         loss = get_loss(block, q_input_ids)
         loss_list.append(loss)
-        block = recover_mp_block(block, hp_layers, raw_dtype)
+        block = recover_mp_block(block, mp_layers, raw_dtype)
         if is_hpex_available():
             htcore.mark_step()
-        logger.debug(f"{hp_layers}, {loss}, {avg_bits}")
+        logger.debug(f"{mp_layers}, {loss}, {avg_bits}")
+
+    # get the worst loss
+    block = create_mp_block(block, mp_layers, self.recipe_mp_dtype)
+    # get average bits
+    avg_bits = get_avg_bits(block)
+    avg_bits_list.append(avg_bits)
+    loss = get_loss(block, q_input_ids)
+    loss_list.append(loss)
+    block = recover_mp_block(block, mp_layers, raw_dtype)
+    if is_hpex_available():
+        htcore.mark_step()
+    logger.debug(f"{mp_layers}, {loss}, {avg_bits}")
 
     # get combination with lowest loss
     best_loss = float("inf")

workspace/quantize.py

@@ -72,7 +72,6 @@ def initialize_model_and_tokenizer(model_name_or_path):
     parser.add_argument("--iters", default=None, type=int, help="iters for autoround.")
     parser.add_argument("--seqlen", default=None, type=int, help="sequence length for autoround.")
     parser.add_argument("--nsamples", default=None, type=int, help="number of samples for autoround.")
-    parser.add_argument("--target_bits", default=5, type=float, help="number of samples for autoround.")
     parser.add_argument("--target_loss_ratio", default=1.2, type=float, help="number of samples for autoround.")
     parser.add_argument(
         "--use_hpu_graph", action="store_true", help="whether to use hpu graph mode to accelerate performance"
@@ -83,6 +82,7 @@ def initialize_model_and_tokenizer(model_name_or_path):
     parser.add_argument(
         "--disable_optimum_habana", action="store_true", help="whether to use adapt_transformers_to_gaudi"
     )
+    parser.add_argument("--target_bits", default=5, type=float, help="number of samples for autoround.")
     parser.add_argument("--mp_ratio", default="1/3", type=str, help="number of samples for autoround.")
     parser.add_argument("--save", action="store_true", help="whether to save the quantized model")
     parser.add_argument("--load", action="store_true", help="whether to load the quantized model")
@@ -226,6 +226,7 @@ def match_pattern(name, pattern):
 
         recipe_results = autoround._generate_recipe(
             mp_config={
+                # "target_bits": float(args.target_bits),
                 "mp_ratio": float(eval(args.mp_ratio)),
             },
         )