utils.py

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
from __future__ import annotations

import glob
import math
import os
import struct
from typing import Dict, Optional

import numpy as np
import paddle
import paddle.distributed as dist
import paddle.incubate.multiprocessing as mp
from paddle.distributed import fleet
from paddle.io import BatchSampler, DataLoader, DistributedBatchSampler
from paddlenlp_ops import get_output
from sklearn.metrics import accuracy_score

from paddlenlp.datasets import InTokensIterableDataset
from paddlenlp.trainer import Trainer, TrainerCallback
from paddlenlp.trainer.trainer_utils import IterableDatasetShard, has_length
from paddlenlp.transformers import (
 AutoTokenizer,
 ChatGLMv2Tokenizer,
 LlamaForCausalLMPipe,
 PretrainedConfig,
)
from paddlenlp.transformers.tokenizer_utils import PretrainedTokenizer
from paddlenlp.utils.log import logger


def compute_metrics(eval_preds):

 flattened_preds = np.array(eval_preds.predictions).flatten()
 flattened_labels = np.array(eval_preds.label_ids).flatten()
 filtered_preds = flattened_preds[flattened_labels != -100]
 filtered_labels = flattened_labels[flattened_labels != -100]
 accuracy = accuracy_score(y_true=filtered_labels, y_pred=filtered_preds)
 return {
 "accuracy": accuracy,
 }


def get_prefix_tuning_params(model):
 if model.base_model_prefix == "chatglm":
 from paddlenlp.peft.prefix import chatglm_postprocess_past_key_value

 num_attention_heads = model.config.num_attention_heads
 num_hidden_layers = model.config.num_hidden_layers
 hidden_size = model.config.hidden_size
 postprocess_past_key_value = chatglm_postprocess_past_key_value
 multi_query_group_num = None
 elif model.base_model_prefix == "chatglm_v2":
 from paddlenlp.peft.prefix import chatglm_postprocess_past_key_value

 num_attention_heads = model.config.num_attention_heads
 num_hidden_layers = model.config.num_layers
 hidden_size = model.config.hidden_size
 postprocess_past_key_value = chatglm_postprocess_past_key_value
 multi_query_group_num = model.config.multi_query_group_num
 elif model.base_model_prefix == "bloom":
 from paddlenlp.peft.prefix import bloom_postprocess_past_key_value

 num_attention_heads = model.config.num_attention_heads
 num_hidden_layers = model.config.n_layer
 hidden_size = model.config.n_embed
 postprocess_past_key_value = bloom_postprocess_past_key_value
 multi_query_group_num = None
 elif model.base_model_prefix == "llama":
 from paddlenlp.peft.prefix import llama_postprocess_past_key_value

 num_attention_heads = model.config.n_head
 num_hidden_layers = model.config.n_layer
 hidden_size = model.config.hidden_size
 postprocess_past_key_value = llama_postprocess_past_key_value
 multi_query_group_num = None
 elif model.base_model_prefix == "qwen":
 from paddlenlp.peft.prefix import qwen_postprocess_past_key_value

 num_attention_heads = model.config.num_attention_heads
 num_hidden_layers = model.config.num_hidden_layers
 hidden_size = model.config.hidden_size
 postprocess_past_key_value = qwen_postprocess_past_key_value
 multi_query_group_num = None
 else:
 raise ValueError(f"Unknown base_model_prefix: {model.base_model_prefix}. ")
 return dict(
 num_attention_heads=num_attention_heads,
 num_hidden_layers=num_hidden_layers,
 hidden_size=hidden_size,
 postprocess_past_key_value=postprocess_past_key_value,
 multi_query_group_num=multi_query_group_num,
 )


def get_lora_target_modules(model):
 # Not yet support RowParallelLinear
 if model.base_model_prefix == "chatglm":
 target_modules = [".*query_key_value.*", ".*dense.*", ".*dense_h_to_4h.*", ".*dense_4h_to_h.*"]
 elif model.base_model_prefix == "chatglm_v2":
 target_modules = [
 ".*query.*",
 ".*key.*",
 ".*value.*",
 ".*dense.*",
 ".*dense_h_to_4h.*",
 ".*dense_4h_to_h.*",
 ]
 elif model.base_model_prefix == "bloom":
 target_modules = [".*query_key_value.*", ".*dense.*", ".*dense_h_to_4h.*", ".*dense_4h_to_h.*"]
 elif model.base_model_prefix == "llama" or isinstance(model, LlamaForCausalLMPipe):
 target_modules = [
 ".*q_proj.*",
 ".*v_proj.*",
 ".*k_proj.*",
 ".*o_proj.*",
 ".*gate_proj.*",
 ".*down_proj.*",
 ".*up_proj.*",
 ]
 elif model.base_model_prefix == "opt":
 target_modules = [
 ".*project_in.*",
 ".*project_out.*",
 ".*q_proj.*",
 ".*k_proj.*",
 ".*v_proj.*",
 ".*qkv_proj.*",
 ".*out_proj.*",
 ".*linear1.*",
 ".*linear2.*",
 ]
 elif model.base_model_prefix == "qwen":
 target_modules = [
 ".*attn.c_attn.*",
 ".*attn.c_proj.*",
 ".*mlp.w1.*",
 ".*mlp.w2.*",
 ".*mlp.c_proj.*",
 ]
 else:
 raise ValueError(f"Unknown base_model_prefix: {model.base_model_prefix}.")
 return target_modules


class InTokensIterDatasetCallback(TrainerCallback):
 """
 A [`TrainerCallback`] that handles early stopping.

 """

 def on_step_end(self, args, state, control, **kwargs):
 train_dataloader = kwargs["train_dataloader"]
 if isinstance(train_dataloader.dataset, InTokensIterableDataset):
 dataset = train_dataloader.dataset
 elif isinstance(train_dataloader.dataset, IterableDatasetShard) and isinstance(
 train_dataloader.dataset.dataset, InTokensIterableDataset
 ):
 dataset = train_dataloader.dataset.dataset
 else:
 raise ValueError(
 "Unexpected dataset format: InTokensIterDatasetCallback expectes `paddlenlp.datasets.InTokensIterableDataset`"
 )
 if state.trial_params is None:
 state.trial_params = {}
 state.trial_params["intokens_global_step"] = dataset.intokens_global_step


class CausalLMTrainer(Trainer):
 def __init__(self, do_generation: bool, gen_args, data_args, **kwargs):
 super().__init__(**kwargs)
 self.do_generation = do_generation
 self.gen_args = gen_args
 self.data_args = data_args

 def prediction_step(
 self,
 model,
 inputs,
 prediction_loss_only: bool,
 ignore_keys=None,
 ):
 if prediction_loss_only or self.args.pipeline_parallel_degree > 1:
 return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)
 elif not self.do_generation:
 loss, logits, labels = super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)
 # argmax here to avoid gather all logits, which is too memory-consuming.
 # keepdim in order to maintain the same shape as logits
 if isinstance(logits, (list, tuple)):
 logits = logits[0]
 return (loss, logits.argmax(axis=-1, keepdim=True), labels)

 loss = None

 model.eval()
 with paddle.no_grad():
 generated_tokens = model.generate(
 input_ids=inputs["input_ids"],
 attention_mask=inputs["attention_mask"] if "attention_mask" in inputs else None,
 position_ids=inputs["position_ids"] if "position_ids" in inputs else None,
 max_length=max(self.data_args.max_length - inputs["input_ids"].shape[-1], 1),
 decode_strategy="sampling",
 top_k=self.gen_args.top_k,
 top_p=self.gen_args.top_p,
 bos_token_id=self.tokenizer.bos_token_id,
 eos_token_id=self.tokenizer.eos_token_id,
 pad_token_id=self.tokenizer.pad_token_id,
 use_cache=True,
 )[0]
 all_preds = []
 for pred_tokens in generated_tokens:
 pred_tokens = pred_tokens[pred_tokens != self.tokenizer.pad_token_id].tolist()
 all_preds.append(pred_tokens)
 max_pred_length = max([len(x) for x in all_preds])
 for index, preds in enumerate(all_preds):
 all_preds[index] = preds + [-100] * (max_pred_length - len(preds))
 all_preds = paddle.to_tensor(all_preds)

 if "labels" in inputs:
 all_labels = paddle.to_tensor(inputs["labels"])
 else:
 all_labels = None

 return (loss, all_preds, all_labels)

 def log(self, logs: Dict[str, float], **kwargs) -> None:
 if "loss" in logs:
 logs["ppl"] = np.exp(logs["loss"])
 if "eval_loss" in logs:
 logs["eval_ppl"] = np.exp(logs["eval_loss"])

 super(CausalLMTrainer, self).log(logs, **kwargs)

 def get_ptq_dataloader(self, ptq_ds):
 if self.args.world_size <= 1:
 ptq_sampler = BatchSampler(
 dataset=ptq_ds,
 shuffle=True,
 batch_size=self.args.per_device_train_batch_size,
 drop_last=self.args.dataloader_drop_last,
 )
 else:
 ptq_sampler = DistributedBatchSampler(
 self.train_dataset,
 batch_size=self.args.per_device_train_batch_size,
 shuffle=True,
 num_replicas=self.args.dataset_world_size,
 rank=self.args.dataset_rank,
 drop_last=self.args.dataloader_drop_last,
 )
 ptq_dataloader = DataLoader(
 ptq_ds,
 batch_sampler=ptq_sampler,
 collate_fn=self.data_collator,
 num_workers=self.args.dataloader_num_workers,
 )
 return ptq_dataloader

 def ptq_loop(
 self,
 dataloader: DataLoader,
 description: str,
 max_eval_iters: Optional[int] = -1,
 ):
 if isinstance(dataloader, paddle.io.DataLoader):
 batch_size = dataloader.batch_sampler.batch_size
 else:
 raise ValueError("Only support for paddle.io.DataLoader")

 if has_length(dataloader):
 logger.info(f" Num examples = {self.num_examples(dataloader)}")
 if max_eval_iters > 0:
 logger.info(f" Total {description} steps = {max_eval_iters}")
 else:
 logger.info(f" Total {description} steps = {len(dataloader)}")
 else:
 logger.info(" Num examples: Unknown")
 if max_eval_iters > 0:
 logger.info(f" Total {description} steps = {max_eval_iters}")

 logger.info(f" Pre device batch size = {batch_size}")
 logger.info(f" Total Batch size = {batch_size * self.args.dataset_world_size}")
 self.model.eval()
 with paddle.no_grad():
 for step, inputs in enumerate(dataloader):
 self.prediction_step(model=self.model, inputs=inputs, prediction_loss_only=True, ignore_keys=None)
 if max_eval_iters > 0 and step >= max_eval_iters - 1:
 break


def get_infer_model_path(input_dir, model_prefix):
 if dist.get_world_size() > 1:
 local_rank = dist.get_rank()
 return os.path.join(input_dir, "rank_{}".format(local_rank), model_prefix)
 else:
 return os.path.join(input_dir, model_prefix)


def generate_rank_mapping(output_filename):
 ring_id = -1
 try:
 hcg = fleet.get_hybrid_communicate_group()
 model_parallel_group = hcg.get_model_parallel_group()
 ring_id = model_parallel_group.id
 except Exception:
 pass

 if ring_id == -1:
 return

 world_size = dist.get_world_size()
 with open(output_filename, "w") as f:
 f.write("[ring_id -> ranks]\n")
 f.write(",".join(map(str, [0] + list(range(world_size)))) + "\n")
 f.write(",".join(map(str, [ring_id] + list(range(world_size)))) + "\n")

 f.write("[rank -> ring_ids]\n")
 for i in range(world_size):
 f.write("{},0,{}\n".format(i, ring_id))


def deserialize_from_file(fp):
 x_type = fp.read(1)
 x_type_out = struct.unpack("c", x_type)[0]
 # data
 data_list = []
 if x_type_out == b"0":
 data = fp.read(4)
 data_out = struct.unpack("f", data)[0]
 while data:
 data_out = struct.unpack("f", data)[0]
 data_list.append(data_out)
 data = fp.read(4)
 elif x_type_out == b"1":
 data = fp.read(8)
 while data:
 data_out = struct.unpack("l", data)[0]
 data_list.append(data_out)
 data = fp.read(8)
 elif x_type_out == b"2":
 data = fp.read(4)
 while data:
 data_out = struct.unpack("i", data)[0]
 data_list.append(data_out)
 data = fp.read(4)
 else:
 print("type error")
 data_arr = np.array(data_list)
 return data_arr


def get_alibi_slopes(num_heads):
 closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
 base = 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3)))
 powers = np.arange(1, 1 + closest_power_of_2)
 slopes = np.power(base, powers)

 if closest_power_of_2 != num_heads:
 extra_base = 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3)))
 num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
 extra_powers = np.arange(1, 1 + 2 * num_remaining_heads, 2)
 slopes = np.concatante([slopes, np.power(extra_base, extra_powers)], axis=0)

 return slopes.astype("float32")


def pad_batch_data(insts, pad_id=0, return_seq_len=False, pad_style="right"):
 """Pad sequences to the max sequence length in batch."""
 max_len = max(map(len, insts))
 if pad_style == "left":
 inst_data = np.array([[pad_id] * (max_len - len(inst)) + list(inst) for inst in insts])
 else:
 inst_data = np.array([list(inst) + [pad_id] * (max_len - len(inst)) for inst in insts])

 if return_seq_len:
 seq_len = np.array([len(inst) for inst in insts])
 return inst_data.astype("int64").reshape([-1, max_len]), seq_len
 else:
 return inst_data.astype("int64").reshape([-1, max_len])


def dybatch_preprocess(
 tokenizer,
 texts: list[str],
 src_length: int,
 max_length: int,
 architectures: str,
 top_p: float,
 temperature: float,
 eos_token_id: int | list[list[int]],
 pre_caches_length: int = 0,
 benchmark: bool = False,
):
 """Pre-process generation inputs."""
 inputs = {}
 if "chatglmforcausallm" == architectures.lower():
 input_ids = []
 position_ids = []

 for text in texts:
 tokens = tokenizer(
 text,
 return_tensors="np",
 padding=True,
 max_length=src_length,
 # if use chat_template, it will not add special_tokens
 add_special_tokens=tokenizer.chat_template is None or isinstance(tokenizer, ChatGLMv2Tokenizer),
 )
 input_ids.append(tokens["input_ids"][0])
 position_ids.append(tokens["position_ids"][0])

 pad_token_id = tokenizer([tokenizer.pad_token], return_tensors="np")["input_ids"][0][0]
 inputs["input_ids"], seq_len = pad_batch_data(input_ids, pad_id=pad_token_id, return_seq_len=True)
 bs = inputs["input_ids"].shape[0]
 max_len = max(map(len, input_ids))

 inst_data_pos = []
 for i in range(len(position_ids)):
 inst_data_pos.append(np.array([list(inst) + [0] * (max_len - len(inst)) for inst in position_ids[i]]))
 inputs["position_ids"] = paddle.to_tensor(np.array(inst_data_pos))
 elif "gpt" in architectures:
 input_ids = []
 if isinstance(texts, str):
 texts = [texts]

 for text in texts:
 tokens = tokenizer(
 text,
 return_tensors="np",
 padding=False,
 max_length=src_length,
 return_attention_mask=False,
 return_token_type_ids=False,
 )
 input_ids.append(tokens["input_ids"][0])

 pad_token_id = tokenizer([tokenizer.pad_token], return_tensors="np")["input_ids"][0][-1]
 inputs["input_ids"], seq_len = pad_batch_data(input_ids, pad_id=pad_token_id, return_seq_len=True)
 bs = inputs["input_ids"].shape[0]
 max_len = max(map(len, input_ids))

 position_ids = paddle.arange(sum(seq_len), dtype="int64")
 pre_len = seq_len[0]
 for length in seq_len[1:]:
 position_ids[pre_len : length + pre_len] = position_ids[pre_len : length + pre_len] - pre_len
 pre_len += length
 inputs["position_ids"] = position_ids
 else:
 input_ids = []
 if isinstance(texts, str):
 texts = [texts]

 for text in texts:
 tokens = tokenizer(
 text,
 return_tensors="np",
 padding=False,
 max_length=src_length,
 return_attention_mask=False,
 return_token_type_ids=False,
 )
 input_ids.append(tokens["input_ids"][0])

 pad_token_id = tokenizer([tokenizer.pad_token], return_tensors="np")["input_ids"][0][-1]
 inputs["input_ids"], seq_len = pad_batch_data(input_ids, pad_id=pad_token_id, return_seq_len=True)
 bs = inputs["input_ids"].shape[0]
 max_len = max(map(len, input_ids))

 position_ids = paddle.zeros(shape=[bs, max_length + src_length], dtype="int64")

 for i in range(bs):
 position_ids[i, pre_caches_length : pre_caches_length + seq_len[i]] = paddle.arange(seq_len[i])
 inputs["position_ids"] = position_ids

 tgt_ids = [input[-1:] for input in input_ids]
 tgt_pos = []
 for i, valid_len in enumerate(map(len, input_ids)):
 tgt_pos.append(valid_len - 1)

 step_idx = [
 0,
 ] * bs
 tgt_pos = np.array(tgt_pos).astype("int64")

 if isinstance(eos_token_id, int):
 eos_token_id = [eos_token_id]

 inputs["eos_token_id"] = np.array(eos_token_id * bs).reshape(-1, 1).astype("int64")

 inputs["top_p"] = (
 np.array(
 [
 top_p,
 ]
 * bs
 )
 .reshape(-1, 1)
 .astype("float32")
 )
 inputs["temperature"] = (
 np.array(
 [
 temperature,
 ]
 * bs
 )
 .reshape(-1, 1)
 .astype("float32")
 )
 inputs["seq_len_encoder"] = seq_len.astype("int32").reshape(-1, 1)
 inputs["seq_len_decoder"] = (seq_len + pre_caches_length).astype("int32").reshape(-1, 1)
 inputs["step_idx"] = np.array(step_idx).astype("int64").reshape(-1, 1)
 inputs["tgt_ids"] = np.array(tgt_ids).astype("int64").reshape(-1, 1)
 inputs["tgt_pos"] = tgt_pos.reshape(-1, 1)
 inputs["max_length"] = np.array(max_length - pre_caches_length).astype("int64").reshape((-1, 1))
 inputs["min_length"] = (
 np.array(
 [
 1
 if not benchmark
 else max_length
 - pre_caches_length, # Note(Zhengzekang): When in benchmark mode, we need to set a fixed decode length.
 ]
 * bs
 )
 .astype("int64")
 .reshape((-1, 1))
 )
 inputs["penalty_score"] = (
 np.array(
 [
 1.0,
 ]
 * bs
 )
 .astype("float32")
 .reshape((-1, 1))
 )
 inputs["frequency_score"] = (
 np.array(
 [
 0.0,
 ]
 * bs
 )
 .astype("float32")
 .reshape((-1, 1))
 )
 inputs["presence_score"] = (
 np.array(
 [
 0.0,
 ]
 * bs
 )
 .astype("float32")
 .reshape((-1, 1))
 )
 inputs["stop_flags"] = (
 np.array(
 [
 0,
 ]
 * bs
 )
 .astype("bool")
 .reshape((-1, 1))
 )
 inputs["stop_nums"] = np.array([bs]).astype("int64")
 return inputs


def load_real_time_tokens():
 tokens = []
 files = glob.glob(os.path.join("./real_time_save.*"))
 for j in range(1, len(files) + 1):
 filename = "./real_time_save.temp_ids_rank_0_step_{}".format(j)
 if not os.path.exists(filename):
 break
 fp = open(filename, "rb+")
 fp.read(1)
 data_list = deserialize_from_file(fp)
 fp.close()
 tokens.append(np.array(data_list).reshape(-1, 1))
 os.system("rm -f ./real_time_save.temp_ids_rank_*")
 tokens = np.concatenate(tokens, axis=1)
 return tokens


def init_chat_template(
 tokenizer: PretrainedTokenizer, model_name_or_path: str, chat_template_file: Optional[str] = None
):
 """init chat template for the given tokenizer.

 If is None, it will not use `chat_template.json`;
 If is equal with `model_name_or_path`, it will use the default loading;
 If is directory, it will find the `chat_template.json` under the directory;
 If is file, it will load it.

 Args:
 tokenizer (PretrainedTokenizer): the instance of tokenizer
 model_name_or_path (str): _description_
 chat_template_file (Optional[str], optional): _description_. Defaults to None.
 """
 # 1. use the default chat_template file
 if chat_template_file is None:
 return

 if str(chat_template_file).lower() == "none":
 # delete the chat_template from tokenizer if not use chat_template.
 # why do this: it will load the `chat_template.json` file by default
 tokenizer.chat_template = None
 return

 # it will load the `chat_template.json` file by default, so do nothing
 if chat_template_file == model_name_or_path:
 if tokenizer.chat_template is None:
 logger.warning(f"there is not `chat_template.json` file in the `{model_name_or_path}`")
 return

 if os.path.isdir(chat_template_file):
 local_chat_template_file_path = os.path.join(chat_template_file, "chat_template.json")
 if os.path.exists(local_chat_template_file_path):
 chat_template_file = local_chat_template_file_path
 else:
 logger.warning(f"there is not `chat_template.json` file in the `{model_name_or_path}`")
 return

 if not os.path.exists(chat_template_file):
 logger.warning(f"there is not `chat_template.json` file from path<`{model_name_or_path}`>")
 return

 logger.info(f"loading `chat_template.json` from `{chat_template_file}`")
 tokenizer.init_chat_template(chat_template_file)


def get_model_max_position_embeddings(config: PretrainedConfig) -> Optional[int]:
 names = [
 "max_position_embeddings", # most of models
 "max_sequence_length", # GLM model
 "seq_length", # llama model
 ]
 for name in names:
 max_length = config.get(name, None)
 if max_length is not None:
 return max_length
 return None


def get_default_max_decoding_length(config: PretrainedConfig, default: int = 1024) -> int:
 """get the default max decoding length from config.

 Args:
 config (PretrainedConfig): the instance of PretrainedConfig
 default (int): the default value of max decoding length

 Returns:
 int: the default max_length of decoding length
 """
 max_position_embeddings = get_model_max_position_embeddings(config)
 if max_position_embeddings is None:
 return default
 return max_position_embeddings // 4


def get_default_max_encoding_length(config: PretrainedConfig, default: int = 1024) -> int:
 """get the default max encoding length from config.

 Args:
 config (PretrainedConfig): the instance of PretrainedConfig
 default (int): the default value of max encoding length

 Returns:
 int: the default max_length of encoding length
 """

 max_position_embeddings = get_model_max_position_embeddings(config)
 if max_position_embeddings is None:
 return default
 return max_position_embeddings // 4 * 3


def read_res(model_name_or_path: str, tensor_queue: mp.Queue, result_queue: mp.Queue):
 tokenizer = AutoTokenizer.from_pretrained(
 model_name_or_path,
 )

 paddle.device.set_device("cpu")
 outputs = []
 output_tensor = tensor_queue.get(timeout=1)

 logger.info("Start read result message")
 logger.info(f"Current path is {os.getcwd()}")
 while True:
 get_output(output_tensor, 0, True)
 if output_tensor[0, 0] == -2: # read none
 continue
 bsz = output_tensor[1, 0].numpy()
 output_numpy = output_tensor[2 : bsz + 2].numpy()
 output_numpy[output_numpy == -1] = 2
 outputs.append(output_numpy)
 if output_tensor[0, 0] == -1:
 break
 output = np.concatenate(outputs, axis=1).tolist()
 seqs = tokenizer.batch_decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=False)
 for i, seq in enumerate(seqs):
 result_queue.put([i, seq])

 logger.info("Finish read result message")