[GGUF support step3]patch for double quant (#473)

Author: n1ck-guo

.azure-pipelines/scripts/ut/run_ut.sh

@@ -6,6 +6,8 @@ echo "set up UT env..."
 pip install pytest-cov pytest-html
 pip install -r /auto-round/test/requirements.txt
 pip list
+# install latest gguf for ut test
+git clone https://github.com/ggml-org/llama.cpp.git && cd llama.cpp/gguf-py && pip install .
 
 cd /auto-round/test || exit 1
 find . -type f -exec sed -i '/sys\.path\.insert(0, "\.\.")/d' {} +

auto_round/autoround.py

@@ -14,6 +14,7 @@
 
 import os
 import re
+import sys
 
 import torch
 import copy
@@ -350,7 +351,17 @@ def _set_device_for_matching_module(self, name, device):
         else:
             module.tuning_device = device
 
+    def _dq_check(self):
+        """Reset the default value of super_bits and super_group_size"""
+        from auto_round.export.export_to_gguf.config import GGUF_CONFIG
+        if self.data_type.endswith("_dq"):
+            gguf_config = GGUF_CONFIG[f"gguf:q{self.bits}_k_s"]
+            self.super_bits = gguf_config["super_bits"] if self.super_bits is None else self.super_bits
+            self.super_group_size = gguf_config["super_group_size"] \
+                if self.super_group_size is None else self.super_group_size
+            
     def check_configs(self):
+
         """Checks if the configurations are valid.
 
         Raises:
@@ -392,6 +403,7 @@ def check_configs(self):
                     f"reset gradient_accumulate_steps to {self.gradient_accumulate_steps}"
                     f" as nsamples must equal or greater"
                     f" than gradient_accumulate_steps * batch_size")
+        self._dq_check()
 
     # def _check_format_compatibility(self, format):  ##TODO
     #     ##check lm_head, mixed_bits, bits, each layer supporting, etc
@@ -491,9 +503,11 @@ def remove_duplicates(lst):
             save_format_ = format.replace(":", "-").replace("_", "-")
             save_folder = os.path.join(output_dir, save_format_) if len(formats) > 1 else output_dir
             self.save_quantized(save_folder, format=format, inplace=inplace, **kwargs)
+            
             folders.append(save_folder)
 
         return model, folders
+
     def quantize(self):
         """Quantize the model and return the quantized model along with layer configurations.
         the entry of AutoRound.
@@ -1538,6 +1552,13 @@ def save_quantized(self, output_dir=None, format="auto_round", inplace=True, **k
                     )
                     format = "auto_round"
 
+        if re.search("q\d_k", format) and not self.data_type.endswith("_dq"):
+            logger.error(
+                f"datatype<{self.data_type}> not support to export {format} format."
+                " Please change export format or data_type."
+            )
+            sys.exit(-1)
+
         if self.low_cpu_mem_usage:
             self.model = self.model.to('cpu')
 

auto_round/export/export_to_gguf/config.py

@@ -18,24 +18,24 @@
 
 GGUF_CONFIG["gguf:q4_1"] = {"bits": 4, "act_bits": 16, "group_size": 32, "asym": True, "data_type": "int"}
 
-# GGUF_CONFIG["gguf:q4_k_s"] = {
-#     "bits": 4,
-#     "act_bits": 16,
-#     "super_group_size": 8,
-#     "super_bits": 6,
-#     "group_size": 32,
-#     "asym": True,
-#     "data_type": "int_asym_dq"
-# }
+GGUF_CONFIG["gguf:q4_k_s"] = {
+    "bits": 4,
+    "act_bits": 16,
+    "super_group_size": 8,
+    "super_bits": 6,
+    "group_size": 32,
+    "asym": True,
+    "data_type": "int_asym_dq"
+}
 
-# GGUF_CONFIG["gguf:q2_k_s"] = {
-#     "bits": 2,
-#     "act_bits": 16,
-#     "super_group_size": 16,
-#     "super_bits": 4,
-#     "group_size": 16,
-#     "asym": True,
-#     "data_type": "int_asym_dq"
-# }
+GGUF_CONFIG["gguf:q2_k_s"] = {
+    "bits": 2,
+    "act_bits": 16,
+    "super_group_size": 16,
+    "super_bits": 4,
+    "group_size": 16,
+    "asym": True,
+    "data_type": "int_asym_dq"
+}
 
 GGUF_CONFIG["gguf:q8_0"] = {"bits": 8, "act_bits": 16, "group_size": 32, "asym": False, "data_type": "int"}

auto_round/export/export_to_gguf/convert.py

@@ -1149,7 +1149,19 @@ def _quant_data(data, data_qtype):
                         layer_name = name[:-len(suffix)]
                         module = get_module(self.model, layer_name)
                         if hasattr(module, "scale"):
+
+                            if hasattr(self, "permute"):
+                                bs = module.scale.shape[0]
+                                for attr in ["scale", "zp", "w_d_scale", "w_d_wmin_m", "w_wmin_m"]:
+                                    if hasattr(module, attr) and getattr(module, attr) is not None:
+                                        attr_tensor = getattr(module, attr)
+                                        ori_shape = attr_tensor.shape
+                                        attr_tensor = self.modify_tensors(attr_tensor.reshape(bs, -1), name, bid)[0][1]
+                                        attr_tensor = attr_tensor.reshape(ori_shape)
+                                        setattr(module, attr, attr_tensor)
+
                             scale = module.scale
+                            
                             if isinstance(scale, torch.Tensor):
                                 scale = scale.numpy()
                             zp = module.zp if hasattr(module, "zp") else None

auto_round/export/export_to_gguf/export.py

@@ -31,8 +31,8 @@
         "q8_0": gguf.LlamaFileType.MOSTLY_Q8_0,
         "q4_0": gguf.LlamaFileType.MOSTLY_Q4_0,
         "q4_1": gguf.LlamaFileType.MOSTLY_Q4_1,
-        # "q4_k_s": gguf.LlamaFileType.MOSTLY_Q4_K_S,
-        # "q2_k_s": gguf.LlamaFileType.MOSTLY_Q2_K_S,
+        "q4_k_s": gguf.LlamaFileType.MOSTLY_Q4_K_S,
+        "q2_k_s": gguf.LlamaFileType.MOSTLY_Q2_K_S,
         "q8_0": gguf.LlamaFileType.MOSTLY_Q8_0,
         "auto": gguf.LlamaFileType.GUESSED,
     }

auto_round/export/export_to_gguf/quant.py

@@ -39,14 +39,19 @@ def register(cls):
     return register
 
 
-def ggml_quant(data: np.array, ggml_type, scale=None, zp=None, wmin_m=None, d_scale=None, d_wmin_m=None, worker=16):
+def ggml_quant(data: np.array, ggml_type, scale=None, zp=None, wmin_m=None, d_scale=None, d_wmin_m=None):
     block_size, type_size = GGML_QUANT_SIZES[ggml_type]
 
     data = data.astype(np.float32, copy=False)
     shape = data.shape
     n_blocks = data.size // block_size
     blocks = data.reshape((n_blocks, block_size))
 
+    if ggml_type.endswith("_k"):
+        worker = 16
+    else:
+        worker = 0
+
     if worker > 0:
         n_groups = (data.shape[0] // worker) or 1
         blocks = np.array_split(blocks, n_groups, axis=0)

auto_round/script/llm.py

@@ -311,61 +311,13 @@ def setup_eval_parser():
     return args
 
 
-def _gguf_args_check(args):
-    from auto_round.utils import logger
-    from auto_round.export.export_to_gguf.config import GGUF_CONFIG
-
-    formats = args.format.lower().replace(' ', '').split(",")
-    for format in GGUF_CONFIG:
-        if format in formats:
-            from pathlib import Path
-            from auto_round.export.export_to_gguf.convert import Model
-            hparams = Model.load_hparams(Path(args.model))
-            model_architecture = hparams["architectures"][0]
-            try:
-                model_class = Model.from_model_architecture(model_architecture)
-            except NotImplementedError:
-                logger.error(f"Model {model_architecture} is not supported to export GGUF format")
-                sys.exit(1)
-
-            if format.endswith("_k") and ("hidden_size" in hparams and hparams["hidden_size"] % 256 != 0):
-                model_name = args.model.split('/')
-                model_name = model_name[-1] if model_name[-1] else model_name[-2]
-                hidden_size = hparams["hidden_size"]
-                logger.error(
-                    f"Currently only support pure mode for format: {format}. "
-                    f"{model_name} is not supported, cause hidden_size({hidden_size}) % 256 !=0")
-                sys.exit(-1)
-
-            unsupport_list, reset_list = [], []
-            gguf_config = GGUF_CONFIG[format]
-            for k, v in gguf_config.items():
-                if getattr(args, k) != v:
-                    unsupport_list.append(f"{k}={getattr(args, k)}")
-                    reset_list.append(f"{k}={v}")
-                    setattr(args, k, v)
-            if len(unsupport_list) > 0:
-                if len(formats) > 1:
-                    logger.error(
-                        f"format {format} not support for {', '.join(unsupport_list)},"
-                        f" please reset to {', '.join(reset_list)}, and retry")
-                    exit(-1)
-                else:
-                    logger.error(
-                        f"format {format} not support for {', '.join(unsupport_list)},"
-                        f" reset to {', '.join(reset_list)}.")
-            logger.info(f"export format {format}, sym = {not args.asym}, group_size = {args.group_size}")
-
-    return args
-
-
 def tune(args):
     import transformers
 
     from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel, AutoConfig
 
     from auto_round.utils import detect_device, get_library_version
-    from auto_round.utils import logger
+    from auto_round.utils import logger, _gguf_args_check
 
     tasks = args.tasks
     if args.format is None:
@@ -602,9 +554,12 @@ def tune(args):
 
         if args.act_bits <= 8 or eval_gguf_model:
             if eval_gguf_model:
+                # gguf floder only contains one file
                 for file in os.listdir(eval_folder):
                     gguf_file = file
-                user_model = AutoModelForCausalLM.from_pretrained(eval_folder, gguf_file=gguf_file, device_map="auto")
+                model = AutoModelForCausalLM.from_pretrained(
+                    eval_folder, gguf_file=gguf_file, device_map="auto" if use_auto_mapping else None)
+                tokenizer = AutoTokenizer.from_pretrained(eval_folder, gguf_file=gguf_file)
             else:
                 if hasattr(model, "hf_device_map") and len(model.hf_device_map) > 1:
                     from accelerate.big_modeling import dispatch_model
@@ -616,7 +571,8 @@ def tune(args):
                     user_model = model.to(device_str)
 
             if args.eval_task_by_task:
-                eval_task_by_task(user_model, device=device_str, tasks=args.tasks, batch_size=args.eval_bs)
+                eval_task_by_task(
+                    user_model, tokenizer=tokenizer, device=device_str, tasks=args.tasks, batch_size=args.eval_bs)
             else:
                 if args.eval_bs is None or args.eval_bs == "auto":
                     logger.warning("This API does not support auto currently, reset eval_bs to 16")
@@ -660,7 +616,8 @@ def eval(args):
     print(make_table(res))
 
 
-def eval_task_by_task(model, device, tasks, tokenizer=None, batch_size=None, max_batch_size=64, trust_remote_code=True):
+def eval_task_by_task(
+        model, device=None, tasks=None, tokenizer=None, batch_size=None, max_batch_size=64, trust_remote_code=True):
     set_cuda_visible_devices(device)
     device_str, parallelism = get_device_and_parallelism(device)
 

auto_round/script/mllm.py

@@ -22,7 +22,9 @@
     is_debug_mode,
     get_device_and_parallelism,
     set_cuda_visible_devices,
-    logger)
+    logger,
+    _gguf_args_check
+    )
 
 
 os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
@@ -278,57 +280,6 @@ def setup_lmeval_parser():
     args = parser.parse_args()
     return args
 
-def _gguf_args_check(args):
-    from auto_round.utils import logger
-
-    _GGUF_CONFIG = {
-        "gguf:q4_0": {
-            "bits": 4,
-            "act_bits": 16,
-            "group_size": 32,
-            "asym": False,
-        },
-        "gguf:q4_1": {
-            "bits": 4,
-            "act_bits": 16,
-            "group_size": 32,
-            "asym": True,
-        }
-    }
-
-    formats = args.format.lower().replace(' ', '').split(",")
-    for format in _GGUF_CONFIG:
-        if format in formats:
-            from pathlib import Path
-            from auto_round.export.export_to_gguf.convert import Model
-            hparams = Model.load_hparams(Path(args.model))
-            model_architecture = hparams["architectures"][0]
-            try:
-                model_class = Model.from_model_architecture(model_architecture)
-            except NotImplementedError:
-                logger.error(f"Model {model_architecture} is not supported to export GGUF format")
-                sys.exit(1)
-
-            unsupport_list, reset_list = [], []
-            gguf_config = _GGUF_CONFIG[format]
-            for k, v in gguf_config.items():
-                if getattr(args, k) != v:
-                    unsupport_list.append(f"{k}={getattr(args, k)}")
-                    reset_list.append(f"{k}={v}")
-                    setattr(args, k, v)
-            if len(unsupport_list) > 0:
-                if len(formats) > 1:
-                    logger.error(
-                        f"format {format} not support for {', '.join(unsupport_list)},"
-                        f" please reset to {', '.join(reset_list)}, and retry")
-                    exit(-1)
-                else:
-                    logger.error(
-                        f"format {format} not support for {', '.join(unsupport_list)},"
-                        f" reset to {', '.join(reset_list)}.")
-            logger.info(f"export format {format}, sym = {not args.asym}, group_size = {args.group_size}")
-
-    return args
 
 def tune(args):
     import transformers

auto_round/utils.py

@@ -1204,3 +1204,56 @@ def is_debug_mode():
         bool: True if debugging is enabled, False otherwise.
     """
     return sys.gettrace() is not None or sys.flags.debug == 1
+
+
+def _gguf_args_check(args):
+    from auto_round.utils import logger
+    from auto_round.export.export_to_gguf.config import GGUF_CONFIG
+
+    formats = args.format.lower().replace(' ', '').split(",")
+    formats = sorted(formats, key=lambda x:len(x))
+    pattern = re.compile("q\d_k")
+    pre_dq_format = ""
+    for format in GGUF_CONFIG:
+        if format in formats:
+            if re.search(pattern, format):
+                if pre_dq_format and re.search(pattern, format).group() not in pre_dq_format:
+                    logger.error(f"Cannot eport {pre_dq_format} and {format} at the same time.")
+                    sys.exit(-1)
+                else:
+                    pre_dq_format = format
+            
+            if os.path.isdir(args.model):
+                from pathlib import Path
+                from auto_round.export.export_to_gguf.convert import Model
+                hparams = Model.load_hparams(Path(args.model))
+                model_architecture = hparams["architectures"][0]
+                try:
+                    model_class = Model.from_model_architecture(model_architecture)
+                except NotImplementedError:
+                    logger.error(f"Model {model_architecture} is not supported to export GGUF format")
+                    sys.exit(1)
+
+                if re.search(pattern, format) and ("hidden_size" in hparams and hparams["hidden_size"] % 256 !=0):
+                    model_name = args.model.split('/')
+                    model_name = model_name[-1] if model_name[-1] else model_name[-2]
+                    hidden_size = hparams["hidden_size"]
+                    logger.error(
+                        f"Currently only support pure mode for format: {format}. "
+                        f"{model_name} is not supported, cause hidden_size({hidden_size}) % 256 !=0")
+                    sys.exit(-1)
+
+            unsupport_list, reset_list = [], []
+            gguf_config = GGUF_CONFIG[format]
+            for k, v in gguf_config.items():
+                if getattr(args, k) != v:
+                    unsupport_list.append(f"{k}={getattr(args, k)}")
+                    reset_list.append(f"{k}={v}")
+                    setattr(args, k, v)
+            if len(unsupport_list) > 0:
+                logger.error(
+                    f"format {format} not support for {', '.join(unsupport_list)},"
+                    f" reset to {', '.join(reset_list)}.")
+            logger.info(f"export format {format}, sym = {not args.asym}, group_size = {args.group_size}")
+
+    return args

test/requirements.txt

@@ -1,2 +1,3 @@
 addict
 modelscope
+gguf