Support AWQ & GroupWiseQuant for LLMs (#7688)

* Support AWQ & GroupWiseQuant for LLMs

* add docs

* add docs

* upadte

Author: RachelXu7

llm/argument.py

@@ -156,13 +156,20 @@ class QuantArgument:
     do_ptq: bool = field(default=False, metadata={"help": "Whether to use PTQ"})
     ptq_step: int = field(default=32, metadata={"help": "Step for PTQ"})
 
+    weight_quant_method: str = field(
+        default="abs_max_channel_wise",
+        metadata={"help": "Weight quantization method, choosen from ['abs_max_channel_wise', 'groupwise']"},
+    )
+
+    # Pre-quant method Shift related parameters
     shift: bool = field(default=False, metadata={"help": "Whether to use Shift"})
     shift_all_linears: bool = field(default=False, metadata={"help": "Whether to shift all linears"})
     shift_sampler: str = field(
         default="ema", metadata={"help": "The name of shift sampler, choosen from ['ema', 'none']"}
     )
     shift_step: int = field(default=32, metadata={"help": "Sample steps when shift"})
 
+    # Pre-quant methos Smooth related parameters
     smooth: bool = field(default=False, metadata={"help": "Whether to use Smooth"})
     smooth_all_linears: bool = field(default=False, metadata={"help": "Whether to smooth all linears"})
     smooth_sampler: str = field(
@@ -179,6 +186,12 @@ class QuantArgument:
     do_gptq: bool = field(default=False, metadata={"help": "Whether to use GPTQ"})
     gptq_step: int = field(default=8, metadata={"help": "Step for GPTQ"})
 
+    # AWQ related parameters, default for WINT4
+    do_awq: bool = field(default=False, metadata={"help": "Whether to use AWQ Search"})
+    auto_clip: bool = field(default=False, metadata={"help": "Whether to use AutoClip from AWQ"})
+    awq_step: int = field(default=8, metadata={"help": "Step for AWQ Search"})
+    autoclip_step: int = field(default=8, metadata={"help": "Step for AutoClip"})
+
 
 @dataclass
 class GenerateArgument:

llm/docs/quantization.md

@@ -3,12 +3,14 @@
 ## 1.算法介绍
 
 大模型量化将16位、32位浮点数的模型参数或激活量化为4位或8位整数能够有效降低模型存储空间和计算资源需求，同时加速推理速度。工具链量化算法包含：
-- **PTQ**。PaddleSlim 团队自研的自适应Shift-SmoothQuant量化算法，在[SmoothQuant](https://arxiv.org/abs/2211.10438)和[Outlier Suppression+](https://arxiv.org/abs/2304.09145)基础上
+- **PTQ**。PaddleSlim 团队自研的自适应PiecewiseSearchSmooth(PSS)量化算法，在[SmoothQuant](https://arxiv.org/abs/2211.10438)和[Outlier Suppression+](https://arxiv.org/abs/2304.09145)基础上
 新增PieceWiseSearch参数搜索算法并将算法扩展至**所有线性层**，对模型权重和激活分布进行调整，减少后续A8W8 PTQ量化损失。
 
 
 - **GPTQ**。[GPTQ](https://arxiv.org/abs/2210.17323)是业界主流的权重量化算法，可以将大模型权重进行4位整数无损量化，提高模型推理速度。
 
+- **AWQ**。[GPTQ](https://arxiv.org/abs/2306.00978)是业界主流的权重量化算法，可以将大模型权重进行4位整数无损量化，提高模型推理速度。
+
 <div align="center">
     <img width="800" alt="llm" src="https://github.com/PaddlePaddle/PaddleNLP/assets/63761690/fe8f941b-4b35-48ca-814f-96533d7e24ce">
 </div>
@@ -65,12 +67,19 @@ python  finetune_generation.py ./llama/ptq_argument.json
 python  finetune_generation.py ./llama/gptq_argument.json
 ```
 
-### 2.5 量化参数介绍
+### 2.5 AWQ 量化
+
+```
+python  finetune_generation.py ./llama/awq_argument.json
+```
+
+### 2.6 量化参数介绍
 
 <summary>&emsp; 量化参数（QuantArgument）</summary><div>
 
 - `quant_type`: PTQ,QAT量化类型，默认为A8W8。支持A8W8,WINT4，WINT8：A8W8指对激活（输入）进行INT8量化，对模型权重进行INT8量化；WINT4指仅对模型权重进行INT4量化，后续使用WeightOnly进行推理；WINT8指仅对模型权重进行INT8量化，后续使用WeightOnly进行推理。
 - `do_ptq`: 是否进行PTQ量化，默认为False。
+- `weight_quant_method`: 权重量化方式，现可选groupwise或者abs_max_channel_wise。
 - `ptq_step`: PTQ量化步数，也即模型前向次数，默认为32。
 - `shift`: 是否在PTQ量化前进行[Shift策略](https://arxiv.org/abs/2304.09145)，默认为False。使用Shift策略需要设`do_ptq`为True。
 - `shift_all_linear`: 是否对模型中所有Linear层应用Shift，如果为True，将会对非LayerNorm-Linear组合的Linear进行Shift，并且添加两个op，默认为False
@@ -85,6 +94,11 @@ python  finetune_generation.py ./llama/gptq_argument.json
 - `smooth_search_piece`: 使用分段搜索功能时，是否搜索分段数量，默认为False。设为True时，`smooth_k_piece`建议设为6，搜索分段数量耗时较长，如需加速Smooth过程建议关闭。
 - `do_gptq`: 是否进行GPTQ量化，GPTQ对模型进行WINT4量化，相比于普通PTQ量化精度更高，量化时间较长。默认为False。
 - `gptq_step`: GPTQ量化步数，也即模型前向次数，默认为8。
+- `do_awq`: 是否进行AWQ量化，AWQ对模型进行WINT4量化，相比于普通PTQ量化精度更高。默认为False。
+- `auto_clip`: AWQ时是否进行自动搜索截断值并对模型权重进行截断操作，截断操作有利于量化模型精度，但搜索速度较慢。默认为False。
+- `autoclip_step`: AutoClip步数，也即模型前向次数，采样时默认concat每轮数据用来搜索截断值，默认为8。
+
+
 </div>
 
 

llm/finetune_generation.py

@@ -556,7 +556,13 @@ def compute_metrics_do_generation(eval_preds):
             raise NotImplementedError(
                 "PTQ strategy not supported for LoRA model. Please merge lora parameters to pretrain model first."
             )
-        from quant import apply_ptq, apply_shift, apply_smooth, get_ptq_model_config
+        from quant import (
+            apply_autoclip,
+            apply_ptq,
+            apply_shift,
+            apply_smooth,
+            get_ptq_model_config,
+        )
 
         trainer.model.eval()
         trainer.model.config.quantization_config.quant_type = quant_args.quant_type
@@ -575,6 +581,9 @@ def compute_metrics_do_generation(eval_preds):
         if quant_args.smooth:
             apply_smooth(quant_args, trainer, ptq_dataloader, ptq_model_config)
 
+        if quant_args.auto_clip:
+            apply_autoclip(quant_args, trainer, ptq_dataloader)
+
         apply_ptq(quant_args, trainer, ptq_dataloader)
         trainer.save_model(merge_tensor_parallel=training_args.tensor_parallel_degree > 1)
 

llm/llama/awq_argument.json

@@ -0,0 +1,22 @@
+{
+    "model_name_or_path": "./checkpoints/llama_sft_ckpts",
+    "per_device_train_batch_size": 8,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "src_length": 1024,
+    "max_length": 2048,
+    "fp16": true,
+    "fp16_opt_level": "O2",
+    "dataset_name_or_path": "./data",
+    "output_dir": "./checkpoints/llama_ptq_ckpts",
+    "do_eval": true,
+    "eval_with_do_generation": false,
+    "do_ptq": true,
+    "quant_type": "weight_only_int4",
+    "weight_quant_method": "groupwise",
+    "ptq_step": 16,
+    "smooth": true,
+    "auto_clip": true,
+    "autoclip_step": 1,
+    "do_awq": true
+  }

llm/quant.py

@@ -23,6 +23,8 @@
 from paddle.quantization import PTQ, QAT, QuantConfig
 from paddleslim.quant.advanced import (
     GPTQ,
+    AutoClip,
+    AWQSearch,
     EMASampler,
     MultiStepSampler,
     PieceWiseSearch,
@@ -34,11 +36,16 @@
     QuantizedColumnParallelLinear,
     QuantizedRowParallelLinear,
 )
-from paddleslim.quant.observers import AbsMaxChannelWiseWeightObserver, AVGObserver
+from paddleslim.quant.observers import (
+    AbsMaxChannelWiseWeightObserver,
+    AVGObserver,
+    GroupWiseWeightObserver,
+)
 from paddleslim.quant.observers.abs_max_weight import (
     AbsMaxChannelWiseWeightObserverLayer,
 )
 from paddleslim.quant.observers.avg import AVGObserverLayer
+from paddleslim.quant.observers.groupwise import GroupWiseWeightObserverLayer
 
 from paddlenlp.peft import PrefixModelForCausalLM
 from paddlenlp.peft.lora import (
@@ -96,20 +103,23 @@ def apply_shift(quant_args, trainer, ptq_dataloader, ptq_model_config):
         sample_function=shift_sampler,
         shift_all_linears=quant_args.shift_all_linears,
     )
-
-    trainer.ptq_loop(
-        ptq_dataloader,
-        description="Shift",
-        max_eval_iters=quant_args.shift_step,
-    )
-    shift.update_weight()
+    with paddle.no_grad():
+        trainer.ptq_loop(
+            ptq_dataloader,
+            description="Shift",
+            max_eval_iters=quant_args.shift_step,
+        )
+        shift.update_weight()
     del shift, shift_sampler
     logger.info("***** Shift done *****")
 
 
 def apply_smooth(quant_args, trainer, ptq_dataloader, ptq_model_config):
 
-    logger.info("***** Running Smooth *****")
+    if quant_args.do_awq:
+        logger.info("***** Running AWQ *****")
+    else:
+        logger.info("***** Running Smooth *****")
     smooth_sampler = MultiStepSampler() if quant_args.smooth_sampler == "multi_step" else None
     if quant_args.smooth_piecewise_search:
         search_func = PieceWiseSearch(
@@ -123,6 +133,12 @@ def apply_smooth(quant_args, trainer, ptq_dataloader, ptq_model_config):
             weight_quant_method="abs_max_channel_wise",
             act_quant_method="avg",
         )
+    elif quant_args.do_awq:
+        search_func = AWQSearch(
+            n_grid=20,
+            bits_length=4,
+            weight_quant_method=quant_args.weight_quant_method,
+        )
     else:
         search_func = None
     smooth = Smooth(
@@ -132,31 +148,64 @@ def apply_smooth(quant_args, trainer, ptq_dataloader, ptq_model_config):
         smooth_all_linears=quant_args.smooth_all_linears,
         sample_function=smooth_sampler,
         search_function=search_func,
+        smooth_method="awq" if quant_args.do_awq else "smoothquant",
     )
-    trainer.ptq_loop(
-        ptq_dataloader,
-        description="Smooth",
-        max_eval_iters=quant_args.smooth_step,
-    )
+    with paddle.no_grad():
+        trainer.ptq_loop(
+            ptq_dataloader,
+            description="Smooth",
+            max_eval_iters=quant_args.smooth_step,
+        )
 
-    smooth.update_weight()
+        smooth.update_weight()
     del smooth, smooth_sampler, search_func
     logger.info("***** Smooth done *****")
 
 
+def apply_autoclip(quant_args, trainer, ptq_dataloader):
+    """
+    AutoClip
+    """
+    print("-------------------Start AutoClip------------------")
+    sampler = MultiStepSampler()
+    auto_clip = AutoClip(
+        trainer.model,
+        weight_bits=4,
+        weight_quant_method=quant_args.weight_quant_method,
+        sample_function=sampler,
+        n_grid=20,
+        max_shrink=0.5,
+    )
+    with paddle.no_grad():
+        trainer.ptq_loop(
+            ptq_dataloader,
+            description="AutoClip",
+            max_eval_iters=quant_args.autoclip_step,
+        )
+        auto_clip.auto_clip()
+    del sampler, auto_clip
+    logger.info("***** AutoClip done *****")
+
+
 def apply_ptq(quant_args, trainer, ptq_dataloader):
     logger.info("***** Running PTQ *****")
     q_config = QuantConfig(activation=None, weight=None)
+    if quant_args.weight_quant_method == "abs_max_channel_wise":
+        weight_observer = AbsMaxChannelWiseWeightObserver
+    elif quant_args.weight_quant_method == "groupwise":
+        weight_observer = GroupWiseWeightObserver
+    else:
+        raise ValueError("weight_quant_method should be one of ['abs_max_channel_wise', 'groupwise']")
 
     if quant_args.quant_type == "a8w8":
         activation = AVGObserver(quant_bits=8)
-        weight = AbsMaxChannelWiseWeightObserver(quant_bits=8)
+        weight = weight_observer(quant_bits=8)
     elif quant_args.quant_type == "weight_only_int4":
         activation = None
-        weight = AbsMaxChannelWiseWeightObserver(quant_bits=4)
+        weight = weight_observer(quant_bits=4)
     elif quant_args.quant_type == "weight_only_int8":
         activation = None
-        weight = AbsMaxChannelWiseWeightObserver(quant_bits=8)
+        weight = weight_observer(quant_bits=8)
     else:
         raise ValueError("quant_type should be one of ['a8w8', 'weight_only_int4', 'weight_only_int8']")
 
@@ -181,10 +230,12 @@ def apply_ptq(quant_args, trainer, ptq_dataloader):
         if isinstance(cur_layer, AbsMaxChannelWiseWeightObserverLayer):
             if "_observer" not in cur_name:
                 weight_scales[cur_name] = cur_layer.scales().numpy().tolist()
+        if isinstance(cur_layer, GroupWiseWeightObserverLayer):
+            if "_observer" not in cur_name:
+                weight_scales[cur_name] = cur_layer.scales().numpy().tolist()
         if isinstance(cur_layer, AVGObserverLayer):
             if "_observer" not in cur_name:
                 act_scales[cur_name] = cur_layer.scales().numpy().tolist()
-
     weight_scales_path = os.path.join(trainer.args.output_dir, "weight_scales.json")
     with open(weight_scales_path, "w") as f:
         json.dump(weight_scales, f)
@@ -210,12 +261,13 @@ def apply_gptq(quant_args, trainer, ptq_dataloader):
             parent_layer, sub_name = find_parent_layer_and_sub_name(model, cur_name)
             cur_quant_layer = GPTQ(cur_layer)
             setattr(parent_layer, sub_name, cur_quant_layer)
-            trainer.ptq_loop(
-                ptq_dataloader,
-                description="GPTQ",
-                max_eval_iters=quant_args.gptq_step,
-            )
-            cur_quant_layer.fasterquant(percdamp=0.1, groupsize=-1, actorder=True)
+            with paddle.no_grad():
+                trainer.ptq_loop(
+                    ptq_dataloader,
+                    description="GPTQ",
+                    max_eval_iters=quant_args.gptq_step,
+                )
+                cur_quant_layer.fasterquant(percdamp=0.1, groupsize=-1, actorder=True)
             del cur_quant_layer
             setattr(parent_layer, sub_name, cur_layer)
     logger.info("***** GPTQ done *****")