[LLM Inference] add --use_fake_parameter option for ptq fake scales and fix compute error of total_max_length (#8955)

* update some code

* update

* update

* update

* update tune_cublaslt_gemm demo

* fix step in tune_cublaslt_gemm

Author: yuanlehome

csrc/generation/test_tune_cublaslt_gemm.py

@@ -12,14 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from paddlenlp_ops import tune_cublaslt_gemm
 import paddle
+from paddlenlp_ops import tune_cublaslt_gemm
+
+M_tensor = paddle.to_tensor([32768])
+
+# llama3.1-8b
+k1 = [4096, 4096, 4096, 14336]
+n1 = [6144, 4096, 28672, 4096]
+
+# llama3.1-405b mp=8
+k2 = [16384, 16384, 16384, 6656]
+n2 = [2560, 16384, 13312, 16384]
+
+# qwen2-1.5b
+k3 = [1536, 1536, 1536, 8960]
+n3 = [2048, 1536, 17920, 1536]
+
+# qwen2-7b
+k4 = [3584, 3584, 3584, 18944]
+n4 = [4608, 3584, 37888, 3584]
 
-M_tensor = paddle.to_tensor([1024])
-K_tensor = paddle.to_tensor([1024, 2048])
-N_tensor = paddle.to_tensor([4096, 8192])
+K_tensor = paddle.to_tensor(k1 + k2 + k3 + k4)
+N_tensor = paddle.to_tensor(n1 + n2 + n3 + n4)
 
 Dtype = "int8"
-Path = "./search.csv"
+Path = "./cublaslt_gemm_search.csv"
 
 tune_cublaslt_gemm(M_tensor, K_tensor, N_tensor, Dtype, True, False, Path)

csrc/generation/tune_cublaslt_gemm.cu

@@ -759,6 +759,9 @@ void TuneCublasltGemm(const paddle::Tensor& M,
       case 1024:
         step = 1024;
         break;
+      case 8192:
+        step = 4096;
+        break;
     }
   }
 

llm/predict/export_model.py

@@ -29,6 +29,15 @@ class ExportArgument:
     output_path: str = field(default=None, metadata={"help": "The output path of model."})
 
 
+def add_inference_args_to_config(model_config, args):
+    """Add export arguments to config."""
+    model_config.infer_model_block_size = args.block_size
+    model_config.infer_model_max_seq_len = args.total_max_length
+    model_config.infer_model_cachekv_int8_type = args.cachekv_int8_type
+    model_config.infer_model_dtype = args.dtype
+    model_config.infer_model_paddle_commit = paddle.version.commit
+
+
 def main():
     parser = PdArgumentParser((PredictorArgument, ModelArgument, ExportArgument))
     predictor_args, model_args, export_args = parser.parse_args_into_dataclasses()
@@ -60,6 +69,7 @@ def main():
             "cachekv_int8_type": predictor_args.cachekv_int8_type,
         },
     )
+    add_inference_args_to_config(predictor.model.config, predictor_args)
     predictor.model.config.save_pretrained(export_args.output_path)
     if predictor.generation_config is not None:
         predictor.generation_config.save_pretrained(export_args.output_path)

llm/predict/predictor.py

@@ -101,7 +101,7 @@ class PredictorArgument:
             "help": "If benchmark set as `True`, we will force model decode to max_length, which is helpful to compute throughput. "
         },
     )
-
+    use_fake_parameter: bool = field(default=False, metadata={"help": "use fake parameter, for ptq scales now."})
     block_attn: bool = field(default=False, metadata={"help": "whether use block attention"})
     block_size: int = field(default=64, metadata={"help": "the block size for cache_kvs."})
     cachekv_int8_type: str = field(
@@ -124,7 +124,7 @@ class PredictorArgument:
 
     @property
     def total_max_length(self):
-        return self.src_length + self.max_length
+        return 8192  # Maximum sequence length.
 
 
 @dataclass
@@ -948,8 +948,7 @@ def predict(self, input_texts: list[str], return_tokens=False):
         result_queue = mp.Queue()
         tensor_queue = mp.Queue()
 
-        output_tensor = paddle.full(shape=[MAX_BSZ + 2, 1], fill_value=2, dtype="int64")
-        output_tensor = output_tensor.cpu()
+        output_tensor = paddle.full(shape=[MAX_BSZ + 2, 1], fill_value=2, dtype="int64").cpu()
         tensor_queue.put(output_tensor)
 
         read_res_process = mp.Process(
@@ -1074,8 +1073,7 @@ def predict(self, input_texts: list[str], return_tokens=False):
         result_queue = mp.Queue()
         tensor_queue = mp.Queue()
 
-        output_tensor = paddle.full(shape=[MAX_BSZ + 2, 1], fill_value=2, dtype="int64")
-        output_tensor = output_tensor.cpu()
+        output_tensor = paddle.full(shape=[MAX_BSZ + 2, 1], fill_value=2, dtype="int64").cpu()
         tensor_queue.put(output_tensor)
 
         read_res_process = mp.Process(
@@ -1108,10 +1106,11 @@ def predict(self, input_texts: list[str], return_tokens=False):
 
 def get_ptq_multicards_num(directory):
     count = 0
-    prefix = "act_scales_"
-    for filename in os.listdir(directory):
-        if filename.startswith(prefix):
-            count += 1
+    if os.path.exists(directory):
+        prefix = "act_scales_"
+        for filename in os.listdir(directory):
+            if filename.startswith(prefix):
+                count += 1
     return count
 
 
@@ -1204,6 +1203,7 @@ def create_predictor(
             config.model_name_or_path = predictor_args.model_name_or_path
             config.quant_type = predictor_args.quant_type
             config.cachekv_int8_type = predictor_args.cachekv_int8_type
+            config.use_fake_parameter = predictor_args.use_fake_parameter
             config.single_card_ptq = True
             if predictor_args.avx_model:
                 config.avx_type = predictor_args.avx_type
@@ -1381,15 +1381,10 @@ def create_predictor(
 
         elif predictor_args.mode == "static":
             config = AutoConfig.from_pretrained(predictor_args.model_name_or_path)
-            config.quant_type = predictor_args.quant_type
-            config.cachekv_int8_type = predictor_args.cachekv_int8_type
 
             if config.quantization_config.quant_type is not None:
-                predictor_args.quant_type = config.quantization_config.quant_type
-                config.quant_type = config.quantization_config.quant_type
-                if "c8" in config.quant_type:
+                if "c8" in config.quantization_config.quant_type:
                     predictor_args.cachekv_int8_type = "static"
-                    config.cachekv_int8_type = "static"
 
             if "llama" in config.architectures[0].lower():
                 if predictor_args.block_attn:

paddlenlp/experimental/transformers/llama/modeling.py

@@ -43,7 +43,12 @@
     GenerationBlockInferenceModel,
     GenerationInferenceModel,
 )
-from paddlenlp.experimental.transformers.utils import infererence_model_from_pretrained
+from paddlenlp.experimental.transformers.utils import (
+    EmptyActScale,
+    EmptyCacheScale,
+    EmptyWeightScale,
+    infererence_model_from_pretrained,
+)
 from paddlenlp.transformers import LlamaConfig, LlamaPretrainedModel
 from paddlenlp.transformers.conversion_utils import split_param_func
 from paddlenlp.transformers.llama.modeling import LlamaLMHead
@@ -346,7 +351,7 @@ def __init__(self, config: LlamaConfig):
         self.num_layers = config.num_hidden_layers
         self.epsilon = config.rms_norm_eps
         self.max_position_embeddings = config.max_position_embeddings
-        self.quant_type = config.quant_type
+        self.quant_type = config.get("quant_type", "")
 
         self.rope_theta = config.rope_theta
         self.use_neox = True
@@ -364,6 +369,8 @@ def __init__(self, config: LlamaConfig):
             self.smooth = config.quantization_config.smooth
             self.shift_smooth_all_linears = config.quantization_config.shift_smooth_all_linears
 
+        self.use_fake_parameter = config.get("use_fake_parameter", False)
+
         if self.use_weight_only:
             assert (
                 self.quant_type == "weight_only_int8" or self.quant_type == "weight_only_int4"
@@ -894,6 +901,30 @@ def set_state_dict(self, state_dict):
 
             if "a8w8" in self.quant_type:
                 if self.shift_smooth_all_linears:
+                    if self.use_fake_parameter:
+                        if "llama.layers.{}.self_attn.o_proj.shift_bias".format(idx) not in state_dict:
+                            state_dict["llama.layers.{}.self_attn.o_proj.shift_bias".format(idx)] = paddle.zeros(
+                                shape=[
+                                    (self.num_attention_heads // self.config.tensor_parallel_degree)
+                                    * (self.hidden_size // self.num_attention_heads)
+                                ],
+                                dtype=paddle.get_default_dtype(),
+                            )
+                            state_dict["llama.layers.{}.self_attn.o_proj.smooth_weight".format(idx)] = paddle.ones(
+                                shape=[
+                                    (self.num_attention_heads // self.config.tensor_parallel_degree)
+                                    * (self.hidden_size // self.num_attention_heads)
+                                ],
+                                dtype=paddle.get_default_dtype(),
+                            )
+                            state_dict["llama.layers.{}.mlp.down_proj.shift_bias".format(idx)] = paddle.zeros(
+                                shape=[self.intermediate_size // self.config.tensor_parallel_degree],
+                                dtype=paddle.get_default_dtype(),
+                            )
+                            state_dict["llama.layers.{}.mlp.down_proj.smooth_weight".format(idx)] = paddle.ones(
+                                shape=[self.intermediate_size // self.config.tensor_parallel_degree],
+                                dtype=paddle.get_default_dtype(),
+                            )
                     self.transformer_block.linear_shifts[idx].set_value(
                         paddle.to_tensor(state_dict["llama.layers.{}.self_attn.o_proj.shift_bias".format(idx)])
                     )
@@ -908,6 +939,33 @@ def set_state_dict(self, state_dict):
                     )
 
                 if self.shift:
+                    if self.use_fake_parameter:
+                        if "llama.layers.{}.input_layernorm.bias".format(idx) not in state_dict:
+                            state_dict["llama.layers.{}.input_layernorm.bias".format(idx)] = paddle.zeros(
+                                shape=[self.hidden_size], dtype=paddle.get_default_dtype()
+                            )
+                            state_dict["llama.layers.{}.post_attention_layernorm.bias".format(idx)] = paddle.zeros(
+                                [self.hidden_size], dtype=paddle.get_default_dtype()
+                            )
+                            unfused_state_dict["self_attn.q_proj.bias"] = paddle.zeros(
+                                shape=[self.num_attention_heads * (self.hidden_size // self.num_attention_heads)],
+                                dtype=paddle.get_default_dtype(),
+                            )
+                            unfused_state_dict["self_attn.k_proj.bias"] = paddle.zeros(
+                                shape=[self.num_key_value_heads * (self.hidden_size // self.num_attention_heads)],
+                                dtype=paddle.get_default_dtype(),
+                            )
+                            unfused_state_dict["self_attn.v_proj.bias"] = paddle.zeros(
+                                shape=[self.num_key_value_heads * (self.hidden_size // self.num_attention_heads)],
+                                dtype=paddle.get_default_dtype(),
+                            )
+                            unfused_state_dict["mlp.gate_proj.bias"] = paddle.zeros(
+                                shape=[self.intermediate_size], dtype=paddle.get_default_dtype()
+                            )
+                            unfused_state_dict["mlp.up_proj.bias"] = paddle.zeros(
+                                shape=[self.intermediate_size], dtype=paddle.get_default_dtype()
+                            )
+
                     self.transformer_block.ln_biases[idx].set_value(
                         paddle.to_tensor(state_dict["llama.layers.{}.input_layernorm.bias".format(idx)])
                     )
@@ -948,6 +1006,14 @@ def set_state_dict(self, state_dict):
                     self.transformer_block.ffn1_biases[idx].set_value(paddle.to_tensor(concated_ffn1_bias))
 
                     if self.shift_smooth_all_linears:
+                        if self.use_fake_parameter:
+                            if "llama.layers.{}.self_attn.o_proj.bias".format(idx) not in state_dict:
+                                state_dict["llama.layers.{}.self_attn.o_proj.bias".format(idx)] = paddle.zeros(
+                                    [self.hidden_size], dtype=paddle.get_default_dtype()
+                                )
+                                state_dict["llama.layers.{}.mlp.down_proj.layer.bias".format(idx)] = paddle.zeros(
+                                    [self.hidden_size], dtype=paddle.get_default_dtype()
+                                )
                         self.transformer_block.linear_biases[idx].set_value(
                             paddle.to_tensor(state_dict["llama.layers.{}.self_attn.o_proj.bias".format(idx)])
                         )
@@ -981,41 +1047,64 @@ def set_state_dict(self, state_dict):
                 weight_scale_map_dict = scale_map_dict["weight_scale"]
                 cache_scale_map_dict = scale_map_dict["cachekv_scale"]
 
-                act_scale_json_path = os.path.join(self.quant_model_path, "act_scales.json")
-                weight_scale_json_path = os.path.join(self.quant_model_path, "weight_scales.json")
-                if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
-                    act_scale_json_path = os.path.join(
-                        self.quant_model_path, f"act_scales_{self.config.tensor_parallel_rank}.json"
+                if not self.use_fake_parameter:
+                    act_scale_json_path = os.path.join(self.quant_model_path, "act_scales.json")
+                    weight_scale_json_path = os.path.join(self.quant_model_path, "weight_scales.json")
+                    if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
+                        act_scale_json_path = os.path.join(
+                            self.quant_model_path, f"act_scales_{self.config.tensor_parallel_rank}.json"
+                        )
+                        weight_scale_json_path = os.path.join(
+                            self.quant_model_path, f"weight_scales_{self.config.tensor_parallel_rank}.json"
+                        )
+                    act_scale_loader = ActScalesLoader(
+                        act_scale_json_path, act_scale_map_dict, num_of_layers=self.config.num_hidden_layers
                     )
-                    weight_scale_json_path = os.path.join(
-                        self.quant_model_path, f"weight_scales_{self.config.tensor_parallel_rank}.json"
+                    weight_scales_loader = WeightScalesLoader(
+                        weight_scale_json_path,
+                        weight_scale_map_dict,
+                        num_of_layers=self.config.num_hidden_layers,
+                        concat_qkv=True,
+                        concat_ffn1=True,
+                    )
+                else:
+                    act_scale_loader = EmptyActScale(act_scale_map_dict, num_of_layers=self.config.num_hidden_layers)
+                    weight_scales_loader = EmptyWeightScale(
+                        weight_scale_map_dict,
+                        num_of_layers=self.config.num_hidden_layers,
+                        num_head=self.num_attention_heads,
+                        dim_head=self.hidden_size // self.num_attention_heads,
+                        ffn_hidden_size=self.intermediate_size,
+                        num_key_value_heads=self.num_key_value_heads,
+                        mp_size=self.config.tensor_parallel_degree,
                     )
-                act_scale_loader = ActScalesLoader(
-                    act_scale_json_path, act_scale_map_dict, num_of_layers=self.config.num_hidden_layers
-                )
                 self.transformer_block.act_scales = act_scale_loader.scale
 
-                weight_scales_loader = WeightScalesLoader(
-                    weight_scale_json_path,
-                    weight_scale_map_dict,
-                    num_of_layers=self.config.num_hidden_layers,
-                    concat_qkv=True,
-                    concat_ffn1=True,
-                )
-
                 if self.config.cachekv_int8_type == "static":
-                    cache_scale_json_path = os.path.join(self.quant_model_path, "cachekv_scales.json")
-                    if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
-                        cache_scale_json_path = os.path.join(
-                            self.quant_model_path, f"cachekv_scales_{self.config.tensor_parallel_rank}.json"
+                    if not self.use_fake_parameter:
+                        cache_scale_json_path = os.path.join(self.quant_model_path, "cachekv_scales.json")
+                        if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
+                            cache_scale_json_path = os.path.join(
+                                self.quant_model_path, f"cachekv_scales_{self.config.tensor_parallel_rank}.json"
+                            )
+                        cache_scales_loader = CacheScaleLoader(
+                            cache_scale_json_path,
+                            cache_scale_map_dict,
+                            num_of_layers=self.config.num_hidden_layers,
+                            num_heads=self.num_attention_heads // self.config.tensor_parallel_degree,
+                            num_key_value_heads=self.num_key_value_heads // self.config.tensor_parallel_degree,
                         )
-                    cache_scales_loader = CacheScaleLoader(
-                        cache_scale_json_path,
-                        cache_scale_map_dict,
-                        num_of_layers=self.config.num_hidden_layers,
-                        num_heads=self.num_attention_heads // self.config.tensor_parallel_degree,
-                        num_key_value_heads=self.num_key_value_heads // self.config.tensor_parallel_degree,
-                    )
+                    else:
+                        cache_scales_loader = EmptyCacheScale(
+                            cache_scale_map_dict,
+                            num_of_layers=self.config.num_hidden_layers,
+                            num_heads=self.num_attention_heads,
+                            dim_heads=self.hidden_size // self.num_attention_heads,
+                            is_channel_wise=False,
+                            num_key_value_heads=self.num_key_value_heads,
+                            mp_size=self.config.tensor_parallel_degree,
+                        )
+
                     for k, v in cache_scales_loader.scale.items():
                         for i_layer, weight_scale in enumerate(v):
                             weight_scale = weight_scale.astype("float32")