[llm]fix bug in ChatGLM merge tp params and prefix model generation (#6730)

* fix

* fix prefix generation

* fix chatglm tp

Author: lugimzzz

llm/README.md

@@ -167,18 +167,12 @@ python  -u  -m paddle.distributed.launch --gpus "0,1"  finetune_generation.py ./
 我们使用张量并行(TP，Tensor Parallelism)训练过程中，为了节省TP参数合并时间往往在中间checkpoint将参数存储为多个TP参数分片，可以使用提供的分片合并参数脚本进行参数合并。
 
 ```
-python merge_tp_params.py  \
-    --model_name_or_path ./checkpoints/chatglm_v2_sft_ckpts/checkpoint-7163 \
-    --merge_model_path ./checkpoints/chatglm_v2_sft_ckpts/checkpoint_merge \
-    --dtype "float16" \
-    --with_tokenizer
+python merge_tp_params.py \
+    --model_name_or_path ./checkpoints/chatglm_v2_sft_ckpts/checkpoint-100
 ```
 
 **参数：**
-- `model_name_or_path`: 必须，预训练模型名称或者本地的模型路径，用于热启模型和分词器，默认为None。
-- `merge_model_path`: 必须，合并参数后保存路径，默认为None。
-- `dtype`: 必须，模型参数dtype，默认为None。
-- `with_tokenizer`: 是否同时保存分词器，默认为False。
+- `model_name_or_path`: 必须，本地的TP模型参数路径，默认为None。
 - `device`: 运行环境，默认为gpu。
 
 ### 3.7 LoRA参数合并

llm/merge_tp_params.py

@@ -11,52 +11,90 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import importlib
+import os
+
 import paddle
-from paddle.distributed import fleet
 
-from paddlenlp.transformers import AutoModelForCausalLM, AutoTokenizer
+from paddlenlp.transformers import AutoConfig
+from paddlenlp.transformers.auto.modeling import MAPPING_NAMES
+from paddlenlp.utils.log import logger
 
 
 def parse_arguments():
     import argparse
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_name_or_path", default=None, required=True, help="The directory of model.")
-    parser.add_argument("--merge_model_path", default=None, required=True, help="The directory of merged model.")
     parser.add_argument("--device", type=str, default="gpu", help="Device")
-    parser.add_argument("--dtype", type=str, default=None, required=True, help="Model dtype")
-    parser.add_argument("--with_tokenizer", type=bool, default=True, help="Save tokenizer at the same time")
     return parser.parse_args()
 
 
-def merge():
+def load_tp_params(tp_degree, path):
+    tp_state_dict_list = []
+    for tp in range(tp_degree):
+        tp_state_dict = {}
+        tmp = paddle.load(os.path.join(path, f"model_state.tp{tp:0>2d}.pdparams"), return_numpy=True)
+        for k, v in tmp.items():
+            tp_state_dict[k] = v
+        tp_state_dict_list.append(tp_state_dict)
+
+    return tp_state_dict_list
+
+
+def merge_tensor_parallel(model_class, state_dict_list, config) -> None:
+    """the entry of converting config and converting model file
+
+    Args:
+        input_dir (str | None): the input dir which contains `pytorch_model.bin` and `config.json` file
+        config (PretrainedConfig): the PretrainedConfig instance of model
+    """
+    name_action_mappings = model_class._get_tensor_parallel_mappings(config, is_split=False)
+    state_keys_map = model_class._resolve_prefix_keys(name_action_mappings.keys(), state_dict_list[0].keys())
+
+    for k, v in state_keys_map.items():
+        name_action_mappings[v] = name_action_mappings.pop(k)
+
+    state_dict_to_save = {}
+    for key in state_dict_list[0].keys():
+        tensor = state_dict_list[0][key]
+        if key in name_action_mappings:
+            ret = [x[key] for x in state_dict_list]
+            action = name_action_mappings.pop(key)
+            tensor = action(ret)
+
+        state_dict_to_save[key] = tensor
+
+    if len(name_action_mappings) > 0:
+        for x in name_action_mappings.keys():
+            logger.warning(f"key <{x}> need to merge tensor parallel but we can't find in model state.")
+
+    logger.info("Finally, we merging state dict to fellowing tensors.")
+    for k, v in state_dict_to_save.items():
+        logger.info(f"{k}, {v.shape}, {v.dtype}")
+
+    return state_dict_to_save
+
+
+def main():
     args = parse_arguments()
     paddle.set_device(args.device)
-    tensor_parallel_degree = paddle.distributed.get_world_size()
-    tensor_parallel_rank = 0
-    if tensor_parallel_degree > 1:
-        strategy = fleet.DistributedStrategy()
-        strategy.hybrid_configs = {
-            "dp_degree": 1,
-            "mp_degree": tensor_parallel_degree,
-            "pp_degree": 1,
-            "sharding_degree": 1,
-        }
-        fleet.init(is_collective=True, strategy=strategy)
-        hcg = fleet.get_hybrid_communicate_group()
-        tensor_parallel_rank = hcg.get_model_parallel_rank()
-
-    model = AutoModelForCausalLM.from_pretrained(
-        args.model_name_or_path,
-        dtype=args.dtype,
-        tensor_parallel_degree=tensor_parallel_degree,
-        tensor_parallel_rank=tensor_parallel_rank,
-    )
-    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
-    if tensor_parallel_rank == 0:
-        model.save_pretrained(args.merge_model_path, merge_tensor_parallel=tensor_parallel_degree > 1)
-        tokenizer.save_pretrained(args.merge_model_path)
+    config = AutoConfig.from_pretrained(args.model_name_or_path)
+    init_class = config["architectures"][0]
+    import_class = importlib.import_module(f"paddlenlp.transformers.{MAPPING_NAMES[init_class[:-11]]}.modeling")
+    model_class = getattr(import_class, init_class)
+
+    if config.tensor_parallel_degree > 1:
+        tp_state_dict_list = load_tp_params(config.tensor_parallel_degree, args.model_name_or_path)
+        state_dict_to_save = merge_tensor_parallel(
+            model_class=model_class, state_dict_list=tp_state_dict_list, config=config
+        )
+
+        logger.info("Saving")
+        paddle.save(state_dict_to_save, os.path.join(args.model_name_or_path, "model_state.pdparams"))
+    else:
+        logger.info("No need to merge since config.tensor_parallel_degree <= 1.")
 
 
 if __name__ == "__main__":
-    merge()
+    main()

llm/opt/pt_argument.json

@@ -116,7 +116,6 @@ def __init__(self, args):
                 model=self.model,
                 prefix_path=self.args.prefix_path,
                 postprocess_past_key_value=prefix_tuning_params["postprocess_past_key_value"],
-                pad_attention_mask=prefix_tuning_params["pad_attention_mask"],
             )
         self.model.eval()
         self.tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, padding_side="left")

llm/predict_generation.py

@@ -167,7 +167,7 @@ def apply_gptq(quant_args, trainer, ptq_dataloader):
     for cur_name, cur_layer in model.named_sublayers():
         if type(cur_layer) in [paddle.nn.Linear, ColumnParallelLinear, RowParallelLinear]:
             num_layer += 1
-            logger.info("GPTQ layer", num_layer, cur_name)
+            logger.info(f"GPTQ layer: {num_layer}, {cur_name}")
             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)

llm/quant.py

@@ -121,15 +121,30 @@ def generate(self, **kwargs):
     def _prepare_inputs_for_generation(self, *args, **kwargs):
         model_kwargs = self.model_prepare_inputs_for_generation(*args, **kwargs)
         attention_mask = model_kwargs["attention_mask"]
+        batch_size = model_kwargs["input_ids"].shape[0]
         if self.pad_attention_mask is not None:
             attention_mask = self.pad_attention_mask(
                 model_kwargs["input_ids"].shape, self.prefix_config.num_prefix_tokens, attention_mask
             )
         else:
-            prefix_attention_mask = paddle.ones(
-                [model_kwargs["input_ids"].shape[0], self.prefix_config.num_prefix_tokens], dtype=attention_mask.dtype
-            )
-            attention_mask = paddle.concat((prefix_attention_mask, attention_mask), axis=1)
+            if len(attention_mask.shape) == 2:
+                prefix_attention_mask = paddle.ones(
+                    [batch_size, self.prefix_config.num_prefix_tokens], dtype=attention_mask.dtype
+                )
+            elif len(attention_mask.shape) == 3:
+                batch_size, src_seq_len, tgt_seq_len = attention_mask.shape
+                prefix_attention_mask = paddle.ones(
+                    [batch_size, src_seq_len, self.prefix_config.num_prefix_tokens], dtype=attention_mask.dtype
+                )
+            elif len(attention_mask.shape) == 4:
+                batch_size, num_heads, src_seq_len, tgt_seq_len = attention_mask.shape
+                prefix_attention_mask = paddle.ones(
+                    [batch_size, num_heads, src_seq_len, self.prefix_config.num_prefix_tokens],
+                    dtype=attention_mask.dtype,
+                )
+            else:
+                raise ValueError(f"Unexpected attention_mask shape: {attention_mask.shape}")
+            attention_mask = paddle.concat((prefix_attention_mask, attention_mask), axis=-1)
         model_kwargs["attention_mask"] = attention_mask
 
         if "past_key_values" in self.forward_keys:
@@ -139,7 +154,6 @@ def _prepare_inputs_for_generation(self, *args, **kwargs):
         else:
             raise NotImplementedError("Model does not support past_key_values either cache")
         if model_kwargs[key] is None:
-            batch_size = model_kwargs["input_ids"].shape[0]
             past_key_values = self._get_past_key_values(batch_size)
             model_kwargs[key] = past_key_values
         return model_kwargs

paddlenlp/peft/prefix/prefix_model.py

@@ -751,7 +751,7 @@ def forward(self, hidden_states):
 
 
 class ChatGLMForCausalLM(ChatGLMPretrainedModel):
-    _keys_to_ignore_on_save = [r"lm_head.weight"]
+    _keys_to_ignore_on_save = [r"lm_head.decoder_weight"]
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config: ChatGLMConfig):