[PaddleNLP 3.0] Refactor RLHF (#8617)

* refactor PPO & RM

llm/.gitignore

@@ -3,7 +3,6 @@ infer.json
 output.json 
 
 # data
-data/
 AdvertiseGen.tar.gz
 
 # checkpoints

llm/Alignment/PPO/comm_utils.py

@@ -0,0 +1,403 @@
+# 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.
+
+import sys
+
+import paddle
+import paddle.distributed as dist
+
+from paddlenlp.trainer.plugins.unified_checkpoint import flatten_list
+from paddlenlp.trainer.trainer import Trainer, logger
+from paddlenlp.trainer.utils.helper import nested_broadcast_tensor_with_empty
+from paddlenlp.utils.distributed import distributed_gather
+
+global_dev_id = 0 if paddle.get_device() == "cpu" else int(paddle.get_device().split(":")[1])
+
+
+def offload_tensor_to_cpu(tensors):
+    if isinstance(tensors, dict):
+        for _, v in tensors.items():
+            offload_tensor_to_cpu(v)
+    elif isinstance(tensors, paddle.Tensor):
+        if tensors.place.is_gpu_place():
+            cpu_tensor = tensors._copy_to(paddle.CUDAPinnedPlace(), False)
+            tensors.value().get_tensor()._share_data_with(cpu_tensor.value().get_tensor())
+    else:
+        logger.warning(f"Can't parse for type {type(tensors)}")
+        return tensors
+
+
+def reload_tensor_to_gpu(tensors):
+    if isinstance(tensors, dict):
+        for _, v in tensors.items():
+            reload_tensor_to_gpu(v)
+    elif isinstance(tensors, paddle.Tensor):
+        if tensors._is_initialized() and not tensors.place.is_gpu_place():
+            gpu_tensor = tensors._copy_to(paddle.CUDAPlace(global_dev_id), False)
+            tensors.value().get_tensor()._share_data_with(gpu_tensor.value().get_tensor())
+    else:
+        logger.warning(f"Can't parse for type {type(tensors)}")
+        return tensors
+
+
+def cleanup_tensor_space(tensors):
+    if isinstance(tensors, dict):
+        for _, v in tensors.items():
+            cleanup_tensor_space(v)
+    elif isinstance(tensors, paddle.Tensor):
+        tensors._clear_data()
+    else:
+        logger.warning(f"Can't parse for type {type(tensors)}")
+        return tensors
+
+
+def data_group_split(tensors, group):
+    if group is None:
+        return tensors
+    if isinstance(tensors, (list, tuple)):
+        return type(tensors)(data_group_split(t, group) for t in tensors)
+    elif isinstance(tensors, dict):
+        new_dict = {}
+        for k, v in tensors.items():
+            new_dict[k] = data_group_split(v, group)
+        return new_dict
+    elif isinstance(tensors, paddle.Tensor):
+        return tensors.split(group.nranks)[group.rank]
+    else:
+        logger.warning(f"Can't parse for type {type(tensors)}")
+        return tensors
+
+
+def data_group_merge(tensors, group):
+    if group is None:
+        return tensors
+
+    if isinstance(tensors, (list, tuple)):
+        return type(tensors)(data_group_merge(t, group) for t in tensors)
+    elif isinstance(tensors, dict):
+        new_dict = {}
+        for k, v in tensors.items():
+            new_dict[k] = data_group_merge(v, group)
+        return new_dict
+    elif isinstance(tensors, paddle.Tensor):
+        tensor_list = []
+        all_gather_nd(tensor_list, tensors, group=group, padded=True)
+        return paddle.concat(tensor_list)
+    else:
+        logger.warning(f"Can't parse for type {type(tensors)}")
+        return tensors
+
+
+def group_rank_guard(group, rank=0):
+    def decorator(func):
+        def wrapper_func(*args, **kwargs):
+            if group.rank == rank:
+                ret = func(*args, **kwargs)
+                dist.barrier()
+            else:
+                ret = None
+                dist.barrier()
+            ret = nested_broadcast_tensor_with_empty(ret, group=group)
+            return ret
+
+        return wrapper_func
+
+    return decorator
+
+
+def repad_rl_batches(batches, input_lengths):
+    if batches.get("position_ids", None) is not None:
+        v = batches["position_ids"]
+        for x in range(v.shape[0]):
+            v[x, input_lengths[x] :] = 1
+        batches["position_ids"] = v
+    for key in list(batches.keys()):
+        if batches[key].shape[0] != input_lengths.shape[0]:
+            batches[key] = batches[key].mean()
+
+    return batches
+
+
+# https://stackoverflow.com/questions/12594148/skipping-execution-of-with-block
+class SkipWithBlock(Exception):
+    pass
+
+
+class SkipContextManager:
+    def __init__(self, skip):
+        self.skip = skip
+
+    def __enter__(self):
+        if self.skip:
+            sys.settrace(lambda *args, **keys: None)
+            frame = sys._getframe(1)
+            frame.f_trace = self.trace
+
+    def trace(self, frame, event, arg):
+        raise SkipWithBlock()
+
+    def __exit__(self, type, value, traceback):
+        if type is None:
+            return  # No exception
+        if issubclass(type, SkipWithBlock):
+            return True  # Suppress special SkipWithBlock exception
+
+
+def all_gather_nd(tensor_list, tensor, group=None, padded=False):
+    """
+    Gathers tensor arrays of different lengths in a list.
+    The length dimension is 0. This supports any number of extra dimensions in the tensors.
+    All the other dimensions should be equal between the tensors.
+
+    Args:
+        tensor (Tensor): Tensor to be broadcast from current process.
+
+    Returns:
+        (Tensor): output list of tensors that can be of different sizes
+    """
+    if len(tensor.shape) == 0:
+        tensor = tensor.reshape([1])
+        dist.all_gather(tensor_list, tensor, group=group)
+        return tensor_list
+
+    world_size = group.nranks
+    local_size = paddle.to_tensor(tensor.shape, place=tensor.place)
+    all_sizes = [paddle.zeros_like(local_size) for _ in range(world_size)]
+    dist.all_gather(all_sizes, local_size, group=group)
+
+    # max_length = max(size[0] for size in all_sizes)
+
+    # length_diff = max_length.item() - local_size[0].item()
+    # if length_diff:
+    #     pad_size = (length_diff, *tensor.size()[1:])
+    #     padding = paddle.zeros(pad_size, place=tensor.place(), dtype=tensor.dtype)
+    #     tensor = padle.concat((tensor, padding))
+
+    max_length = max(size[-1] for size in all_sizes)
+
+    length_diff = max_length.item() - local_size[-1].item()
+    if length_diff:
+        pad_size = (*tensor.shape[:-1], length_diff)
+        padding = paddle.zeros(pad_size, dtype=tensor.dtype)
+        tensor = paddle.concat([tensor, padding], axis=-1)
+
+    all_tensors_padded = []
+    dist.all_gather(all_tensors_padded, tensor, group=group)
+    # all_tensors = []
+    if padded:
+        tensor_list.extend(all_tensors_padded)
+        return all_tensors_padded
+
+    for tensor_, size in zip(all_tensors_padded, all_sizes):
+        tensor_list.append(tensor_[..., : size[-1]])
+    return tensor_list
+
+
+def export_evaluate_model(self: Trainer, train_model, eval_model, **kwargs):
+    if eval_model is None:
+        return None
+
+    with_offload = kwargs.pop("with_offload", False)
+    train_tp_size = max(train_model.config.tensor_parallel_degree, 1)
+    eval_tp_size = max(eval_model.config.tensor_parallel_degree, 1)
+    eval_tp_rank = max(eval_model.config.tensor_parallel_rank, 0)
+
+    hcg = dist.fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    pp_group = hcg.get_pipe_parallel_group()
+    sd_group = hcg.get_sharding_parallel_group()
+    dp_group = hcg.get_data_parallel_group()
+
+    global_rank = paddle.distributed.get_rank()
+
+    train_state_dict = train_model.state_dict()
+    eval_state_dict = eval_model.state_dict()
+
+    if dp_group.rank <= 0 and sd_group.rank <= 0:
+        train_pp_size = pp_group.nranks
+        if eval_tp_size > 1 and train_tp_size != eval_tp_size:
+            raise ValueError("Only support for the same tensor_parallel_degree for train and eval model for now.")
+
+        # 单卡情况
+        # tp->single
+        # tp+pp -> single
+        if eval_tp_size == 1:
+            if train_pp_size == 1 and train_tp_size > 1:
+                # tp ->single
+                logger.error("using tp to single eval model.")
+                # state = train_model.merge_tensor_parallel()
+                tp_actions = train_model.get_tensor_parallel_convert_actions(
+                    train_model.config,
+                    loaded_state_dict_keys=eval_state_dict.keys(),
+                    is_split=False,
+                    ignore_error=False,
+                )
+
+                is_dst = global_rank == 0
+                for key in eval_state_dict.keys():
+                    tensor = train_state_dict[key]
+                    if key in tp_actions:
+                        ret = distributed_gather(tensor, dst=0, group=tp_group, offload=False)
+                        action = tp_actions.pop(key)
+                        tensor = action(ret) if is_dst else None
+                    else:
+                        tensor = tensor._copy_to(paddle.CPUPlace(), False) if is_dst else None
+
+                    if tensor is not None:
+                        eval_state_dict[key].set_value(tensor)
+
+                    if not eval_state_dict[key]._is_initialized():
+                        v = eval_state_dict[key]
+                        t = paddle._C_ops.full_like(v, 0, v.dtype, paddle.CUDAPlace(global_dev_id))
+                        v.get_tensor()._share_data_with(t.get_tensor())
+
+                    if with_offload:
+                        offload_tensor_to_cpu(train_state_dict[key])
+            else:
+                # single to single
+                # tp+pp -> single
+                raise ValueError("Not support yet.")
+
+        def create_send_recv_table(train_keys, eval_keys):
+            recv_table = []
+            send_table = []
+            if pp_group.rank == 0:
+                for key in eval_keys:
+                    recv_table.append((key, global_rank))
+
+            for key in train_keys:
+                send_table.append((key, global_rank))
+
+            all_recv, all_send = [], []
+            paddle.distributed.all_gather_object(all_recv, [recv_table], group=pp_group)
+            paddle.distributed.all_gather_object(all_send, [send_table], group=pp_group)
+            all_recv = flatten_list(all_recv)
+            all_send = flatten_list(all_send)
+
+            send_dict = {}
+            for k, v in all_send:
+                send_dict[k] = v
+
+            table = []
+            for k, v in all_recv:
+                # key, send, recv
+                table.append([k, send_dict.pop(k), v])
+            assert len(send_dict) == 0, f"Some key can't be recv {send_dict.keys()}"
+            return table
+
+            # pp0tp0 -> pp0tp0
+            # pp0tp1 -> pp0tp1
+            # pp1tp0 -> pp0tp0
+            # pp1tp1 -> pp0tp1
+
+        # tp情况
+        # tp+pp->tp
+        self.timers and self.timers("export-merge-pp").start()
+        if eval_tp_size > 1 and train_pp_size > 1:
+            table = create_send_recv_table(train_state_dict.keys(), eval_state_dict.keys())
+
+            for key, src_rank, dst_rank in table:
+                # Init tensor for model is cleaned
+                if not eval_state_dict[key]._is_initialized():
+                    v = eval_state_dict[key]
+                    t = paddle._C_ops.full_like(v, 0, v.dtype, paddle.CUDAPlace(global_dev_id))
+                    v.get_tensor()._share_data_with(t.get_tensor())
+
+                if src_rank == dst_rank and global_rank == src_rank:
+                    eval_state_dict[key].copy_(train_state_dict[key], True)
+                else:
+                    if global_rank == src_rank:
+                        dist.stream.send(train_state_dict[key], dst=dst_rank)
+
+                    if global_rank == dst_rank:
+                        dist.stream.recv(eval_state_dict[key], src=src_rank)
+
+                # Offload train model if need
+                if global_rank == src_rank and with_offload:
+                    offload_tensor_to_cpu(train_state_dict[key])
+
+        self.timers and self.timers("export-merge-pp").stop()
+        self.timers and self.timers("export-broadcast-pp").start()
+        if pp_group.nranks > 1:
+            paddle.distributed.parallel.sync_params_buffers(
+                eval_model, comm_group=pp_group, src_rank=pp_group.ranks[0], fuse_params=False
+            )
+        self.timers and self.timers("export-broadcast-pp").stop()
+    else:
+        # 其他 DP rank 的state dict, 适配 offload 和初始化
+        self.timers and self.timers("export-offload-and-init").start()
+        if with_offload:
+            for key in list(train_state_dict.keys()):
+                offload_tensor_to_cpu(train_state_dict[key])
+        for k, v in eval_state_dict.items():
+            if not v._is_initialized():
+                t = paddle._C_ops.full_like(v, 0, v.dtype, paddle.CUDAPlace(global_dev_id))
+                v.get_tensor()._share_data_with(t.get_tensor())
+        self.timers and self.timers("export-offload-and-init").stop()
+
+    paddle.distributed.barrier()
+    self.timers and self.timers("export-broadcast-sd-dp").start()
+    if eval_tp_size == 1:
+        for _, tensor in eval_state_dict.items():
+            paddle.distributed.broadcast(tensor, src=0, group=None, sync_op=True)
+    else:
+        if sd_group.nranks > 1:
+            if dp_group.rank <= 0:
+                paddle.distributed.parallel.sync_params_buffers(
+                    eval_model, comm_group=sd_group, src_rank=sd_group.ranks[0], fuse_params=False
+                )
+        if dp_group.nranks > 1:
+            paddle.distributed.parallel.sync_params_buffers(
+                eval_model, comm_group=dp_group, src_rank=dp_group.ranks[0], fuse_params=False
+            )
+    self.timers and self.timers("export-broadcast-sd-dp").stop()
+    # paddle.save(eval_state_dict, f"./tmp/eval_{sd_group.rank}_tp_{eval_tp_rank}_pp_{pp_group.rank}.pdparams")
+    # paddle.save(train_state_dict, f"./tmp/train_{sd_group.rank}_tp_{tp_group.rank}_pp_{pp_group.rank}.pdparams")
+    # paddle.distributed.barrier()
+    # exit(-1)
+
+    old_dp_workers = self.args.world_size // (max(sd_group.nranks, 1) * max(dp_group.nranks, 1))
+    group_nums = self.args.logical_process_index // old_dp_workers * eval_tp_size + eval_tp_rank
+
+    if not hasattr(self, "_policy_model_eval_group") or self._policy_model_eval_group is None:
+        self._policy_model_eval_group = create_data_trans_group(global_rank, group_nums)
+
+    return None
+
+
+def create_data_trans_group(global_rank, group_nums):
+    all_split_table = []
+    paddle.distributed.all_gather_object(all_split_table, [(global_rank, group_nums)])
+    all_split_table = flatten_list(all_split_table)
+    split_dict = {}
+    for k, v in all_split_table:
+        split_dict[k] = v
+
+    split_ranks = {}
+    for k, v in all_split_table:
+        if v in split_ranks:
+            split_ranks[v].append(k)
+        else:
+            split_ranks[v] = [k]
+
+    group = None
+    for k, ranks in split_ranks.items():
+        gp = paddle.distributed.new_group(ranks=ranks)
+        if global_rank in ranks:
+            group = gp
+
+    return group
+
+
+Trainer.export_evaluate_model = export_evaluate_model

llm/Alignment/PPO/data/__init__.py

@@ -0,0 +1,33 @@
+# 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 paddle.io import Dataset
+
+from .alpaca import *
+from .base import *
+from .preference import *
+from .prompt_only import *
+from .safe_rlhf import *
+from .supervised import *
+
+
+class DummyDataset(Dataset):
+    def __init__(self, length: int) -> None:
+        self.length = length
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, index: int):
+        return {}

llm/Alignment/PPO/data/alpaca.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Stanford Alpaca dataset for supervised instruction fine-tuning."""
+
+from __future__ import annotations
+
+from datasets import load_dataset
+
+from .base import RawDataset, RawSample
+
+__all__ = ["AlpacaDataset"]
+
+
+class AlpacaDataset(RawDataset):
+    NAME: str = "alpaca"
+    ALIASES: tuple[str, ...] = ("stanford-alpaca",)
+
+    def __init__(self, path: str | None = None) -> None:
+        self.data = load_dataset(path or "tatsu-lab/alpaca", split="train")
+
+    def __getitem__(self, index: int) -> RawSample:
+        data = self.data[index]
+        input = (  # pylint: disable=redefined-builtin
+            " ".join((data["instruction"], data["input"])) if data["input"] else data["instruction"]
+        )
+        answer = data["output"]
+        return RawSample(input=input, answer=answer)
+
+    def __len__(self) -> int:
+        return len(self.data)

llm/Alignment/PPO/data/base.py

@@ -0,0 +1,520 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Base dataset class."""
+
+from __future__ import annotations
+
+import abc
+import bisect
+import copy
+import os
+import warnings
+from fractions import Fraction
+from typing import Any, Callable, ClassVar, Collection, Iterable, Iterator, List
+from weakref import WeakValueDictionary
+
+import numpy as np
+import paddle
+from paddle.io import Dataset, IterableDataset, Subset
+from paddle.io.dataloader.collate import default_collate_fn
+from tqdm import tqdm
+from typing_extensions import NotRequired  # Python 3.11+
+from typing_extensions import TypedDict  # Python 3.10+
+
+from paddlenlp.transformers.tokenizer_utils import (
+    PaddingStrategy,
+    PretrainedTokenizerBase,
+    TruncationStrategy,
+)
+
+# from safe_rlhf.utils import is_main_process
+
+
+__all__ = [
+    "RawDataset",
+    "RawSample",
+    "TokenizedDataset",
+    "CollatorBase",
+    "parse_dataset",
+]
+
+IGNORE_INDEX: int = -100
+PROMPT_BEGIN: str = "BEGINNING OF CONVERSATION: "
+PROMPT_USER: str = "USER: {input} "
+PROMPT_ASSISTANT: str = "ASSISTANT:"  # should not have a space at the end
+PROMPT_INPUT: str = PROMPT_BEGIN + PROMPT_USER + PROMPT_ASSISTANT
+
+
+def format_prompt(
+    input: str | list[str],  # pylint: disable=redefined-builtin
+    eos_token: str,
+) -> str:
+    if isinstance(input, str):
+        input = [input]
+    elif not isinstance(input, list):
+        raise ValueError(f"Unsupported type of `input`: {type(input)}. Expected: str or list[str].")
+
+    if len(input) % 2 != 1:
+        raise ValueError(
+            "The length of `input` must be odd, while `input` must end at the user question.",
+        )
+
+    buffer = [PROMPT_BEGIN]
+    for i, line in enumerate(input):
+        if i % 2 == 0:
+            # User input
+            buffer.extend((PROMPT_USER.format(input=line), PROMPT_ASSISTANT))
+        else:
+            # Assistant response
+            buffer.extend((line, eos_token))
+
+    return "".join(buffer)
+
+
+def left_padding(sequences, padding_value=0):
+    arrs = [np.asarray(seq) for seq in sequences]
+    max_length = max([len(seq) for seq in sequences])
+    bs = len(sequences)
+    data = np.full([bs, max_length], padding_value, dtype=arrs[0].dtype)
+    for i, arr in enumerate(arrs):
+        data[i, -len(arr) :] = arr
+    return data
+
+
+def right_padding(sequences, padding_value=0):
+    arrs = [np.asarray(seq) for seq in sequences]
+    max_length = max([len(seq) for seq in sequences])
+    bs = len(sequences)
+    data = np.full([bs, max_length], padding_value, dtype=arrs[0].dtype)
+    for i, arr in enumerate(arrs):
+        data[i, : len(arr)] = arr
+    return data
+
+
+class ConcatDataset(Dataset):
+    r"""Dataset as a concatenation of multiple datasets.
+
+    This class is useful to assemble different existing datasets.
+
+    Args:
+        datasets (sequence): List of datasets to be concatenated
+    """
+    datasets: List[Dataset]
+    cumulative_sizes: List[int]
+
+    @staticmethod
+    def cumsum(sequence):
+        r, s = [], 0
+        for e in sequence:
+            l = len(e)
+            r.append(l + s)
+            s += l
+        return r
+
+    def __init__(self, datasets: Iterable[Dataset]) -> None:
+        super(ConcatDataset, self).__init__()
+        self.datasets = list(datasets)
+        assert len(self.datasets) > 0, "datasets should not be an empty iterable"  # type: ignore[arg-type]
+        for d in self.datasets:
+            assert not isinstance(d, IterableDataset), "ConcatDataset does not support IterableDataset"
+        self.cumulative_sizes = self.cumsum(self.datasets)
+
+    def __len__(self):
+        return self.cumulative_sizes[-1]
+
+    def __getitem__(self, idx):
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError("absolute value of index should not exceed dataset length")
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        return self.datasets[dataset_idx][sample_idx]
+
+    @property
+    def cummulative_sizes(self):
+        warnings.warn(
+            "cummulative_sizes attribute is renamed to " "cumulative_sizes", DeprecationWarning, stacklevel=2
+        )
+        return self.cumulative_sizes
+
+
+def parse_dataset(string: str) -> tuple[str, dict[str, Any]]:
+    """Parse dataset path and its proportion and optionally additional arguments from a string.
+
+    Args:
+        string (str): Dataset string in the format of ``dataset_name[:proportion[:dataset_path]]``.
+    """
+    if string.count(":") > 2:
+        raise ValueError(
+            f"Invalid dataset string `{string}`, "
+            "should be in the format of `dataset_name[:proportion[:dataset_path]]`.",
+        )
+    name, colon, proportion = string.partition(":")
+    if not colon:
+        return name, {"proportion": 1.0}
+    proportion, colon, path = proportion.partition(":")
+    if "/" in proportion:
+        left, right = proportion.split("/")
+        left, right = left.strip(), right.strip()
+        if right == "-1":
+            proportion = Fraction(int(left), -1)
+            if proportion == 0:
+                proportion = 0.0
+        else:
+            proportion = float(left) / float(right)
+    elif proportion != "":
+        proportion = float(proportion)
+    else:
+        proportion = 1.0
+    if not colon:
+        return name, {"proportion": proportion}
+    if not path:
+        raise ValueError(f"Invalid dataset path `{path}`.")
+    path = os.path.expanduser(path)
+    return name, {"proportion": proportion, "path": path}
+
+
+class RawSample(TypedDict, total=False):
+    """Raw sample type.
+
+    For SupervisedDataset, should provide (input, answer) or (dialogue).
+    For PreferenceDataset, should provide (input, answer, other_answer, better).
+    For SafetyPreferenceDataset, should provide (input, answer, other_answer, safer, is_safe, is_other_safe).
+    For PromptOnlyDataset, should provide (input).
+
+    When input is a list, it would be processed as a dialogue.
+    """
+
+    # Texts
+    input: NotRequired[str | list[str]]  # either `input` or `dialogue` should be provided
+    """User input text."""
+    answer: NotRequired[str]
+    """Assistant answer text."""
+    other_answer: NotRequired[str]
+    """Other assistant answer text via resampling."""
+    dialogue: NotRequired[list[str]]  # either `input` or `dialogue` should be provided
+    """Dialogue history."""
+
+    # Flags
+    better: NotRequired[bool]
+    """Whether ``answer`` is better than ``other_answer``."""
+    safer: NotRequired[bool]
+    """Whether ``answer`` is safer than ``other_answer``."""
+    is_safe: NotRequired[bool]
+    """Whether ``answer`` is safe."""
+    is_other_safe: NotRequired[bool]
+    """Whether ``other_answer`` is safe."""
+
+
+class RawDataset(Dataset):
+    """Dataset that provides raw text samples."""
+
+    NAME: ClassVar[str]
+    """Name of the dataset."""
+    ALIASES: ClassVar[Collection[str]]
+    """Name aliases of the dataset."""
+    __REGISTRY: ClassVar[WeakValueDictionary[str, RawDataset]] = WeakValueDictionary()
+    """Registry of all subclasses."""
+    __ALIAS_NAME_MAPPING: ClassVar[dict[str, str]] = {}
+    """Mapping from aliases to names."""
+
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+        name = getattr(cls, "NAME", None)
+        if name is None:
+            return
+
+        if not isinstance(name, str):
+            raise ValueError("`NAME` should be a string.")
+        if not name:
+            raise ValueError("`NAME` should not be empty.")
+        if name in cls.__REGISTRY or name in cls.__ALIAS_NAME_MAPPING:
+            raise ValueError(f"Duplicate dataset name `{name}`.")
+        cls.__REGISTRY[name] = cls
+
+        aliases = set(cls.__dict__.get("ALIASES", ()))  # do not inherit aliases from parent class
+        for alias in aliases:
+            if not isinstance(alias, str):
+                raise ValueError("`ALIASES` should be a set of strings.")
+            if not alias:
+                raise ValueError("`ALIASES` should not contain empty strings.")
+            if alias in cls.__REGISTRY or alias in cls.__ALIAS_NAME_MAPPING:
+                raise ValueError(f"Duplicate dataset alias `{alias}`.")
+        cls.__ALIAS_NAME_MAPPING.update(dict.fromkeys(aliases, name))
+
+    @staticmethod
+    def load(name: str, /, *args: Any, **kwargs: Any) -> RawDataset:
+        """Load a raw dataset by name."""
+        normalized_name = RawDataset.__ALIAS_NAME_MAPPING.get(name, name)
+        try:
+            cls = RawDataset.__REGISTRY[normalized_name]
+        except KeyError as ex:
+            raise ValueError(
+                f"Unknown dataset name `{name}`. "
+                "You should implement a raw dataset with that name and import it in your code. "
+                f'Available datasets are: {", ".join(sorted(RawDataset.__REGISTRY.keys()))}.',
+            ) from ex
+        return cls(*args, **kwargs)
+
+    make = load  # alias
+
+    @abc.abstractmethod
+    def __getitem__(self, index: int) -> RawSample:
+        """Get a data sample by index."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Get the number of samples in the dataset."""
+        raise NotImplementedError
+
+    def __iter__(self) -> Iterator[RawSample]:
+        """Iterate over the dataset."""
+        return (self[index] for index in range(len(self)))
+
+    def split_train_test(
+        self,
+        split_ratio: float = 0.2,
+        shuffle: bool = True,
+        seed: int = 42,
+    ) -> tuple[Dataset[RawSample], Dataset[RawSample]]:
+        """Split the dataset into train dataset and test dataset by split ratio."""
+        num_samples = len(self)
+        indices = list(range(num_samples))
+        num_test_samples = round(num_samples * split_ratio)
+        if shuffle:
+            rng = np.random.default_rng(seed)
+            rng.shuffle(indices)
+
+        train_dataset = Subset(self, indices[:-num_test_samples])
+        test_dataset = Subset(self, indices[-num_test_samples:])
+        return train_dataset, test_dataset
+
+
+class TokenizedDataset(Dataset):
+    """Dataset that provides tokenized samples."""
+
+    raw_datasets: list[RawDataset]
+    """Raw datasets."""
+    tokenizer: PretrainedTokenizerBase
+    """Tokenizer."""
+    rawdata: list[RawSample]
+    """Merged raw data samples."""
+    data: list[dict[str, paddle.Tensor]]
+    """Tokenized data samples."""
+
+    _SENTINEL: Any = object()
+
+    def __init__(  # pylint: disable=too-many-branches
+        self,
+        dataset_names_and_attributes: dict[str, float | dict[str, Any]] | Iterable[tuple[str, float | dict[str, Any]]],
+        tokenizer: PretrainedTokenizerBase,
+        lazy_tokenization: bool = True,
+        seed: int = 42,
+    ) -> None:
+        if not isinstance(dataset_names_and_attributes, dict):
+            dataset_names_and_attributes = tuple(dataset_names_and_attributes)
+            dataset_names = [name for name, _ in dataset_names_and_attributes]
+            if len(dataset_names) != len(set(dataset_names)):
+                raise ValueError(
+                    f"Dataset names should be unique, but got {dataset_names}.",
+                )
+
+        super().__init__()
+        self.dataset_names_and_proportion: dict[str, float | Fraction] = {}
+        self.raw_datasets = []
+        for name, attributes in dict(dataset_names_and_attributes).items():
+            if isinstance(attributes, float):
+                kwargs = {"proportion": attributes}
+            elif isinstance(attributes, dict):
+                kwargs = dict(attributes)  # copy
+            else:
+                raise TypeError(
+                    f"Dataset `{name}` attributes should be a float or a dict, " f"got {type(attributes).__name__}.",
+                )
+            proportion = kwargs.pop("proportion", 1.0)
+            if isinstance(proportion, Fraction):
+                if not (proportion < 0 and proportion.denominator == 1):
+                    raise ValueError(
+                        f"Dataset `{name}` proportion should be a negative integer "
+                        f"represents `num_samples / -1`, got {proportion}.",
+                    )
+            else:
+                proportion = float(proportion)
+                if proportion < 0.0:
+                    raise ValueError(
+                        f"Dataset `{name}` proportion should be no less than 0.0, " f"got {proportion}.",
+                    )
+            if proportion == 0.0:
+                continue
+            self.dataset_names_and_proportion[name] = proportion
+            self.raw_datasets.append(RawDataset.load(name, **kwargs))
+
+        self.tokenizer = tokenizer
+        self.seed = seed
+
+        merged_rawdata = self._merge_raw_datasets(seed=seed)
+        self.rawdata = [merged_rawdata[i] for i in range(len(merged_rawdata))]
+        if lazy_tokenization:
+            self.data = [self._SENTINEL for _ in range(len(self.rawdata))]
+        else:
+            self.data = list(
+                map(
+                    self.preprocess,
+                    tqdm(
+                        self.rawdata,
+                        desc="Preprocessing raw dataset...",
+                        disable=False,  # not is_main_process(),
+                    ),
+                ),
+            )
+
+    def __getitem__(self, index: int) -> dict[str, paddle.Tensor]:
+        """Get a tokenized data sample by index."""
+        data = self.data[index]
+        if data is self._SENTINEL:
+            raw_sample = self.rawdata[index]
+            data = self.preprocess(raw_sample)
+            self.data[index] = data
+        return data
+
+    def __len__(self) -> int:
+        """Get the number of samples in the dataset."""
+        return len(self.data)
+
+    @abc.abstractmethod
+    def preprocess(self, raw_sample: RawSample) -> dict[str, paddle.Tensor]:
+        """Pre-process a raw sample into a tokenized sample."""
+        raise NotImplementedError
+
+    def _merge_raw_datasets(self, seed: int | None = None) -> Dataset[RawSample]:
+        """Merge multiple raw datasets into one dataset."""
+        if seed is None:
+            seed = self.seed
+
+        num_datasets = len(self.raw_datasets)
+        datasets = []
+        for raw_dataset, seed_seq in zip(
+            self.raw_datasets,
+            np.random.SeedSequence(seed).spawn(num_datasets),
+        ):
+            num_raw_samples = len(raw_dataset)
+            proportion = self.dataset_names_and_proportion[raw_dataset.NAME]
+            if isinstance(proportion, Fraction):
+                num_selected_samples = abs(int(proportion))
+            else:
+                num_selected_samples = round(num_raw_samples * proportion)
+
+            indices = []
+            while num_selected_samples >= num_raw_samples > 0:
+                indices.extend(range(num_raw_samples))
+                num_selected_samples -= num_raw_samples
+            if num_selected_samples > 0:
+                rng = np.random.default_rng(seed_seq)
+                indices.extend(
+                    rng.choice(
+                        len(raw_dataset),
+                        size=num_selected_samples,
+                        replace=False,
+                    ).tolist(),
+                )
+            datasets.append(Subset(raw_dataset, indices))
+
+        if num_datasets == 1:
+            return datasets[0]
+        return ConcatDataset(datasets)
+
+    def tokenize(
+        self,
+        text: str,
+        add_special_tokens: bool = True,
+        padding: bool | str | PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        truncation: bool | str | TruncationStrategy = TruncationStrategy.LONGEST_FIRST,
+        max_length: int | None = None,
+    ) -> paddle.Tensor:  # size = (L,)
+        """Tokenize a text string into a tensor representation."""
+        if max_length is None:
+            max_length = self.tokenizer.model_max_length
+
+        return self.tokenizer(
+            text,
+            add_special_tokens=add_special_tokens,
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            return_tensors="np",
+        )["input_ids"][0]
+
+    def get_collator(self) -> Callable[[list[dict[str, paddle.Tensor]]], dict[str, paddle.Tensor]]:
+        """Get a collator function for the dataset."""
+        return default_collate_fn
+
+    def resize(self, size: int) -> None:
+        """Resize the dataset to the given size."""
+        if size < 0:
+            raise ValueError(f"Invalid dataset size: {size}")
+        old_size = len(self)
+        if size == old_size:
+            return
+        if size > old_size:
+            num_replicates = (size + old_size - 1) // old_size
+            self.data = self.data * num_replicates
+            self.rawdata = self.rawdata * num_replicates
+        self.data = self.data[:size]
+        self.rawdata = self.rawdata[:size]
+
+    def split_train_test(
+        self,
+        split_ratio: float = 0.2,
+        shuffle: bool = True,
+        seed: int | None = None,
+    ) -> tuple[Dataset[dict[str, paddle.Tensor]], Dataset[dict[str, paddle.Tensor]]]:
+        """Split the dataset into train dataset and test dataset by split ratio."""
+        if seed is None:
+            seed = self.seed
+
+        num_samples = len(self)
+        indices = list(range(num_samples))
+        num_test_samples = round(num_samples * split_ratio)
+        if shuffle:
+            rng = np.random.default_rng(seed)
+            rng.shuffle(indices)
+
+        memo = {id(attr): attr for attr in self.__dict__.values()}
+        train_dataset = copy.deepcopy(self, memo=memo.copy())
+        test_dataset = copy.deepcopy(self, memo=memo.copy())
+        train_dataset.data = [self.data[i] for i in indices[:-num_test_samples]]
+        train_dataset.rawdata = [self.rawdata[i] for i in indices[:-num_test_samples]]
+        test_dataset.data = [self.data[i] for i in indices[-num_test_samples:]]
+        test_dataset.rawdata = [self.rawdata[i] for i in indices[-num_test_samples:]]
+
+        return train_dataset, test_dataset
+
+
+class CollatorBase(metaclass=abc.ABCMeta):
+    pad_token_id: int  # The id of the padding token for the tokenizer.
+
+    def __init__(self, pad_token_id: int) -> None:
+        """Initialize a collator."""
+        self.pad_token_id = pad_token_id
+
+    @abc.abstractmethod
+    def __call__(self, samples: list[dict[str, paddle.Tensor]]) -> dict[str, paddle.Tensor]:
+        """Collate a list of samples into a batch."""
+        raise NotImplementedError

llm/Alignment/PPO/data/preference.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Dataset class for preference training."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+import paddle
+from typing_extensions import TypedDict  # Python 3.10+
+
+from .base import (
+    CollatorBase,
+    RawSample,
+    TokenizedDataset,
+    format_prompt,
+    right_padding,
+)
+
+__all__ = [
+    "PreferenceDataset",
+    "PreferenceCollator",
+    "PreferenceSample",
+    "PreferenceBatch",
+]
+
+
+class PreferenceSample(TypedDict, total=True):
+    better_input_ids: paddle.Tensor  # size = (L,)
+    worse_input_ids: paddle.Tensor  # size = (L,)
+
+
+class PreferenceBatch(TypedDict, total=True):
+    better_input_ids: paddle.Tensor  # size = (B, L)
+    better_attention_mask: paddle.Tensor  # size = (B, L)
+
+    worse_input_ids: paddle.Tensor  # size = (B, L)
+    worse_attention_mask: paddle.Tensor  # size = (B, L)
+
+
+class PreferenceDataset(TokenizedDataset):
+    def preprocess(self, raw_sample: RawSample, index=0) -> PreferenceSample:
+        prompt = format_prompt(input=raw_sample["input"], eos_token=self.tokenizer.eos_token)
+        better_answer = raw_sample["answer"]
+        worse_answer = raw_sample["other_answer"]
+        better = raw_sample["better"]
+        if not better:
+            better_answer, worse_answer = worse_answer, better_answer
+
+        better_input_ids = self.tokenize(prompt + better_answer + self.tokenizer.eos_token)
+        worse_input_ids = self.tokenize(prompt + worse_answer + self.tokenizer.eos_token)
+
+        if (
+            better_input_ids.shape == worse_input_ids.shape
+            and np.all(np.equal(better_input_ids, worse_input_ids)).item()
+        ):
+            # raise ValueError(
+            import warnings
+
+            warnings.warn(
+                "Two responses get the same `input_ids` after tokenization.\n\n"
+                f"Prompt: {prompt}\n\n"
+                f"Better answer: {better_answer}\n\n"
+                f"Worse answer: {worse_answer}",
+            )
+        return {
+            "better_input_ids": better_input_ids,  # size = (L,)
+            "worse_input_ids": worse_input_ids,  # size = (L,)
+        }
+
+    def get_collator(self) -> Callable[[list[dict[str, paddle.Tensor]]], dict[str, paddle.Tensor]]:
+        return PreferenceCollator(self.tokenizer.pad_token_id)
+
+
+class PreferenceCollator(CollatorBase):
+    def __call__(self, samples: list[PreferenceSample]) -> PreferenceBatch:
+        input_ids = [sample["better_input_ids"] for sample in samples] + [
+            sample["worse_input_ids"] for sample in samples
+        ]  # size = (2 * B, L)
+        attention_mask = [np.ones(input_id.shape, dtype=bool) for input_id in input_ids]  # size = (2 * B, L)
+
+        input_ids = right_padding(input_ids, padding_value=self.pad_token_id)  # size = (2 * B, L)
+        attention_mask = right_padding(attention_mask, padding_value=0)  # size = (2 * B, L)
+
+        bs = len(samples)
+        (better_input_ids, worse_input_ids,) = (  # size = (B, L)  # size = (B, L)
+            input_ids[:bs],
+            input_ids[bs:],
+        )
+        (better_attention_mask, worse_attention_mask,) = (  # size = (B, L)  # size = (B, L)
+            attention_mask[:bs],
+            attention_mask[bs:],
+        )
+
+        return {
+            "better_input_ids": better_input_ids,  # size = (B, L)
+            "better_attention_mask": better_attention_mask,  # size = (B, L)
+            "worse_input_ids": worse_input_ids,  # size = (B, L)
+            "worse_attention_mask": worse_attention_mask,  # size = (B, L)
+        }

llm/Alignment/PPO/data/prompt_only.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Dataset class for preference training."""
+
+from __future__ import annotations
+
+from typing import Callable, Hashable
+
+import numpy as np
+import paddle
+from paddle.io import Dataset, Subset
+from typing_extensions import TypedDict  # Python 3.10+
+
+from .base import CollatorBase, RawSample, TokenizedDataset, format_prompt, left_padding
+
+__all__ = [
+    "PromptOnlyDataset",
+    "PromptOnlyCollator",
+    "PromptOnlySample",
+    "PromptOnlyBatch",
+]
+
+
+class PromptOnlySample(TypedDict, total=True):
+    input_ids: paddle.Tensor  # size = (L,)
+
+
+class PromptOnlyBatch(TypedDict, total=True):
+    input_ids: paddle.Tensor  # size = (B, L)
+    attention_mask: paddle.Tensor  # size = (B, L)
+
+
+class PromptOnlyDataset(TokenizedDataset):
+    def preprocess(self, raw_sample: RawSample) -> PromptOnlySample:
+        prompt = format_prompt(input=raw_sample["input"], eos_token=self.tokenizer.eos_token)
+        input_ids = self.tokenize(prompt)
+        return {
+            "input_ids": input_ids,  # size = (L,)
+        }
+
+    def get_collator(self) -> Callable[[list[dict[str, paddle.Tensor]]], dict[str, paddle.Tensor]]:
+        return PromptOnlyCollator(self.tokenizer.pad_token_id)
+
+    def _merge_raw_datasets(self, seed: int | None = None) -> Dataset[RawSample]:
+        """Merge multiple raw datasets into one dataset and remove duplicates."""
+
+        def to_hashable(raw_sample: RawSample) -> Hashable:
+            input = raw_sample["input"]  # pylint: disable=redefined-builtin
+            return input if isinstance(input, str) else tuple(input)
+
+        merged = super()._merge_raw_datasets(seed)
+        inputs = {to_hashable(merged[i]): i for i in range(len(merged))}
+        return Subset(merged, sorted(inputs.values()))
+
+
+class PromptOnlyCollator(CollatorBase):
+    def __call__(self, samples: list[PromptOnlySample]) -> PromptOnlyBatch:
+        input_ids = [sample["input_ids"] for sample in samples]
+        attention_mask = [np.ones(input_id.shape, dtype=bool) for input_id in input_ids]
+
+        input_ids = left_padding(input_ids, padding_value=self.pad_token_id)
+        attention_mask = left_padding(attention_mask, padding_value=0)
+        return {
+            "input_ids": input_ids,  # size = (B, L)
+            "attention_mask": attention_mask,  # size = (B, L)
+        }

llm/Alignment/PPO/data/safe_rlhf.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Safe-RLHF preference datasets."""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from datasets import load_dataset
+
+from .base import RawDataset, RawSample
+
+__all__ = [
+    "SafeRLHFDataset",
+    "SafeRLHFTrainDataset",
+    "SafeRLHFTestDataset",
+    "SafeRLHF10KTrainDataset",
+]
+
+
+class SafeRLHFDataset(RawDataset):
+    SPLIT: ClassVar[str]
+    PATH: ClassVar[str]
+
+    def __init__(self, path: str | None = None) -> None:
+        self.data = load_dataset(path or self.PATH, split=self.SPLIT)
+
+    def __getitem__(self, index: int) -> RawSample:
+        data = self.data[index]
+        return RawSample(
+            input=data["prompt"],
+            answer=data["response_0"],
+            other_answer=data["response_1"],
+            better=int(data["better_response_id"]) == 0,
+            safer=int(data["safer_response_id"]) == 0,
+            is_safe=bool(data["is_response_0_safe"]),
+            is_other_safe=bool(data["is_response_1_safe"]),
+        )
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+
+class SafeRLHFTrainDataset(SafeRLHFDataset):
+    NAME: str = "PKU-SafeRLHF/train"
+    ALIASES: tuple[str, ...] = ("PKU-Alignment/PKU-SafeRLHF/train",)
+    PATH: str = "PKU-Alignment/PKU-SafeRLHF"
+    SPLIT: str = "train"
+
+
+class SafeRLHFTestDataset(SafeRLHFDataset):
+    NAME: str = "PKU-SafeRLHF/test"
+    ALIASES: tuple[str, ...] = ("PKU-Alignment/PKU-SafeRLHF/test",)
+    PATH: str = "PKU-Alignment/PKU-SafeRLHF"
+    SPLIT: str = "test"
+
+
+class SafeRLHF30KTrainDataset(SafeRLHFDataset):
+    NAME: str = "PKU-SafeRLHF-30K/train"
+    ALIASES: tuple[str, ...] = ("PKU-Alignment/PKU-SafeRLHF-30K/train",)
+    PATH: str = "PKU-Alignment/PKU-SafeRLHF-30K"
+    SPLIT: str = "train"
+
+
+class SafeRLHF30KTestDataset(SafeRLHFDataset):
+    NAME: str = "PKU-SafeRLHF-30K/test"
+    ALIASES: tuple[str, ...] = ("PKU-Alignment/PKU-SafeRLHF-30K/test",)
+    PATH: str = "PKU-Alignment/PKU-SafeRLHF-30K"
+    SPLIT: str = "test"
+
+
+class SafeRLHF10KTrainDataset(SafeRLHFDataset):
+    NAME: str = "PKU-SafeRLHF-10K/train"
+    ALIASES: tuple[str, ...] = ("PKU-Alignment/PKU-SafeRLHF-10K/train",)
+    PATH: str = "PKU-Alignment/PKU-SafeRLHF-10K"
+    SPLIT: str = "train"

llm/Alignment/PPO/data/supervised.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Dataset class for preference training."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import paddle
+from typing_extensions import TypedDict  # Python 3.10+
+
+from .base import (
+    IGNORE_INDEX,
+    PROMPT_ASSISTANT,
+    PROMPT_BEGIN,
+    PROMPT_USER,
+    CollatorBase,
+    RawSample,
+    TokenizedDataset,
+    format_prompt,
+    right_padding,
+)
+
+__all__ = [
+    "SupervisedDataset",
+    "SupervisedCollator",
+    "SupervisedSample",
+    "SupervisedBatch",
+]
+
+
+class SupervisedSample(TypedDict, total=True):
+    input_ids: paddle.Tensor  # size = (L,)
+    labels: paddle.Tensor  # size = (L,)
+
+
+class SupervisedBatch(TypedDict, total=True):
+    input_ids: paddle.Tensor  # size = (B, L)
+    labels: paddle.Tensor  # size = (B, L)
+    attention_mask: paddle.Tensor  # size = (B, L)
+
+
+class SupervisedDataset(TokenizedDataset):
+    def preprocess(self, raw_sample: RawSample) -> SupervisedSample:
+        if raw_sample.get("input") is None and raw_sample.get("dialogue") is None:
+            raise ValueError("Either `input` or `dialogue` must be provided.")
+        if raw_sample.get("input") is not None and raw_sample.get("dialogue") is not None:
+            raise ValueError("At most one of `input` and `dialogue` can be provided.")
+
+        if raw_sample.get("input") is not None:
+            input = raw_sample["input"]  # pylint: disable=redefined-builtin
+            if not isinstance(input, str):
+                raise ValueError(f"Unsupported type of `input`: {type(input)}. Expected: str.")
+            prompt = format_prompt(input=input, eos_token=self.tokenizer.eos_token)
+            answer = raw_sample["answer"]
+            text = prompt + answer + self.tokenizer.eos_token
+
+            input_ids = self.tokenize(text)
+            labels = input_ids.copy()
+            # Mask non-assistant input
+            labels[: len(self.tokenize(prompt))] = IGNORE_INDEX
+            return {"input_ids": input_ids, "labels": labels}
+
+        dialogue = raw_sample["dialogue"]  # is not None
+        text = PROMPT_BEGIN
+        offsets = [0]
+        for i, line in enumerate(dialogue):
+            if i % 2 == 0:
+                # User input
+                text += PROMPT_USER.format(input=line) + PROMPT_ASSISTANT
+            else:
+                # Assistant input
+                text += line + self.tokenizer.eos_token
+            input_ids = self.tokenize(text)
+            offsets.append(len(input_ids))
+
+        labels = input_ids.copy()
+        # Mask non-assistant input
+        for begin, end in zip(offsets[::2], offsets[1::2]):
+            labels[begin:end] = IGNORE_INDEX
+
+        return {
+            "input_ids": input_ids,  # size = (L,)
+            "labels": labels,  # size = (L,)
+        }
+
+    def get_collator(
+        self, shift: bool = False
+    ) -> Callable[[list[dict[str, paddle.Tensor]]], dict[str, paddle.Tensor]]:
+        return SupervisedCollator(self.tokenizer.pad_token_id, shift)
+
+
+class SupervisedCollator(CollatorBase):
+    def __init__(self, pad_token_id: int, shift: bool = False) -> None:
+        """Initialize a collator."""
+        super().__init__(pad_token_id)
+        self._shift = shift
+
+    def __call__(self, samples: list[SupervisedSample]) -> SupervisedBatch:
+        input_ids = right_padding(
+            [sample["input_ids"] for sample in samples],
+            padding_value=self.pad_token_id,
+        )
+        labels = right_padding(
+            [sample["labels"] for sample in samples],
+            padding_value=IGNORE_INDEX,
+        )
+        if self._shift:
+            input_ids = input_ids[:, :-1]
+            labels = labels[:, 1:]
+        attention_mask = input_ids != self.pad_token_id
+        return {
+            "input_ids": input_ids,  # size = (B, L)
+            "labels": labels,  # size = (B, L)
+            "attention_mask": attention_mask,  # size = (B, L)
+        }

llm/Alignment/PPO/infer_utils.py

@@ -0,0 +1,383 @@
+# 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 copy
+import inspect
+import types
+from contextlib import contextmanager
+
+import paddle
+import paddle.distributed as dist
+from comm_utils import cleanup_tensor_space, offload_tensor_to_cpu, reload_tensor_to_gpu
+from paddle.utils import try_import
+from trainer_utils import guard_set_args
+
+import paddlenlp
+from paddlenlp.trainer.trainer import Trainer, logger
+from paddlenlp.transformers import PretrainedModel, PretrainedTokenizer
+from paddlenlp.transformers.model_utils import dtype_guard
+
+
+class Predictor:
+    def __init__(self, config, model: PretrainedModel = None, tokenizer: PretrainedTokenizer = None):
+        self.model_config = model.config
+        self.config = config
+        self.tokenizer = tokenizer
+        self.model = model
+        self.is_available = False
+        self._weights_mapping = None
+        # TODO(guosheng): Removde dependency on llm.Predictor
+        # 1. buffer_maker creates caches and other buffer inputs can be shared
+        # among multi time prediction. define caches and extra inputs creation
+        # method instead of using predictor.__init__
+        # 2. inputs_processer creates caches and other inputs can be shared among
+        # multi time prediction. define caches and extra inputs creation method
+        # instead of using predictor.__init__
+        from predictor import InferencePredictorMixin
+
+        self._buffer_maker = types.MethodType(InferencePredictorMixin.__init__, self)
+        self._inputs_processer = types.MethodType(InferencePredictorMixin._preprocess, self)
+
+    @staticmethod
+    def create_predictor(trainer):
+        from predictor import (
+            PdArgumentParser,
+            PredictorArgument,
+            get_model_max_position_embeddings,
+        )
+
+        # create infer model
+        # NOTE: infer model use static name param_attr to create and cannot be
+        # created multiple times.
+        def create_infer_model(model, dtype, set_state=False):
+            from models.infer_model_utils import patch_infer_generate
+
+            # apply patches to make FuseMT adapt
+            patch_infer_generate(
+                eos_token_id=trainer.tokenizer.eos_token_id, pad_token_id=trainer.tokenizer.pad_token_id
+            )
+            config = copy.deepcopy(model.config)
+            hcg = dist.fleet.get_hybrid_communicate_group()  # may differ with training
+            config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
+            config.tensor_parallel_rank = hcg.get_model_parallel_rank()
+            config.weight_only_quant_bits = -1
+            config.quant_type = None
+            config.use_cachekv_int8 = False
+            config.single_card_ptq = True
+            infer_model_cls = getattr(paddlenlp.experimental.transformers, model.__class__.__name__ + "InferenceModel")
+            # ori_init_weights = infer_model_cls.init_weights
+            # infer_model_cls.init_weights = lambda self: None
+            with dtype_guard(dtype):
+                infer_model = infer_model_cls(config)
+            # infer_model_cls.init_weights = ori_init_weights
+
+            if set_state:
+                state_dict = {}
+                for k, v in model.state_dict().items():
+                    # state_dict[k] = np.from_dlpack(paddle.utils.dlpack.to_dlpack(v))
+                    state_dict[k] = v.numpy()
+                infer_model.set_state_dict(state_dict)
+            return infer_model
+
+        # to avoid oom, clear param of infer_model imediately
+        ori_creat_param = paddle.nn.Layer.create_parameter
+
+        def _create_param(self, *args, **kwargs):
+            param = ori_creat_param(self, *args, **kwargs)
+            param._clear_data()
+            # param._clear()
+            return param
+
+        paddle.nn.Layer.create_parameter = _create_param
+        # trainer might use an extra model instead of trainer.model for eval
+        eval_model = getattr(trainer, "_inner_eval_model", None)
+        infer_model = create_infer_model(trainer.model if eval_model is None else eval_model, dtype=trainer.amp_dtype)
+        paddle.nn.Layer.create_parameter = ori_creat_param
+        # for k, v in infer_model.state_dict().items():
+        #     v._clear()
+
+        # create predictor
+        parser = PdArgumentParser((PredictorArgument,))
+        predictor_args = parser.parse_dict(
+            {
+                "src_length": get_model_max_position_embeddings(  # can be changed dynamically by predictor.input_length
+                    trainer.model.config if eval_model is None else eval_model.config
+                ),
+                "max_length": trainer.args.max_length,
+                "dtype": trainer.amp_dtype,
+                "batch_size": trainer.args.per_device_train_batch_size,
+                # infer model do not support top_k, and differ with non-infer model
+                # generation which gets default top_K=50 using generation_config.top_k
+                "top_p": trainer.args.top_p,
+                "temperature": trainer.args.temperature,
+                "repetition_penalty": trainer.args.repetition_penalty,
+            }
+        )[0]
+        policy_predictor = Predictor(predictor_args, model=infer_model, tokenizer=trainer.tokenizer)
+        return policy_predictor
+
+    def _create_caches(self):
+        """inputs can be reused among multiple predictions, such as cache"""
+        if hasattr(self, "cache_kvs_shape"):  # has created cache
+            input_length = getattr(self, "input_length", 0)
+            # TODO(guosheng): better way to get history max cahce length, we can
+            # not get cahce length form cache tensor when not know cache layout
+            if input_length <= self.config.src_length:  # reuse cahce
+                return
+            else:  # create longer cache
+                self._clear_caches()
+        self.config.src_length = getattr(self, "input_length", self.config.src_length)
+        if not hasattr(self, "_buffer_attrs"):
+            pre_attrs = set(self.__dict__.keys())
+        self.cache_kvs_shape = self.model.get_cache_kvs_shape(
+            self.model_config, self.config.batch_size, self.config.total_max_length
+        )
+        self._buffer_maker(self.config, self.tokenizer)
+        if not hasattr(self, "_buffer_attrs"):
+            self._buffer_attrs = set(self.__dict__.keys()) - pre_attrs
+
+    def _clear_caches(self):
+        # del or offload
+        for attr in self._buffer_attrs:
+            delattr(self, attr)
+
+    def disable(self, model, onload_model=True):
+        # clear caches
+        self._clear_caches()
+        # clear params
+        for _, param in self.model.state_dict().items():
+            param._clear_data()
+            # param._clear()
+        if onload_model:
+            model.to(paddle.device.get_device())
+        self.is_available = False
+
+    def enable(self, model, offload_model=True):
+        if self.is_available:
+            return
+        # set params
+        self.set_state_dict(model, offload_model)
+        self.is_available = True
+
+    @paddle.no_grad()
+    def set_state_dict(self, model, offload_model=True):
+        offload_place = paddle.CUDAPinnedPlace()
+        state_dict = {}
+        for k, v in model.state_dict().items():
+            state_dict[k] = v
+
+        if getattr(self, "_weights_mapping", None) is None:
+            self._weights_mapping = self.model.get_weights_mapping()
+
+        # non_share_params = []
+        for k, v in self._weights_mapping.items():
+            param, (convert_fun, args) = k, v
+            args = [state_dict[name] for name in args]
+            value = convert_fun(*args)
+            if offload_model:
+                for arg in args:
+                    # shared params no need to offload
+                    if value is not arg:
+                        cpu_arg = arg._copy_to(offload_place, blocking=False)
+                        cpu_arg._share_buffer_to(arg)
+            if not isinstance(value, paddle.Tensor):
+                param.set_value(value)
+            # elif isinstance(value.place, paddle.CUDAPlace):
+            elif value.place.is_gpu_place():
+                # NOTE: _share_buffer_to seems do not work
+                # value._share_buffer_to(param)
+                # value._share_underline_tensor_to(param)
+                param.get_tensor()._share_data_with(value.get_tensor())
+            else:
+                param.copy_(value, True)
+
+        paddle.device.cuda.synchronize()
+
+    def _preprocess(self, source):
+        # make cache when infer happens to get actual shape to save memory
+        self._create_caches()
+        with guard_set_args(self.config, {"src_length": getattr(self, "input_length", self.config.src_length)}):
+            inputs = self._inputs_processer(source)
+        # We want to use a defined input_length to create cache and input_ids.
+        # However predictor could not use a specified length to pad currently.
+        # Thus we use this way to let get the actual input length.
+        self.infer_input_length = inputs["input_ids"].shape[-1]
+        return inputs
+
+    @paddle.no_grad()
+    def _infer(self, inputs):
+        for key in inputs.keys():
+            if paddle.is_tensor(inputs[key]):
+                continue
+            if isinstance(inputs[key], list):
+                if paddle.is_tensor(inputs[key]):
+                    continue
+                inputs[key] = [paddle.to_tensor(item) for item in inputs[key]]
+            else:
+                inputs[key] = paddle.to_tensor(inputs[key])
+
+        inputs["cache_kvs"] = self.cache_kvs
+        return self.model.generate(**inputs)
+
+    def _postprocess(self, predictions):
+        return predictions
+
+    @paddle.no_grad()
+    def predict(self, input_texts: str | list[str]):
+        tokenized_source = self._preprocess(input_texts)
+        predictions = self._infer(tokenized_source)
+        decoded_predictions = self._postprocess(predictions)
+        return decoded_predictions
+
+
+policy_predictor: Predictor = None
+
+
+@contextmanager
+def infer_guard(trainer, offload_model=True):
+    # trainer might use an extra model instead of trainer.model for eval
+    eval_model = getattr(trainer, "_inner_eval_model", None)
+    model = trainer.model if eval_model is None else eval_model
+
+    # PipelineParallel does not support inference speedup
+    if not getattr(trainer, "use_fusemt", False) or isinstance(
+        model, (dist.fleet.meta_parallel.PipelineLayer, dist.fleet.model.PipelineParallel)
+    ):
+        yield
+        return
+
+    try:
+        try_import("paddlenlp_ops")
+    except:
+        logger.warning("paddlenlp_ops does not exist, please install paddlenlp_ops for generation speedup.")
+        yield
+        return
+
+    global policy_predictor
+    if policy_predictor is None:
+        policy_predictor = Predictor.create_predictor(trainer)
+    if not policy_predictor.is_available:
+        policy_predictor.enable(model, offload_model=offload_model)
+
+    # TODO(guosheng): patch for dist.all_recude to use tp group, fix it later
+    ori_all_reduce = dist.all_reduce
+    ori_broadcast = dist.broadcast
+    hcg = dist.fleet.get_hybrid_communicate_group()
+    dist.all_reduce = lambda x: ori_all_reduce(x, group=hcg.get_model_parallel_group())
+    dist.broadcast = lambda x, rank: ori_broadcast(
+        x, src=hcg.get_model_parallel_group_src_rank(), group=hcg.get_model_parallel_group()
+    )
+    yield
+    dist.all_reduce = ori_all_reduce
+    dist.broadcast = ori_broadcast
+
+    policy_predictor.disable(model, onload_model=offload_model)
+
+
+class InferEvalModel:
+    """For faster generation, not support PipelineParallel yet."""
+
+    def __init__(self, trainer: Trainer):
+        # trainer might use an extra model instead of trainer.model for eval
+        eval_model = getattr(trainer, "_inner_eval_model", None)
+        self.model: PretrainedModel = trainer.model if eval_model is None else eval_model
+        self.tokenizer: PretrainedTokenizer = trainer.tokenizer
+        self.trainer = trainer
+
+    def enable(self):
+        trainer = self.trainer
+        if trainer.model is not self.model:
+            trainer.export_evaluate_model(
+                trainer.model,
+                self.model,
+                with_offload="train_model" in trainer.args.offload_level,
+            )
+        else:
+            reload_tensor_to_gpu(self.model.state_dict())
+
+    def disable(self):
+        trainer = self.trainer
+        if trainer.model is not self.model:
+            cleanup_tensor_space(self.model.state_dict())
+        else:
+            offload_tensor_to_cpu(self.model.state_dict())
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.model, name)
+
+    def eval(self):
+        self.model.eval()
+
+    def train(self):
+        self.model.train()
+
+    def __call__(self, *args, **kwargs):
+        # assert model is on GPU
+        assert policy_predictor is None or not policy_predictor.is_available
+        return self.model(*args, **kwargs)
+
+    def generate(self, *args, **kwargs):
+        if policy_predictor is None or not policy_predictor.is_available:
+            return self.model.generate(*args, **kwargs)
+
+        arg_dict = inspect.signature(self.model.generate).bind(*args, **kwargs).arguments
+        input_ids = arg_dict["input_ids"]
+        generation_config = arg_dict["generation_config"]
+        # convert text and tokenize again to convert left padding to right padding
+        # remove this if inputs is right padding
+        # TODO(guosheng): allow to use right padding to infer directly
+        prompts = self.tokenizer.batch_decode(input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        # decoded prompts has been applied with chat_template
+        # NOTE(guosheng): Whether to add special token should be checked, None
+        # chat_template would not add special token in predictor, since it assumes
+        # chat_template includes special tokens. While Beaver dataset tokenization
+        # does not use chat_template, it uses hard coded template which excludes
+        # special tokens.
+        with guard_set_args(
+            policy_predictor.tokenizer,
+            {
+                # predictor use right padding for infer model by default
+                # "padding_side": "right",
+                # "chat_template": None
+            },
+        ):
+            # NOTE: right padding in predictor according to prompt might have a
+            # different length with input_ids, espically when input_ids has more
+            # paddings than the necessary. Thus pass input_length to predictor to:
+            # 1. use a consistent length to replace input_ids back to output to
+            #    keep the same padding format. however predictor could not use a
+            #    specified length to pad currently
+            # 2. allow to use a dynamic length for memory efficiency (by a smaller
+            #    cache)
+            policy_predictor.input_length = input_ids.shape[-1]
+            outputs = policy_predictor.predict(prompts)
+
+        if generation_config.trunc_input:
+            outputs = (outputs[0][:, policy_predictor.infer_input_length :],)
+            return outputs
+
+        if policy_predictor.input_length != policy_predictor.infer_input_length:
+            outputs = (paddle.concat([input_ids, outputs[0][:, policy_predictor.infer_input_length :]], axis=-1),)
+            return outputs
+
+        outputs = (outputs[0],)
+        if self.tokenizer.padding_side == "left":
+            # convert back to left padding inputs
+            outputs[0][:, : input_ids.shape[-1]] = input_ids
+        return outputs

llm/Alignment/PPO/models/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Utility functions for score models."""
+
+from .score_model_utils import AutoModelForScore, ScoreModelOutput
+
+__all__ = ["AutoModelForScore", "ScoreModelOutput"]

llm/Alignment/PPO/models/infer_model_utils.py

@@ -0,0 +1,185 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Utilities for inference model."""
+
+import numpy as np
+import paddle
+
+
+def patch_paddlenlp_ops(eos_token_id, pad_token_id):
+    import paddlenlp_ops
+
+    paddlenlp_ops.save_with_output = lambda *args, **kwargs: None
+
+    # TODO(guosheng): update the custom op code directly.
+    ori_set_ends = paddlenlp_ops.set_stop_value_multi_ends
+
+    def _set_ends(topk_ids, stop_flags, end_ids, mode):
+        # infer model uses eos_token_id to pad and discriminate ending,
+        # patch to use pad_token_id to pad to unify with non-infer model.
+        topk_ids_out, stop_flags_out = ori_set_ends(topk_ids, stop_flags, end_ids, mode)
+        if pad_token_id != eos_token_id:
+            topk_ids_out = paddle.where(stop_flags, pad_token_id, topk_ids_out)
+        return topk_ids_out, stop_flags_out
+
+    paddlenlp_ops.set_stop_value_multi_ends = _set_ends
+
+
+def patch_infer_generate(eos_token_id, pad_token_id):
+    """patches for inference model to make FuseMT adapt"""
+    # patch paddlenlp_ops, maybe update the custom op code directly later
+    # NOTE: should patch paddlenlp_ops before infer model import
+    patch_paddlenlp_ops(eos_token_id, pad_token_id)
+
+    # patch get_weights_mapping for InferenceModel
+    patch_infer_model()
+
+    # patch GenerationInferenceModel.sample
+    from paddlenlp.experimental.transformers.generation_utils import (
+        GenerationInferenceModel,
+    )
+
+    ori_update_model_kwargs = GenerationInferenceModel.update_model_kwargs_for_generation
+
+    def _update_model_kwargs(self, *args, **kwargs):
+        # update_model_kwargs_for_generation only returns , hack to.
+        model_kwargs = ori_update_model_kwargs(self, *args, **kwargs)
+        next_tokens = model_kwargs["next_tokens"]
+        all_input_ids = paddle.concat([model_kwargs["all_input_ids"], next_tokens], axis=1)
+        model_kwargs["next_tokens"] = all_input_ids
+        model_kwargs["all_input_ids"] = None
+        return model_kwargs
+
+    GenerationInferenceModel.update_model_kwargs_for_generation = _update_model_kwargs
+
+
+_model_weights_mapping_dict = {}
+
+
+def register_model(model_cls_name):
+    def mark_cls_name(func):
+        # Do not register here although we can, otherwise infer model would import
+        # before paddlenlp_ops.
+        _model_weights_mapping_dict[model_cls_name] = func
+        return func
+
+    return mark_cls_name
+
+
+def patch_infer_model():
+    import paddlenlp.experimental.transformers as infer_transformers
+
+    for model_cls_name, get_weights_mapping in _model_weights_mapping_dict.items():
+        model_cls = getattr(infer_transformers, model_cls_name)
+        model_cls.get_weights_mapping = get_weights_mapping
+
+
+@register_model("LlamaForCausalLMInferenceModel")
+def get_weights_mapping(self):
+    """model to infer model"""
+    head_size = self.config.hidden_size // self.config.num_attention_heads
+
+    def _concat_qkv(q, k, v):
+        if isinstance(q, paddle.Tensor):
+            concated_qkv_weight = paddle.concat([q, k, v], axis=-1).T.reshape(
+                [
+                    3 * (self.config.num_attention_heads // self.config.tensor_parallel_degree) * (head_size),
+                    self.config.hidden_size,
+                ]
+            )
+        else:
+            concated_qkv_weight = (
+                np.concatenate(
+                    [q, k, v],
+                    axis=-1,
+                )
+                .transpose(1, 0)
+                .reshape(
+                    3 * (self.config.num_attention_heads // self.config.tensor_parallel_degree) * (head_size),
+                    self.config.hidden_size,
+                )
+            )
+
+        return concated_qkv_weight
+
+    def _concat_ffn1(w1, w2):
+        if isinstance(w1, paddle.Tensor):
+            concated_ffn1_weight = paddle.concat([w1, w2], axis=-1)
+        else:
+            concated_ffn1_weight = np.concatenate([w1, w2], axis=-1)
+        return concated_ffn1_weight
+
+    identity = lambda x: x
+
+    weight_mapping = {}
+    weight_mapping[self.lm_head.weight] = [
+        identity,
+        [
+            "lm_head.weight",
+        ],
+    ]
+    weight_mapping[self.llama.embed_tokens.weight] = [
+        identity,
+        [
+            "llama.embed_tokens.weight",
+        ],
+    ]
+    weight_mapping[self.llama.norm.weight] = [
+        identity,
+        [
+            "llama.norm.weight",
+        ],
+    ]
+    for idx in range(self.config.num_hidden_layers):
+        weight_mapping[self.llama.transformer_block.qkv_weights[idx]] = [
+            _concat_qkv,
+            [
+                f"llama.layers.{idx}.self_attn.q_proj.weight",
+                f"llama.layers.{idx}.self_attn.k_proj.weight",
+                f"llama.layers.{idx}.self_attn.v_proj.weight",
+            ],
+        ]
+        weight_mapping[self.llama.transformer_block.ffn1_weights[idx]] = [
+            _concat_ffn1,
+            [
+                f"llama.layers.{idx}.mlp.gate_proj.weight",
+                f"llama.layers.{idx}.mlp.up_proj.weight",
+            ],
+        ]
+        weight_mapping[self.llama.transformer_block.linear_weights[idx]] = [
+            identity,
+            [
+                f"llama.layers.{idx}.self_attn.o_proj.weight",
+            ],
+        ]
+        weight_mapping[self.llama.transformer_block.ffn2_weights[idx]] = [
+            identity,
+            [
+                f"llama.layers.{idx}.mlp.down_proj.weight",
+            ],
+        ]
+        weight_mapping[self.llama.transformer_block.ln_scales[idx]] = [
+            identity,
+            [
+                f"llama.layers.{idx}.input_layernorm.weight",
+            ],
+        ]
+        weight_mapping[self.llama.transformer_block.ffn_ln_scales[idx]] = [
+            identity,
+            [
+                f"llama.layers.{idx}.post_attention_layernorm.weight",
+            ],
+        ]
+    return weight_mapping

llm/Alignment/PPO/models/model_pp.py

@@ -0,0 +1,261 @@
+# 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.
+
+import paddle
+import paddle.nn as nn
+from paddle.distributed.fleet.meta_parallel import LayerDesc
+
+from paddlenlp.transformers import LlamaForCausalLM, LlamaForCausalLMPipe
+from paddlenlp.transformers.llama.modeling import LlamaDecoderLayer
+from paddlenlp.transformers.llama.modeling_pp import (
+    LlamaRMSNormPipe,
+    parse_args,
+    return_args,
+)
+
+from .pp_model_utils import fwd_args_to_dict, get_expected_keys, pad_batches_inputs
+from .ppo_model_utils import (
+    RLHFPPOMixedLoss,
+    RLHFValueLoss,
+    create_loss,
+    make_position_ids,
+)
+from .score_model_utils import ScoreModelMixin
+
+# patches for base pipe model
+# non-pipe model class, can be used to parse and convert forward args
+# mainly used for generation with PipelienParallel model
+LlamaForCausalLMPipe._non_pipe_model_class = LlamaForCausalLM
+LlamaForCausalLMPipe._non_pipe_decoder_layer_class = LlamaDecoderLayer
+
+
+class LlamaPolicyPipe(LlamaForCausalLMPipe):
+    # TODO(guosheng): maybe make a Mixin is better
+
+    @fwd_args_to_dict
+    def _prepare_pipeline_inputs_func(self, inputs):
+        # first_stage_keys = ["input_ids", "attention_mask"]
+        first_stage_keys = ["input_ids", "attention_mask", "position_ids"]
+        # last_stage_keys = [
+        #     "labels", "input_ids", "log_probs", "advantages", "sequence_mask"
+        # ]
+        # TODO(guosheng): make input keys same with model arg names, maybe we
+        # can use inspect and set as global var which can then be used here and
+        # in PPOTrainer.
+        last_stage_keys = ["labels", "input_ids", "old_log_probs", "reward_advantages", "sequence_mask"]
+
+        if type(inputs) is dict:
+            # for left padding, position_ids is nececessary
+            if "position_ids" not in inputs:
+                inputs["position_ids"] = make_position_ids(inputs["attention_mask"])
+            # ppo-loss and ptx-loss need different labels, and data iter provides
+            # corrensponding data, thus add the not provided fields here.
+            # policy trian and infer has different inputs, infer uses position_ids.
+            # for key in last_stage_keys:
+            for key in first_stage_keys + last_stage_keys:
+                if key not in inputs:
+                    inputs[key] = None
+            return [
+                get_expected_keys(inputs, first_stage_keys),
+                get_expected_keys(inputs, last_stage_keys),
+            ]
+
+        for data in inputs:
+            # for key in last_stage_keys:
+            for key in first_stage_keys + last_stage_keys:
+                if key not in data:
+                    if key == "position_ids":
+                        data[key] = make_position_ids(data["attention_mask"])
+                        continue
+                    data[key] = None
+        # keys = list(inputs[0].keys())
+        inputs_batch = {key: [data.get(key) for data in inputs] for key in first_stage_keys + last_stage_keys}
+        # NOTE(guosheng): PipelineParallel requires send/recv tensors among
+        # micro-batches/accu-steps have the same shape. Thus pad here, maybe
+        # should make data collator do padding and pad optionally here, since
+        # padding strategy may not be clear here.
+        # 1. For input_ids/attention_mask/labels (prompt+target) padding:
+        # Some data fields, such as input_ids/attention_mask/labels, should
+        # have same shape after padding, and each of them cannot pad only
+        # according to its own max length which might be different since the
+        # filed value is None for different batches/tasks.
+        src_tgt_keys = ["input_ids", "attention_mask", "labels", "position_ids"]
+        max_len = max([x.shape[-1] for x in inputs_batch["input_ids"]])
+        pad_len = [max_len - x.shape[-1] for x in inputs_batch["input_ids"]]
+        for key in src_tgt_keys:
+            # Do not pad position_ids with 0 since 0s in position_ids has special
+            # usage in reward model. We use 1 to pad.
+            padding_value = self._ignore_index if key == "labels" else 1 if key == "position_ids" else 0
+            inputs_batch[key] = pad_batches_inputs(inputs_batch[key], padding_value, pad_len=pad_len)
+        # 2. For old_log_probs/reward_advantages/sequence_mask (target) padding:
+        # hard to pad acorss batches, think in some cases one batch might have the
+        # longest prompt+target length but the shortest target lengh, which might
+        # cause mismatch between inputs with prompt+target length and labels with
+        # target length. NOTE: however trick can be used here, label fields with
+        # target length such as old_log_probs/reward_advantages/sequence_mask do
+        # not need to join comm and thus there is no need to keep same shape among
+        # batches of accumulation steps, they just need to pad as prompt+target
+        # fields such as input_ids.
+        tgt_keys = ["old_log_probs", "reward_advantages", "sequence_mask"]
+        for key in tgt_keys:
+            padding_value = 0
+            inputs_batch[key] = pad_batches_inputs(inputs_batch[key], padding_value, pad_len=pad_len)
+        # for key, value in inputs_batch.items():
+        #     padding_value = self._ignore_index if key == "labels" else 0
+        #     max_len = max_len if key in [
+        #         "input_ids", "attention_mask", "labels"
+        #     ] else None
+        #     inputs_batch[key] = pad_batches_inputs(value, padding_value, max_len)
+        return [
+            get_expected_keys(inputs_batch, first_stage_keys),
+            get_expected_keys(inputs_batch, last_stage_keys),
+        ]
+
+    def __init__(self, config, **kwargs):
+        # NOTE: make _sequential_layers/_single_to_pp_mapping/_pp_to_single_mapping
+        # instance attrs instead of class attrs to support more than one pipeline
+        # models. Maybe make all sequential_layers add once.
+        self._sequential_layers = []
+        self._single_to_pp_mapping = None
+        self._pp_to_single_mapping = None
+        # To be consistent with score model init and allow hyper-param be passed
+        # using __init__/from_pretrained
+        self._init_kwargs = kwargs
+        super().__init__(config)
+        self._ignore_index = self._loss_fn.sft_criterion.ignore_index
+
+    def get_loss_fn(self, config):
+        return create_loss(RLHFPPOMixedLoss, config, self._init_kwargs)
+
+    @property
+    def head_out_meta(self):
+        """mainly for eval/generation with PipelineParallel"""
+        # None means to use actual data info
+        return paddle.static.InputSpec(shape=[None, None, self.config.vocab_size], dtype=None)
+
+
+class _LlamaRMSNormPipe(LlamaRMSNormPipe):
+    """
+    We need position_ids for reward model, so wrap LlamaRMSNormPipe to pass position_ids
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def forward(self, args):
+        hidden_states, attention_mask, position_ids, alibi = parse_args(args)
+        return return_args(self.norm(hidden_states), attention_mask, position_ids)
+
+
+# LayerDesc of PipelineParallel requires head to be a nn.Layer
+class ValueHead(nn.Layer, ScoreModelMixin):
+    def __init__(self, config, **kwargs):
+        super().__init__()
+        self.config = config
+        self.init_score_head(config, hidden_size=config.hidden_size, **kwargs)
+
+    def forward(self, args):
+        # attention_mask passed from pre-stage is shaped (bs, 1, seq_len, seq_len)
+        hidden_state, attention_mask, position_ids, alibi = parse_args(args)
+        outputs = self.get_score(
+            hidden_state, attention_mask=attention_mask, position_ids=position_ids, return_dict=True
+        )
+        return outputs
+
+
+class LlamaValuePipe(LlamaForCausalLMPipe):
+    # TODO(guosheng): maybe make a Mixin is better
+
+    @fwd_args_to_dict
+    def _prepare_pipeline_inputs_func(self, inputs):
+        # ValueHead/get_score needs original attention_mask or position_ids,
+        # while attention_mask passed from pre-stage is not the original, thus
+        # hack for position_ids here.
+        # Maybe add position_ids into inputs later and use position_ids instead
+        # of attention_mask to get score not only for pipeline parallel.
+        first_stage_keys = ["input_ids", "attention_mask", "position_ids"]
+        # TODO(guosheng): make input keys same with model arg names, maybe we
+        # can use inspect and set as global var which can then be used here and
+        # in PPOTrainer.
+        last_stage_keys = ["old_reward_values", "reward_returns", "sequence_mask"]
+
+        if type(inputs) is dict:
+            if "position_ids" not in inputs:
+                inputs["position_ids"] = make_position_ids(inputs["attention_mask"])
+
+            return [
+                get_expected_keys(inputs, first_stage_keys),
+                get_expected_keys(inputs, last_stage_keys),
+            ]
+
+        for data in inputs:
+            if "position_ids" not in data:
+                data["position_ids"] = make_position_ids(data["attention_mask"])
+        # keys = list(inputs[0].keys())
+        inputs_batch = {key: [data.get(key) for data in inputs] for key in first_stage_keys + last_stage_keys}
+        # 1. For input_ids/attention_mask (prompt+target) padding:
+        # src_tgt_keys = ["input_ids", "attention_mask"]
+        src_tgt_keys = ["input_ids", "attention_mask", "position_ids"]
+        max_len = max([x.shape[-1] for x in inputs_batch["input_ids"]])
+        pad_len = [max_len - x.shape[-1] for x in inputs_batch["input_ids"]]
+        for key in src_tgt_keys:
+            # Do not pad position_ids with 0 since 0s in position_ids has special
+            # usage in reward model. We use 1 to pad.
+            padding_value = self._ignore_index if key == "labels" else 1 if key == "position_ids" else 0
+            inputs_batch[key] = pad_batches_inputs(inputs_batch[key], padding_value, pad_len=pad_len)
+        # 2. For old_reward_values/reward_returns/sequence_mask (target) padding:
+        tgt_keys = ["old_reward_values", "reward_returns", "sequence_mask"]
+        for key in tgt_keys:
+            padding_value = 0
+            inputs_batch[key] = pad_batches_inputs(inputs_batch[key], padding_value, pad_len=pad_len)
+        # for key, value in inputs_batch.items():
+        #     inputs_batch[key] = pad_batches_inputs(value, padding_value=0)
+        # if "position_ids" not in inputs[0]:
+        #     inputs_batch["position_ids"] = [
+        #         make_position_ids(attention_mask) for attention_mask in inputs_batch["attention_mask"]
+        #     ]
+        return [
+            get_expected_keys(inputs_batch, first_stage_keys),
+            get_expected_keys(inputs_batch, last_stage_keys),
+        ]
+
+    def __init__(self, config, **kwargs):
+        # NOTE: make _sequential_layers/_single_to_pp_mapping/_pp_to_single_mapping
+        # instance attrs instead of class attrs to support more than one pipeline
+        # models. Maybe make all sequential_layers add once.
+        self._sequential_layers = []
+        self._single_to_pp_mapping = None
+        self._pp_to_single_mapping = None
+        # To be consistent with score model init and allow hyper-param be passed
+        # using __init__/from_pretrained
+        self._init_kwargs = kwargs
+        super().__init__(config)
+
+    def add_head(self, config):
+        init_kwargs = self._init_kwargs
+        # hack to replace original RMSNormPipe to support ValueHead inputs
+        norm_prefix = self._sequential_layers.pop(-1)["name_prefix"]
+        self.add_sequential_layer(LayerDesc(_LlamaRMSNormPipe, config=config), norm_prefix)
+        self.add_sequential_layer(LayerDesc(ValueHead, config, **init_kwargs), "")
+
+    def get_loss_fn(self, config):
+        return create_loss(RLHFValueLoss, config, self._init_kwargs)
+
+    @property
+    def head_out_meta(self):
+        # None means to use actual data info
+        return (
+            paddle.static.InputSpec(shape=[None, None, 1], dtype=None),
+            paddle.static.InputSpec(shape=[None, 1], dtype=None),
+        )

llm/Alignment/PPO/models/pp_model_utils.py

@@ -0,0 +1,92 @@
+# 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.
+
+import importlib
+import inspect
+
+import paddle
+
+
+def fwd_step_patch(func, output, self, *args, **kwargs):
+    # training patch
+    if self.training and self.is_pipeline_last_stage():
+        if getattr(self, "_step_losses", None):
+            self._step_losses.append(output.detach())
+        else:
+            self._step_losses = [output.detach()]
+
+
+def make_wrapper(func, pre_patch=None, post_patch=None):
+    def wrapper(*args, **kwargs):
+        if pre_patch is not None:
+            pre_patch(func, None, *args, **kwargs)
+        output = func(*args, **kwargs)
+        if post_patch is not None:
+            post_patch(func, output, *args, **kwargs)
+        return output
+
+    return wrapper
+
+
+funcs = [(paddle.distributed.fleet.model.PipelineParallel._forward_step, fwd_step_patch)]
+
+for func in funcs:
+    fun, patch = func
+    module = importlib.import_module(fun.__module__)
+    cls_name = fun.__qualname__[: -len(fun.__name__) - 1]
+    wrap_fun = make_wrapper(fun, post_patch=patch)
+    cls_obj = getattr(module, cls_name)
+    setattr(cls_obj, fun.__name__, wrap_fun)
+
+
+@paddle.no_grad()
+def pad_batches_inputs(inputs, padding_value=0, max_len=None, pad_len=None):
+    """Pad length for tensors shaped [bs, seq_len] to [bs, max(seq_lens)]"""
+    if pad_len is not None:
+        pad_len = [pad_len] * len(inputs) if isinstance(pad_len, int) else pad_len
+    elif max_len is None:
+        # max_len = max([x.shape[-1] for x in inputs if x is not None])
+        max_len = max([x.shape[-1] if isinstance(x, paddle.Tensor) else 0 for x in inputs])
+        pad_len = [max_len - x.shape[-1] if isinstance(x, paddle.Tensor) else 0 for x in inputs]
+    for i in range(len(inputs)):
+        x = inputs[i]
+        # if x is None or x.shape[-1] == max_len:
+        if not isinstance(x, paddle.Tensor) or x.shape[-1] == max_len:
+            continue
+        inputs[i] = paddle.concat([x, paddle.full([x.shape[0], pad_len[i]], padding_value, dtype=x.dtype)], -1)
+    return inputs
+
+
+def get_expected_keys(inputs, keys):
+    ret = tuple([inputs.get(k, None) for k in keys if k in inputs])
+    if len(ret) == 1:
+        ret = ret[0]
+    return ret
+
+
+def fwd_args_to_dict(fun):
+    def _impl(self, *args, **kwargs):
+        try:
+            return fun(self, *args, **kwargs)
+        except TypeError:
+            # otherwise, inputs is any valid format of non_pipe_model forward args,
+            # convert to dict, to support more args format in prediction_pipeline_step
+            # assume no arg is inspect.Parameter.VAR_KEYWORD
+            arg_dict = (
+                inspect.signature(self._non_pipe_model_class.forward).bind(*((self,) + args), **kwargs).arguments
+            )
+            arg_dict.pop("self")
+            return fun(self, arg_dict)
+
+    return _impl

llm/Alignment/PPO/models/ppo_model.py

@@ -0,0 +1,117 @@
+# 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 paddlenlp.transformers import LlamaForCausalLM, PretrainedConfig
+
+from .ppo_model_utils import PolicyOutput, RLHFPPOMixedLoss, RLHFValueLoss, ValueOutput
+from .score_model import LlamaModelForScore
+
+
+# TODO(guosheng): create Mixin and make model classes using metaclass.
+class LlamaPolicyModel(LlamaForCausalLM):
+    def __init__(self, config: PretrainedConfig, **kwargs):
+        super().__init__(config)
+        self.loss_fn = RLHFPPOMixedLoss(config, **kwargs)
+
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=False,
+        past_key_values=None,
+        log_probs=None,
+        advantages=None,
+        sequence_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        outputs = super().forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            labels=None,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs[0]
+        loss = None
+        if labels is not None or advantages is not None:
+            loss = self.loss_fn(logits, (labels, input_ids, log_probs, advantages, sequence_mask))
+        if not return_dict:
+            return (loss,) + outputs if loss is not None else outputs
+
+        return PolicyOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class LlamaValueModel(LlamaModelForScore):
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.loss_fn = RLHFValueLoss(config, **kwargs)
+
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        use_cache=False,
+        past_key_values=None,
+        old_values=None,
+        returns=None,
+        sequence_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        outputs = super().forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        reward_values, rewards = outputs
+        loss = None
+        if returns is not None:
+            loss = self.loss_fn(reward_values, old_values, returns, sequence_mask)
+        if not return_dict:
+            return (loss,) + outputs if loss is not None else outputs
+
+        return ValueOutput(
+            loss=loss,
+            value=reward_values,
+            reward=rewards,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

llm/Alignment/PPO/models/ppo_model_utils.py

@@ -0,0 +1,274 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Utilities for score models."""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+# use LlamaPretrainingCriterion as common PretrainingCriterion
+from paddlenlp.transformers import LlamaPretrainingCriterion as PretrainingCriterion
+from paddlenlp.transformers.model_outputs import ModelOutput
+
+
+@dataclass
+class PolicyOutput(ModelOutput):
+    loss: Optional[paddle.Tensor] = None
+    logits: paddle.Tensor = None
+    # logits_entropy: Optional[paddle.Tensor] = None
+    past_key_values: Optional[Tuple[Tuple[paddle.Tensor]]] = None
+    hidden_states: Optional[Tuple[paddle.Tensor]] = None
+    attentions: Optional[Tuple[paddle.Tensor]] = None
+    cross_attentions: Optional[Tuple[paddle.Tensor]] = None
+
+
+@dataclass
+class ValueOutput(ModelOutput):
+    loss: Optional[paddle.Tensor] = None
+    value: paddle.Tensor = None
+    reward: paddle.Tensor = None
+    past_key_values: Optional[Tuple[Tuple[paddle.Tensor]]] = None
+    hidden_states: Optional[Tuple[paddle.Tensor]] = None
+    attentions: Optional[Tuple[paddle.Tensor]] = None
+    cross_attentions: Optional[Tuple[paddle.Tensor]] = None
+
+
+def merge_fwd_labels(loss_cls):
+    """
+    PipelineParallel and trainer.criterion both use labels as tuple, thus wrap.
+    """
+    ori_fwd = loss_cls.forward
+
+    def loss_fwd(self, predict, labels):
+        return ori_fwd(self, predict, *labels)
+
+    fwd_params = inspect.signature(ori_fwd).parameters
+    # forward(self, predict, label1, label2, ...)
+    loss_cls.label_names = list(fwd_params.keys())[2:]
+    loss_cls.label_default_values = {}
+    for label_name in loss_cls.label_names:
+        if fwd_params[label_name].default is not inspect.Parameter.empty:
+            loss_cls.label_default_values[label_name] = fwd_params[label_name].default
+    loss_cls.forward = loss_fwd
+    return loss_cls
+
+
+def create_loss(loss_cls, config, extra_args, merge_labels=None):
+    """
+    loss_cls(paddle.nn.Layer): loss class
+    config(PratrainedConfig): model config, to be consistent with loss defined
+        in transformers
+    extra_args(dict): create loss with more args not in config
+    merge_labels: use a wrapped loss_cls whose label args are merged into one arg,
+        this is useful to PipelineParallel and trainer.criterion since they only
+        support loss format corresponding to this format.
+    """
+    # TODO(guosheng): merge_labels if loss_cls not
+    ori_fwd = loss_cls.forward
+    if merge_labels:
+        fwd_params = inspect.signature(ori_fwd).parameters
+        if len(fwd_params.keys()) > 3:  # merge_fwd_labels has not done
+            loss_cls = merge_fwd_labels(loss_cls)
+    # forward(self, predict, label1, label2, ...)
+    loss_arg_names = list(inspect.signature(loss_cls.__init__).parameters.keys())[2:]
+    if isinstance(extra_args, dict):
+        loss_kwargs = dict([(name, extra_args[name]) for name in loss_arg_names if name in extra_args])
+    else:
+        # create from TrainingArguments
+        loss_kwargs = dict([(name, getattr(extra_args, name)) for name in loss_arg_names if hasattr(extra_args, name)])
+    loss = loss_cls(config, **loss_kwargs)
+    return loss
+
+
+@paddle.no_grad()
+def make_position_ids(attention_mask, source=None):
+    if len(attention_mask.shape) == 4:  # causal mask
+        position_ids_p1 = attention_mask.cast(paddle.int64).sum(-1)
+        position_ids = position_ids_p1 - 1
+        position_ids = paddle.where(position_ids == -1, position_ids_p1, position_ids)
+        return position_ids[:, 0, :]
+    assert len(attention_mask.shape) == 2  # padding mask
+    attention_mask_bool = attention_mask
+    attention_mask = attention_mask.cast(paddle.int64)
+    position_ids = attention_mask.cumsum(-1) - 1
+    # Make padding positions in source be 0, since reward model use position_ids
+    # plus with padding size (number of 0s) in source to calculate end offsets.
+    # It does not matter when source is left padding and target is right padding
+    # which is the output of non-FuseMT generation, while when using FuseMT whose
+    # output is right padding source and right padding target, we have to set
+    # padding positions in source be 0 to make compatible.
+    if source is not None:
+        src_len = position_ids[:, source.shape[-1] - 1].unsqueeze(-1)
+        position_ids = paddle.where(
+            paddle.logical_and(paddle.logical_not(attention_mask_bool), position_ids <= src_len),
+            attention_mask,
+            position_ids,
+        )
+        return position_ids
+    position_ids = paddle.where(position_ids == -1, attention_mask, position_ids)
+    return position_ids
+
+
+@paddle.no_grad()
+def make_attention_mask(input_ids, pad_id, unk_id=None, past_key_values_length=0, causal_mask=True):
+    attention_mask = input_ids != pad_id
+    if unk_id is not None and pad_id != unk_id:
+        attention_mask = paddle.logical_and(attention_mask, input_ids != unk_id)
+    if not causal_mask:
+        return attention_mask
+
+    batch_size, target_length = input_ids.shape  # target_length: seq_len
+    mask = paddle.tril(paddle.ones((target_length, target_length), dtype="bool"))
+    if past_key_values_length > 0:
+        # [tgt_len, tgt_len + past_len]
+        mask = paddle.concat([paddle.ones([target_length, past_key_values_length], dtype="bool"), mask], axis=-1)
+    # [bs, 1, tgt_len, tgt_len + past_len]
+    causal_mask = mask[None, None, :, :].expand([batch_size, 1, target_length, target_length + past_key_values_length])
+
+    attention_mask = attention_mask[:, None, None, :]
+    expanded_attn_mask = attention_mask & causal_mask
+    return expanded_attn_mask
+
+
+def gather_log_probabilities(logits: paddle.Tensor, labels: paddle.Tensor) -> paddle.Tensor:
+    """Gather log probabilities of the given labels from the logits."""
+    log_probs = F.log_softmax(logits, axis=-1)
+    log_probs_labels = paddle.take_along_axis(log_probs, axis=-1, indices=labels.unsqueeze(axis=-1))
+    return log_probs_labels.squeeze(axis=-1)
+
+
+class RLHFPPOLoss(nn.Layer):
+    def __init__(self, config, clip_range_ratio=0.2):
+        super().__init__()
+        self.clip_range_ratio = clip_range_ratio
+        self.config = config
+
+    def actor_loss_fn(
+        self, log_probs: paddle.Tensor, old_log_probs: paddle.Tensor, advantages: paddle.Tensor, mask: paddle.Tensor
+    ) -> paddle.Tensor:
+        # policy gradient loss
+        ratio = paddle.exp(log_probs - old_log_probs)
+        pg_loss1 = -advantages * ratio
+        pg_loss2 = -advantages * paddle.clip(
+            ratio,
+            1.0 - self.clip_range_ratio,
+            1.0 + self.clip_range_ratio,
+        )
+        return paddle.sum(paddle.maximum(pg_loss1, pg_loss2) * mask) / mask.sum()
+
+    def forward(self, logits, input_ids, old_log_probs, reward_advantages, sequence_mask):
+        log_probs = gather_log_probabilities(logits[:, :-1], input_ids[:, 1:])
+        if log_probs.shape[1] == old_log_probs.shape[1]:
+            # labels (old_log_probs, reward_advantages, sequence_mask) has
+            # src+tgt-1 length, valid length is determined by sequence_mask
+            pass
+        elif log_probs.shape[1] < old_log_probs.shape[1]:
+            # labels (old_log_probs, reward_advantages, sequence_mask) has
+            # src+tgt length and the last one is a padding to be consistent
+            # with input_ids
+            assert log_probs.shape[1] == old_log_probs.shape[1] - 1
+            log_probs = paddle.concat([log_probs, paddle.zeros([log_probs.shape[0], 1], dtype=log_probs.dtype)], -1)
+        else:
+            # labels (old_log_probs, reward_advantages, sequence_mask) has tgt length
+            log_probs = log_probs[:, -old_log_probs.shape[1] :]
+        actor_loss = self.actor_loss_fn(
+            log_probs,
+            old_log_probs,
+            reward_advantages,
+            sequence_mask,
+        )
+        return actor_loss
+
+
+@merge_fwd_labels
+class RLHFPPOMixedLoss(nn.Layer):
+    """provide two losses, one for PPO loss, the other for SFT loss."""
+
+    def __init__(self, config, ptx_coeff=16, clip_range_ratio=0.2):
+        super(RLHFPPOMixedLoss, self).__init__()
+        self.ptx_coeff = ptx_coeff
+        self.ppo_criterion = RLHFPPOLoss(config, clip_range_ratio)
+        self.sft_criterion = PretrainingCriterion(config)
+
+    def forward(self, logits, labels, input_ids, old_log_probs, reward_advantages, sequence_mask):
+        logits = logits if isinstance(logits, paddle.Tensor) else logits[0]
+        loss = None
+        # sft, pt loss
+        if labels is not None:
+            loss = self.ptx_coeff * self.sft_criterion(logits, labels)
+        # ppo loss
+        if reward_advantages is not None:
+            loss = self.ppo_criterion(logits, input_ids, old_log_probs, reward_advantages, sequence_mask)
+
+        return loss
+
+
+@merge_fwd_labels
+class RLHFValueLoss(nn.Layer):
+    def __init__(self, config, clip_range_value=5.0):
+        super().__init__()
+        self.clip_range_value = clip_range_value
+        self.config = config
+
+    def critic_loss_fn(
+        self,
+        values: paddle.Tensor,
+        old_values: paddle.Tensor,
+        returns: paddle.Tensor,
+        mask: paddle.Tensor,
+    ) -> paddle.Tensor:
+        """Compute critic loss."""
+        # TODO(guosheng): use paddle.clip when its min/max can support more than
+        # 0D Tensor
+        values_clipped = paddle.minimum(
+            paddle.maximum(values, old_values - self.clip_range_value), old_values + self.clip_range_value
+        )
+        vf_loss1 = paddle.square(values - returns)
+        vf_loss2 = paddle.square(values_clipped - returns)
+        return 0.5 * paddle.sum(paddle.maximum(vf_loss1, vf_loss2) * mask) / mask.sum()
+
+    def forward(self, reward_values, old_reward_values, reward_returns, sequence_mask):
+        reward_values = reward_values if isinstance(reward_values, paddle.Tensor) else reward_values[0]
+        reward_values = reward_values.squeeze(axis=-1)[:, :-1]
+        if reward_values.shape[1] == old_reward_values.shape[1]:
+            # labels (old_reward_values, reward_returns, sequence_mask) has
+            # src+tgt-1 length, valid length is determined by sequence_mask
+            pass
+        elif reward_values.shape[1] < old_reward_values.shape[1]:
+            # labels (old_reward_values, reward_returns, sequence_mask) has
+            # src+tgt length and the last one is a padding to be consistent
+            # with input_ids
+            assert reward_values.shape[1] == old_reward_values.shape[1] - 1
+            reward_values = paddle.concat(
+                [reward_values, paddle.zeros([reward_values.shape[0], 1], dtype=reward_values.dtype)], -1
+            )
+        else:
+            # labels (old_reward_values, reward_returns, sequence_mask) has
+            # tgt length
+            reward_values = reward_values[:, -old_reward_values.shape[1] :]
+        reward_critic_loss = self.critic_loss_fn(
+            reward_values,
+            old_reward_values,
+            reward_returns,
+            sequence_mask,
+        )
+
+        return reward_critic_loss

llm/Alignment/PPO/models/score_model.py

@@ -0,0 +1,140 @@
+# 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
+
+from typing import Any
+
+import paddle
+from paddle import nn
+
+import paddlenlp
+from paddlenlp.transformers import (
+    LlamaConfig,
+    LlamaModel,
+    LlamaPretrainedModel,
+    PretrainedConfig,
+    PretrainedModel,
+)
+from paddlenlp.transformers.conversion_utils import (
+    StateDictNameMapping,
+    init_name_mappings,
+)
+
+from .score_model_utils import ScoreModelMixin, ScoreModelOutput
+
+
+class LlamaModelForScore(ScoreModelMixin, LlamaPretrainedModel):
+    _keys_to_ignore_on_load_missing = ["lm_head.weight"]
+
+    def __init__(self, config: PretrainedConfig, **kwargs: Any) -> None:
+        super().__init__(config)
+        self.llama = LlamaModel(config)
+
+        # config.architectures = [self.__class__.__name__]
+        self.init_score_head(config, hidden_size=config.hidden_size, **kwargs)
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.llama.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.llama.embed_tokens = value
+
+    def get_decoder(self) -> PretrainedModel:
+        return self.llama
+
+    def set_decoder(self, decoder: PretrainedModel) -> None:
+        self.llama = decoder
+
+    def forward(  # pylint: disable=too-many-arguments
+        self,
+        input_ids: paddle.Tensor,
+        attention_mask: paddle.Tensor,
+        position_ids: paddle.Tensor | None = None,
+        past_key_values: list[paddle.Tensor] | None = None,
+        inputs_embeds: paddle.Tensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> tuple[paddle.Tensor, paddle.Tensor] | ScoreModelOutput:
+        assert attention_mask is not None
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.llama(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = outputs[0]  # size = (B, L, E)
+        return self.get_score(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            return_dict=return_dict,
+        )
+
+    @classmethod
+    def _get_name_mappings(cls, config: LlamaConfig) -> list[StateDictNameMapping]:
+        mappings: list[StateDictNameMapping] = []
+        model_mappings = [
+            ["embed_tokens.weight"],
+            ["norm.weight"],
+        ]
+        for layer_index in range(config.num_hidden_layers):
+            layer_mappings = [
+                [f"layers.{layer_index}.self_attn.q_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.self_attn.k_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.self_attn.v_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.self_attn.o_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.self_attn.rotary_emb.inv_freq"],
+                [f"layers.{layer_index}.mlp.gate_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.mlp.down_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.mlp.up_proj.weight", None, "transpose"],
+                [f"layers.{layer_index}.input_layernorm.weight"],
+                [f"layers.{layer_index}.post_attention_layernorm.weight"],
+            ]
+            model_mappings.extend(layer_mappings)
+
+        init_name_mappings(mappings=model_mappings)
+        # base-model prefix "LlamaModel"
+        if "LlamaModel" not in config.architectures:
+            for mapping in model_mappings:
+                mapping[0] = "model." + mapping[0]
+                mapping[1] = "llama." + mapping[1]
+            model_mappings.append(["lm_head.weight", "lm_head.weight", "transpose"])
+            model_mappings.extend(
+                [
+                    ["score_head.weight", "score_head.weight", "transpose"],
+                    ["normalizer.var", "normalizer.var"],
+                    ["normalizer.mean", "normalizer.mean"],
+                    ["normalizer.count", "normalizer.count"],
+                ]
+            )
+
+        mappings = [StateDictNameMapping(*mapping, index=index) for index, mapping in enumerate(model_mappings)]
+        return mappings
+
+
+paddlenlp.transformers.LlamaModelForScore = LlamaModelForScore

llm/Alignment/PPO/models/score_model_utils.py

@@ -0,0 +1,417 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2023 PKU-Alignment Team. 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.
+"""Utilities for score models."""
+
+from __future__ import annotations
+
+import importlib
+import io
+import json
+from abc import abstractmethod
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import paddle
+import paddle.distributed as dist
+import paddle.nn as nn
+
+from paddlenlp.transformers import PretrainedConfig
+from paddlenlp.transformers.auto.modeling import (
+    MAPPING_NAMES,
+    _BaseAutoModelClass,
+    get_init_configurations,
+    get_name_mapping,
+    is_standard_config,
+)
+from paddlenlp.transformers.model_outputs import ModelOutput
+
+
+class AutoModelForScore(_BaseAutoModelClass):
+    CONFIGURATION_MODEL_MAPPING = get_init_configurations()
+    _pretrained_model_dict = CONFIGURATION_MODEL_MAPPING
+    # "ForScore" might be more consistent with other tasks, while this suffix is
+    # consistent with Beaver now.
+    _task_suffix = "ModelForScore"
+    _name_mapping: OrderedDict[str, Any] = get_name_mapping(_task_suffix)
+    _score_module_name: str = "models.score_model"
+
+    @classmethod
+    def _get_model_class_from_config(cls, pretrained_model_name_or_path, config_file_path, config=None):
+        if config is None:
+            with io.open(config_file_path, encoding="utf-8") as f:
+                config = json.load(f)
+
+        # Get class name corresponds to this configuration
+        if is_standard_config(config):
+            architectures = config["architectures"]
+            init_class = architectures.pop() if len(architectures) > 0 else None
+        else:
+            init_class = config.pop("init_class", None)
+
+        # Get class name corresponds to this task, since we might init Score
+        # model with CausalLM model.
+        init_class = init_class[:-5] if init_class is not None and init_class.endswith("Model") else init_class
+        model_name = None
+        if init_class:
+            for model_flag, name in MAPPING_NAMES.items():
+                if model_flag in init_class:
+                    model_name = model_flag + "Model"
+                    break
+        if model_name is None:
+            raise AttributeError(
+                f"Unable to parse 'architectures' or 'init_class' from {config_file_path}. Also unable to infer model class from 'pretrained_model_name_or_path'"
+            )
+        init_class = cls._name_mapping[model_name + "_Import_Class"]
+        # module_name = cls._name_mapping[init_class]
+        module_name = cls._score_module_name
+
+        import_class = importlib.import_module(module_name)
+        model_class = getattr(import_class, init_class)
+        return model_class
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        return cls._from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+
+@dataclass
+class ScoreModelOutput(ModelOutput):
+    """
+    Output of the score model.
+
+    Args:
+        scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, score_dim, sequence_length)`):
+            Prediction scores of the score model.
+        end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, score_dim)`):
+            Prediction scores of the end of the sequence.
+    """
+
+    scores: paddle.Tensor | None = None  # size = (B, L, D)
+    end_scores: paddle.Tensor | None = None  # size = (B, D)
+
+
+class ScoreModelMixin:
+    """Base class for score models."""
+
+    score_head: nn.Linear
+    normalizer: Normalizer
+    do_normalize: bool = False
+    normalize_function: NormalizeFunction = "affine"
+    _initialized: bool = False
+
+    def init_score_head(self, config: PretrainedConfig, hidden_size: int, **kwargs: Any) -> None:
+        """Initialize the score head."""
+        if self._initialized:
+            return
+
+        config.score_dim = kwargs.pop("score_dim", getattr(config, "score_dim", 1))
+        config.bias = kwargs.pop("bias", getattr(config, "bias", False))
+
+        config.score_type = kwargs.pop("score_type", getattr(config, "score_type", "reward"))
+        if config.score_type == "reward":
+            self.normalize_function = "affine"
+        elif config.score_type == "cost":
+            self.normalize_function = "scale"
+        elif config.score_type == "critic":
+            self.normalize_function = "identity"
+        else:
+            raise ValueError(
+                f"Invalid score type: {config.score_type}. Expected one of 'reward', 'cost', or 'critic'.",
+            )
+
+        config.do_normalize = kwargs.pop(
+            "do_normalize",
+            getattr(config, "do_normalize", False),
+        )
+        self.do_normalize = config.do_normalize
+
+        config.normalizer_type = kwargs.pop(
+            "normalizer_type",
+            getattr(config, "normalizer_type", None),
+        )
+        if config.normalizer_type not in {"RunningMeanStd", "ExponentialMovingAverage", None}:
+            raise ValueError(
+                f"Invalid norm type: {config.normalizer_type}."
+                "Expected one of 'RunningMeadStd', 'ExponentialMovingAverage', or None.",
+            )
+        if config.normalizer_type == "ExponentialMovingAverage":
+            config.momentum = kwargs.pop("momentum", getattr(config, "momentum", None))
+        momentum = getattr(config, "momentum", None)
+
+        self.score_head = nn.Linear(hidden_size, config.score_dim, bias_attr=config.bias)
+        self.normalizer = Normalizer.instantiate(
+            normalizer_type=config.normalizer_type,
+            normalize_function=self.normalize_function,
+            shape=(config.score_dim,),
+            momentum=momentum,
+        )
+
+        mean = getattr(config, "mean", None)
+        var = getattr(config, "var", None)
+        self.normalizer.set_mean_var(mean, var)
+
+        self._initialized = True
+
+    def get_score(
+        self,
+        hidden_state: paddle.Tensor,  # size = (B, L, E)
+        attention_mask: paddle.Tensor | None = None,  # size = (B, L)
+        position_ids: paddle.Tensor | None = None,  # size = (B, L)
+        return_dict: bool | None = None,
+    ) -> ScoreModelOutput:
+        """Forward pass of the score model."""
+        scores = self.score_head(hidden_state)  # size = (B, L, D)
+
+        if position_ids is not None:
+            first_pos = paddle.arange(hidden_state.shape[0]).unsqueeze(-1)
+            # Take left padding into account, which has 0s in left and max_len
+            # in right.
+            left_pad_mask = position_ids == 0
+            # position_ids = paddle.where(
+            #     left_pad_mask, position_ids, position_ids + left_pad_mask.sum(-1, keepdim=True) - 1
+            # )
+            # the above limits right padding must not be 0s, the following suits
+            # to both left and right padding with 0s
+            left_pad_num = (
+                paddle.where(left_pad_mask, position_ids.shape[-1] + 100, position_ids).argmin(axis=-1, keepdim=True)
+                - 1
+            )
+            position_ids = left_pad_num + position_ids
+            second_pos = paddle.max(position_ids, axis=-1, keepdim=True)
+            end_pos = paddle.stack([first_pos, second_pos], axis=-1).squeeze(1)
+            end_score = scores.gather_nd(end_pos)
+        else:
+            # attention_mask passed from pipeline pre-stage is shaped (bs, 1, seq_len, seq_len)
+            assert attention_mask is not None and len(attention_mask.shape) == 2
+            end_score = []
+            end_pos = []
+            for i in range(hidden_state.shape[0]):
+                end_index = attention_mask[i].nonzero()[-1].item()
+                end_pos.append((i, end_index))
+                end_score.append(scores[i, end_index])  # size = (D,)
+            end_score = paddle.stack(end_score, axis=0)  # size = (B, D)
+
+        if self.training and self.do_normalize:
+
+            if dist.is_initialized():
+                gathered_end_score_list = []
+                try:
+                    # gather among data parallel group
+                    hcg = dist.fleet.get_hybrid_communicate_group()
+                    group = hcg.get_sharding_parallel_group()
+                    dist.all_gather(gathered_end_score_list, end_score, group)
+                except:
+                    dist.all_gather(gathered_end_score_list, end_score)
+                gathered_end_score = paddle.concat(gathered_end_score_list, axis=0)
+                self.normalizer.update(gathered_end_score)
+            else:
+                self.normalizer.update(end_score)
+            self.config.mean = self.normalizer.mean.tolist()
+            self.config.var = self.normalizer.var.tolist()
+
+        if self.do_normalize:
+            scores = self.normalizer.normalize(scores)
+            end_score = self.normalizer.normalize(end_score)
+
+        if not return_dict:
+            return scores, end_score
+
+        return ScoreModelOutput(
+            scores=scores,  # size = (B, L, D)
+            end_scores=end_score,  # size = (B, D)
+        )
+
+    def set_normalize(self, mode: bool = True) -> None:
+        if self.do_normalize == mode:
+            return
+
+        self.do_normalize = self.config.do_normalize = mode
+
+
+NormalizeFunction = Literal["affine", "scale", "translate", "identity"]
+NormalizerType = Literal["RunningMeanStd", "ExponentialMovingAverage"]
+
+
+class Normalizer(nn.Layer):
+    """Normalize input to have zero mean and unit variance."""
+
+    mean: paddle.Tensor
+    var: paddle.Tensor
+    count: paddle.Tensor
+    normalize_function: NormalizeFunction
+
+    def __init__(
+        self,
+        normalize_function: NormalizeFunction,
+        shape: tuple[int, ...],
+        device: str | None = None,
+    ) -> None:
+        """Initialize."""
+        super().__init__()
+        if normalize_function not in {"affine", "scale", "translate", "identity"}:
+            raise ValueError(
+                f"Invalid normalization function type: {normalize_function}. ",
+                'Expected one of "affine", "scale", "translate", "identity".',
+            )
+        self.normalize_function = normalize_function
+        self.register_buffer("mean", paddle.zeros(shape, dtype=paddle.get_default_dtype()))  # align torch zeros/ones
+        self.register_buffer("var", paddle.ones(shape, dtype=paddle.get_default_dtype()))  # align torch zeros/ones
+        self.register_buffer("count", paddle.zeros(1, dtype=paddle.int64))
+
+    @abstractmethod
+    def update(self, data: paddle.Tensor) -> None:
+        """Update mean and variance."""
+        raise NotImplementedError
+
+    @property
+    def std(self) -> paddle.Tensor:
+        """Return standard deviation."""
+        return self.var.sqrt()
+
+    def set_mean_var(
+        self,
+        mean: paddle.Tensor | list[float] | tuple[float, ...] | None,
+        var: paddle.Tensor | list[float] | tuple[float, ...] | None,
+    ) -> None:
+        """Set mean and variance."""
+        mean = paddle.to_tensor(mean, dtype=self.mean.dtype, place=self.mean.place) if mean is not None else self.mean
+        var = paddle.to_tensor(var, dtype=self.var.dtype, place=self.var.place) if var is not None else self.var
+
+        assert mean.shape == self.mean.shape
+        assert var.shape == self.var.shape
+
+        self.mean = mean
+        self.var = var
+
+    def forward(
+        self,
+        data: paddle.Tensor,
+        epsilon=1e-8,
+    ) -> paddle.Tensor:
+        """Update and normalize input."""
+        if self.training:
+            self.update(data)
+        return self.normalize(data, epsilon=epsilon)
+
+    def normalize(
+        self,
+        data: paddle.Tensor,
+        epsilon=1e-8,
+    ) -> paddle.Tensor:
+        """Normalize input."""
+        if self.normalize_function == "affine":
+            return (data - self.mean.detach()) / (self.std.detach() + epsilon)
+        if self.normalize_function == "scale":
+            return data / (self.std.detach() + epsilon)
+        if self.normalize_function == "translate":
+            return data - self.mean.detach()
+        if self.normalize_function == "identity":
+            return data
+        raise ValueError(
+            f"Invalid normalization function type: {self.normalize_function}. ",
+            'Expected one of "affine", "scale", "translate", "identity".',
+        )
+
+    @classmethod
+    def instantiate(
+        cls,
+        normalizer_type: NormalizerType | None,
+        normalize_function: NormalizeFunction,
+        shape: tuple[int, ...],
+        device: str | None = None,
+        **kwargs: Any,
+    ) -> Normalizer:
+        """Get a normalizer."""
+        if normalizer_type == "RunningMeanStd":
+            return RunningMeanStd(
+                normalize_function,
+                shape=shape,
+                device=device,
+            )
+        if normalizer_type == "ExponentialMovingAverage":
+            return ExponentialMovingAverage(
+                normalize_function,
+                shape=shape,
+                device=device,
+                **kwargs,
+            )
+        if normalizer_type is None:
+            return IdentityNormalizer(
+                normalize_function,
+                shape=shape,
+                device=device,
+            )
+        raise ValueError(
+            f"Invalid normalization function type: {normalizer_type}. "
+            'Expected one of "RunningMeanStd", "ExponentialMovingAverage".',
+        )
+
+
+class RunningMeanStd(Normalizer):
+    """Running mean and standard deviation."""
+
+    def update(self, data: paddle.Tensor) -> None:
+        """Update mean and variance."""
+        batch_mean = data.mean(0)
+        batch_var = data.var(0)
+        batch_count = data.shape[0]
+
+        delta = batch_mean - self.mean
+        total_count = self.count + batch_count
+
+        new_mean = self.mean + delta * batch_count / total_count
+        m_a = self.var * self.count
+        m_b = batch_var * batch_count
+        m2 = (  # pylint: disable=invalid-name
+            m_a + m_b + paddle.square(delta) * (self.count * batch_count / total_count)
+        )
+        new_var = m2 / total_count
+
+        self.mean = new_mean
+        self.var = new_var
+        self.count = total_count
+
+
+class ExponentialMovingAverage(Normalizer):
+    """Exponential moving average."""
+
+    def __init__(
+        self,
+        normalize_function: NormalizeFunction,
+        shape: tuple[int, ...],
+        device: str | None = None,
+        momentum: float = 0.9,
+    ) -> None:
+        super().__init__(normalize_function, shape=shape, device=device)
+        self.momentum = momentum
+
+    def update(self, data: paddle.Tensor) -> None:
+        """Update mean and variance."""
+        batch_mean = data.mean(0)
+        batch_var = data.var(0)
+        batch_count = data.shape[0]
+
+        self.mean = self.momentum * self.mean + (1.0 - self.momentum) * batch_mean
+        self.var = self.momentum * self.var + (1.0 - self.momentum) * batch_var
+        self.count += batch_count  # pylint: disable=no-member
+
+
+class IdentityNormalizer(Normalizer):
+    """Identity normalizer."""
+
+    def update(self, data: paddle.Tensor) -> None:
+        """Update mean and variance."""
+        self.count += data.shape[0]  # pylint: disable=no-member

llm/Alignment/PPO/ppo_main.py

@@ -0,0 +1,511 @@
+# 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.
+
+import copy
+import os
+import sys
+import types
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Any, Dict, Tuple
+
+import paddle
+from data import PromptOnlyDataset, SupervisedDataset, parse_dataset
+from models import AutoModelForScore
+from models.score_model import LlamaModelForScore  # noqa
+from ppo_trainer import PPOTrainer, cleanup_tensor_space, offload_tensor_to_cpu
+
+from paddlenlp.trainer import PdArgumentParser, TrainingArguments, get_last_checkpoint
+from paddlenlp.transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    LlamaTokenizer,
+)
+from paddlenlp.utils.log import logger
+
+
+@dataclass
+class TrainingArguments(TrainingArguments):
+    kl_coeff: float = field(
+        default=0.02,
+        metadata={"help": "The coefficient for the KL divergence between the reference and actor policy."},
+    )
+    clip_range_ratio: float = field(
+        default=0.2,
+        metadata={
+            "help": "The clipping range for ratio between the old and new policy. This is the epsilon parameter in the PPO algorithm."
+        },
+    )
+    clip_range_score: float = field(
+        default=50.0,
+        metadata={
+            "help": "The clipping range for the output of the score model. The reward is clipped into [-clip_range_score, clip_range_score]."
+        },
+    )
+    clip_range_value: float = field(
+        default=5.0,
+        metadata={
+            "help": "The clipping range for the value function. The value is clipped into [value_estimate - clip_range_value, value_estimate + clip_range_value] during training."
+        },
+    )
+    ptx_coeff: float = field(
+        default=0.0,
+        metadata={"help": "The coefficient for the ptx loss."},
+    )
+    update_iters: int = field(
+        default=1,
+        metadata={"help": "The number of repeated updates on a generated batch."},
+    )
+    critic_learning_rate: float = field(
+        default=None,
+        metadata={"help": "Initial learning rate (after the potential warmup period) for the critic model training."},
+    )
+    critic_weight_decay: float = field(
+        default=None,
+        metadata={"help": "Weight decay to for the critic model training."},
+    )
+    critic_lr_scheduler_type: str = field(
+        default=None,
+        metadata={"help": "The scheduler type for critic model."},
+    )
+    critic_warmup_ratio: float = field(
+        default=None,
+        metadata={"help": "Ratio of warm steps over total training steps for the critic lr scheduler."},
+    )
+    critic_recompute: bool = field(
+        default=None,
+        metadata={"help": "Enable gradient checkpointing for critic model."},
+    )
+    normalize_reward: bool = field(
+        default=None,
+        metadata={"help": "Whether to normalize the reward during RL training."},
+    )
+    temperature: float = field(
+        default=1.0,
+        metadata={"help": "The value used to module the next token probabilities."},
+    )
+    # top_k: int = field(
+    #     default=1,
+    #     metadata={"help": "top_k"},
+    # )
+    top_p: float = field(
+        default=0.8,
+        metadata={
+            "help": "If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to`top_p` or higher are kept for generation."
+        },
+    )
+    num_return_sequences: int = field(
+        default=1,
+        metadata={"help": "The number of independently computed returned sequences for each element in the batch."},
+    )
+    repetition_penalty: float = field(
+        default=1.0,
+        metadata={"help": "The parameter for repetition penalty. 1.0 means no penalty."},
+    )
+    per_device_prompt_batch_size: int = field(
+        default=16,
+        metadata={"help": "Batch size (per device) for the training dataloader."},
+    )
+    eval_mode: str = field(
+        default=None,
+        metadata={
+            "help": "eval mode for actor model and reward_critic_model, optional for: None, single, tensor_parallel."
+        },
+    )
+
+    offload_level: str = field(
+        default="",
+        metadata={"help": "Offload model, optional for: eval, reward, optimizer, train_model"},
+    )
+    use_fusemt: bool = field(
+        default=True,
+        metadata={"help": "use inference model to speedup in rollout generation"},
+    )
+
+    # save_generation_output: bool = field(
+    #     default=False,
+    #     metadata={"help": "Whether to save generated text to file when eval"},
+    # )
+
+
+@dataclass
+class ModelArgument:
+    actor_model_name_or_path: str = field(
+        default=None, metadata={"help": "Build-in pretrained model name or the path to local model."}
+    )
+    reward_model_name_or_path: str = field(
+        default=None, metadata={"help": "Build-in pretrained model name or the path to local model."}
+    )
+    reward_critic_model_name_or_path: str = field(
+        default=None, metadata={"help": "Build-in pretrained model name or the path to local model."}
+    )
+    use_flash_attention: bool = field(default=False, metadata={"help": "Whether to use flash attention"})
+
+    # # LoRA related parameters
+    # lora: bool = field(default=False, metadata={"help": "Whether to use LoRA technique"})
+    # lora_path: str = field(default=None, metadata={"help": "Initialize lora state dict."})
+    # lora_rank: int = field(default=8, metadata={"help": "Lora attention dimension"})
+
+    # # prefix tuning related parameters
+    # prefix_tuning: bool = field(default=False, metadata={"help": "Whether to use Prefix technique"})
+    # num_prefix_tokens: int = field(default=128, metadata={"help": "Number of prefix tokens"})
+
+
+@dataclass
+class DataArgument:
+    train_datasets: str = field(default=None, metadata={"help": "Dataset name(s) registered in the raw dataset."})
+    eval_datasets: str = field(default=None, metadata={"help": "Dataset name(s) registered in the raw dataset."})
+    eval_split_ratio: float = field(default=None, metadata={"help": "Ratio of eval data to train data"})
+    ptx_datasets: str = field(default=None, metadata={"help": "Dataset name(s) registered in the raw dataset."})
+    max_length: int = field(
+        default=2048,
+        metadata={
+            "help": "The maximum length that model input tokens can have. When intokens is set to True, it's also the maximum length for InTokens data stream"
+        },
+    )
+
+    @property
+    def parsed_train_datasets(self) -> Tuple[str, Dict[str, Any]]:
+        """Parse dataset path and its proportion and optionally additional arguments from `train_datasets`."""
+        return [parse_dataset(string) for string in self.train_datasets.split(",")]
+
+    @property
+    def parsed_eval_datasets(self) -> Tuple[str, Dict[str, Any]]:
+        """Parse dataset path and its proportion and optionally additional arguments from `eval_datasets`."""
+        if self.eval_datasets is None:
+            return None
+        return [parse_dataset(string) for string in self.eval_datasets.split(",")]
+
+    @property
+    def parsed_ptx_datasets(self) -> Tuple[str, Dict[str, Any]]:
+        """Parse dataset path and its proportion and optionally additional arguments from `ptx_datasets`."""
+        if self.ptx_datasets is None:
+            return None
+        return [parse_dataset(string) for string in self.ptx_datasets.split(",")]
+
+
+def main():
+    # Arguments
+    parser = PdArgumentParser((ModelArgument, DataArgument, TrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    training_args.print_config(model_args, "Model")
+    training_args.print_config(data_args, "Data")
+    if training_args.eval_mode is not None and len(training_args.eval_mode) == 0:
+        training_args.eval_mode = None
+    if training_args.eval_mode is None and training_args.offload_level is not None:
+        training_args.offload_level = training_args.offload_level.replace("eval", "")
+
+    # Setup GPU & distributed training
+    paddle.set_device(training_args.device)
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, world_size: {training_args.world_size}, "
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16 or training_args.bf16}"
+    )
+
+    # Detecting last checkpoint.
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 1:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Load model
+    if training_args.fp16_opt_level == "O2":
+        if training_args.fp16:
+            dtype = "float16"
+        elif training_args.bf16:
+            dtype = "bfloat16"
+        else:
+            raise ValueError("Please specific dtype: --fp16 or --bf16")
+    else:
+        dtype = "float32"
+    training_args.max_length = data_args.max_length
+
+    model_class_lm, model_class_score = AutoModelForCausalLM, AutoModelForScore
+    if training_args.pipeline_parallel_degree > 1:
+        from models.model_pp import LlamaPolicyPipe, LlamaValuePipe
+
+        model_class_lm = LlamaPolicyPipe
+        model_class_score = LlamaValuePipe
+        extra_args = {
+            "ptx_coeff": training_args.ptx_coeff,
+            "clip_range_ratio": training_args.clip_range_ratio,
+        }
+    else:
+        # non-pipe modelForCausalLM does not accept extra_args and use other ways
+        # (StepTrainer.create_criterion) to set hyper-parameters
+        extra_args = {}
+
+    # actor model
+    model_config = AutoConfig.from_pretrained(
+        model_args.actor_model_name_or_path,
+        tensor_parallel_output=False,
+        tensor_parallel_degree=training_args.tensor_parallel_degree,
+        tensor_parallel_rank=training_args.tensor_parallel_rank,
+        dtype=dtype,
+    )
+    if hasattr(model_config, "use_flash_attention"):
+        model_config.use_flash_attention = model_args.use_flash_attention
+
+    # model_config.num_hidden_layers = 2
+
+    actor_model = model_class_lm.from_pretrained(
+        model_args.actor_model_name_or_path,
+        config=model_config,
+        **extra_args,
+        # ptx_coeff=training_args.ptx_coeff,
+        # clip_range_ratio=training_args.clip_range_ratio,
+    )
+    if training_args.eval_mode is not None:
+        config = copy.deepcopy(actor_model.config)
+        if training_args.eval_mode == "single":
+            config.tensor_parallel_degree = -1
+            config.tensor_parallel_rank = 0
+        actor_eval_model = AutoModelForCausalLM.from_config(config)
+        # TODO(guosheng): AutoModel (in `_get_model_class_from_config`) pop out
+        # architecture which is necessary for infer predictor currently
+        config.architectures = actor_model.config.architectures
+        # actor_eval_model = AutoModelForCausalLM.from_pretrained(model_args.actor_model_name_or_path, config=config)
+    else:
+        actor_eval_model = None
+
+    # todo reference model
+    if training_args.eval_mode is not None:
+        config = copy.deepcopy(model_config)
+        if training_args.eval_mode == "single":
+            config.tensor_parallel_degree = -1
+            config.tensor_parallel_rank = 0
+        actor_reference_model = AutoModelForCausalLM.from_pretrained(
+            model_args.actor_model_name_or_path,
+            config=config,
+        )
+    else:
+        actor_reference_model = model_class_lm.from_pretrained(
+            model_args.actor_model_name_or_path,
+            config=model_config,
+        )
+
+    actor_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.actor_model_name_or_path, model_max_length=data_args.max_length, padding_side="left"
+    )
+
+    # reward model
+    model_config = AutoConfig.from_pretrained(
+        model_args.reward_model_name_or_path,
+        tensor_parallel_output=False,
+        tensor_parallel_degree=training_args.tensor_parallel_degree,
+        tensor_parallel_rank=training_args.tensor_parallel_rank,
+        dtype=dtype,
+    )
+    if hasattr(model_config, "use_flash_attention"):
+        model_config.use_flash_attention = model_args.use_flash_attention
+    # model_config.num_hidden_layers = 2
+    # todo
+    if training_args.eval_mode is not None:
+        config = copy.deepcopy(model_config)
+        if training_args.eval_mode == "single":
+            config.tensor_parallel_degree = -1
+            config.tensor_parallel_rank = 0
+        reward_model = AutoModelForScore.from_pretrained(
+            model_args.reward_model_name_or_path,
+            config=config,
+            score_type="reward",
+            do_normalize=training_args.normalize_reward,
+        )
+    else:
+        reward_model = model_class_score.from_pretrained(
+            model_args.reward_model_name_or_path,
+            config=model_config,
+            score_type="reward",
+            do_normalize=training_args.normalize_reward,
+        )
+    reward_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.reward_model_name_or_path, model_max_length=data_args.max_length, padding_side="right"
+    )
+    # critic model
+    if model_args.reward_critic_model_name_or_path is None:
+        model_args.reward_critic_model_name_or_path = model_args.reward_model_name_or_path
+    reward_critic_model = model_class_score.from_pretrained(
+        model_args.reward_critic_model_name_or_path,
+        config=model_config,
+        score_type="critic",
+        do_normalize=False,
+        clip_range_value=training_args.clip_range_value,
+    )
+    reward_critic_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.reward_critic_model_name_or_path, model_max_length=data_args.max_length, padding_side="left"
+    )
+    if training_args.eval_mode is not None:
+        config = copy.deepcopy(reward_critic_model.config)
+        if training_args.eval_mode == "single":
+            config.tensor_parallel_degree = -1
+            config.tensor_parallel_rank = 0
+        reward_critic_eval_model = AutoModelForScore.from_config(config)
+        # reward_critic_eval_model =  AutoModelForScore.from_pretrained(
+        #     model_args.reward_critic_model_name_or_path,config=model_config
+        # )
+    else:
+        reward_critic_eval_model = None
+
+    #         # actor model
+    #         model_config = AutoConfig.from_pretrained(
+    #             model_args.actor_model_name_or_path,
+    #             tensor_parallel_output=False,
+    #             tensor_parallel_degree=training_args.tensor_parallel_degree,
+    #             tensor_parallel_rank=training_args.tensor_parallel_rank,
+    #             dtype=dtype,
+    #         )
+    #         model_config.num_hidden_layers = 2
+    #         if hasattr(model_config, "use_flash_attention"):
+    #             model_config.use_flash_attention = model_args.use_flash_attention
+    #         actor_model = AutoModelForCausalLM.from_pretrained(
+    #             model_args.actor_model_name_or_path,
+    #             config=model_config,
+    #         )
+    #
+    #         if training_args.eval_mode is not None:
+    #             config = copy.deepcopy(actor_model.config)
+    #             if training_args.eval_mode == "single":
+    #                 config.tensor_parallel_degree = -1
+    #                 config.tensor_parallel_rank = 0
+    #             actor_eval_model = AutoModelForCausalLM.from_config(config)
+    #         else:
+    #             actor_eval_model = None
+    #
+    #         # reference model
+    #         actor_reference_model = AutoModelForCausalLM.from_pretrained(
+    #             model_args.actor_model_name_or_path,
+    #             config=model_config,
+    #         )
+    #         actor_tokenizer = AutoTokenizer.from_pretrained(
+    #             model_args.actor_model_name_or_path, model_max_length=data_args.max_length, padding_side="left"
+    #         )
+    #
+    #         # reward model
+    #         model_config = AutoConfig.from_pretrained(
+    #             model_args.reward_model_name_or_path,
+    #             tensor_parallel_output=False,
+    #             tensor_parallel_degree=training_args.tensor_parallel_degree,
+    #             tensor_parallel_rank=training_args.tensor_parallel_rank,
+    #             dtype=dtype,
+    #         )
+    #         model_config.num_hidden_layers = 2
+    #         if hasattr(model_config, "use_flash_attention"):
+    #             model_config.use_flash_attention = model_args.use_flash_attention
+    #         reward_model = AutoModelForScore.from_pretrained(
+    #             model_args.reward_model_name_or_path,
+    #             config=model_config,
+    #             score_type="reward",
+    #             do_normalize=training_args.normalize_reward,
+    #         )
+    #         reward_tokenizer = AutoTokenizer.from_pretrained(
+    #             model_args.reward_model_name_or_path, model_max_length=data_args.max_length, padding_side="right"
+    #         )
+    #
+    #         # critic model
+    #         if model_args.reward_critic_model_name_or_path is None:
+    #             model_args.reward_critic_model_name_or_path = model_args.reward_model_name_or_path
+    #         reward_critic_model = AutoModelForScore.from_pretrained(
+    #             model_args.reward_critic_model_name_or_path, config=model_config, score_type="critic", do_normalize=False
+    #         )
+    #         reward_critic_tokenizer = AutoTokenizer.from_pretrained(
+    #             model_args.reward_critic_model_name_or_path, model_max_length=data_args.max_length, padding_side="left"
+    #         )
+    #
+    #         if training_args.eval_mode is not None:
+    #             config = copy.deepcopy(reward_critic_model.config)
+    #             if training_args.eval_mode == "single":
+    #                 config.tensor_parallel_degree = -1
+    #                 config.tensor_parallel_rank = 0
+    #             reward_critic_eval_model = AutoModelForScore.from_config(config)
+    #         else:
+    #             reward_critic_eval_model = None
+
+    for tokenizer in [actor_tokenizer, reward_tokenizer, reward_critic_tokenizer]:
+        if isinstance(tokenizer, LlamaTokenizer) and tokenizer.pad_token_id is None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    train_ds = PromptOnlyDataset(data_args.parsed_train_datasets, tokenizer=actor_tokenizer)
+    if data_args.eval_datasets is None and data_args.eval_split_ratio:
+        train_ds, dev_ds = train_ds.split_train_test(split_ratio=data_args.eval_split_ratio)
+    elif data_args.eval_datasets is not None:
+        dev_ds = PromptOnlyDataset(data_args.parsed_eval_datasets, tokenizer=actor_tokenizer)
+    else:
+        dev_ds = None
+
+    ptx_ds = (
+        SupervisedDataset(data_args.parsed_ptx_datasets, tokenizer=actor_tokenizer)
+        if data_args.ptx_datasets is not None
+        else None
+    )
+    if ptx_ds is not None:
+        # PretrainingCriterion requires shifted inputs and labels
+        ptx_ds.get_collator = types.MethodType(partial(ptx_ds.get_collator.__func__, shift=True), ptx_ds)
+
+    # offload
+    # cleanup actor_eval_model, reward_critic_eval_model
+    # offload actor_reference_model reward_model
+
+    if training_args.offload_level is not None:
+        if "eval" in training_args.offload_level:
+            cleanup_tensor_space(actor_eval_model.state_dict())
+            cleanup_tensor_space(reward_critic_eval_model.state_dict())
+        if "reward" in training_args.offload_level:
+            # if pp mode, should lazy offload
+            offload_tensor_to_cpu(actor_reference_model.state_dict())
+            offload_tensor_to_cpu(reward_model.state_dict())
+
+    trainer = PPOTrainer(
+        #  (policy_model, reference_model, reward_model, value_model)
+        #   policy_model, sft_model,       reward_model, value_model
+        #  (policy_model, reference_model, reward_model, value_model,
+        #  (policy_model, reference_model, reward_model, value_model, policy_eval_model, value_eval_model
+        #  (actor_model, actor_reference_model, reward_model, reward_critic_model, actor_eval_model, reward_critic_eval_model
+        model=(
+            actor_model,
+            actor_reference_model,
+            reward_model,
+            reward_critic_model,
+            actor_eval_model,
+            reward_critic_eval_model,
+        ),
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=dev_ds,
+        ptx_dataset=ptx_ds,
+        tokenizer=(actor_tokenizer, actor_tokenizer, reward_tokenizer, reward_critic_tokenizer),
+        data_collator=train_ds.get_collator(),
+    )
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+    trainer.train(resume_from_checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    main()

llm/Alignment/PPO/tests/run_model.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2024 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.
+
+import copy
+from dataclasses import dataclass, field
+
+import numpy
+import paddle
+from paddle.distributed import fleet
+from ppo_trainer import Trainer, data_group_merge, data_group_split, group_rank_guard
+
+from paddlenlp.trainer import PdArgumentParser, TrainingArguments
+from paddlenlp.transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoModelForCausalLMPipe,
+)
+
+
+@dataclass
+class ModelArgument:
+    model_name_or_path: str = field(
+        default=None, metadata={"help": "Build-in pretrained model name or the path to local model."}
+    )
+    test_mode: str = field(default="export", metadata={"help": "export data_split or rank_guard."})
+
+
+def test_group_rank_guard(group):
+    @group_rank_guard(group=group, rank=0)
+    def func():
+        tensor = paddle.randn([4, 64])
+        return tensor
+
+    t = func()
+    ret = []
+    paddle.distributed.stream.all_gather(ret, t, group=group)
+
+    for x in ret:
+        assert x._md5sum() == t._md5sum(), f"{x} {t}"
+
+
+def main():
+    # Arguments
+    parser = PdArgumentParser((ModelArgument, TrainingArguments))
+    model_args, training_args = parser.parse_args_into_dataclasses()
+
+    hcg = fleet.get_hybrid_communicate_group()
+    pp_group = hcg.get_pipe_parallel_group()
+    tp_group = hcg.get_model_parallel_group()
+
+    if model_args.test_mode == "rank_guard":
+        test_group_rank_guard(tp_group)
+        return 0
+
+    model_config = AutoConfig.from_pretrained(
+        model_args.model_name_or_path,
+        tensor_parallel_output=False,
+        tensor_parallel_degree=training_args.tensor_parallel_degree,
+        tensor_parallel_rank=training_args.tensor_parallel_rank,
+        dtype="float32",
+    )
+
+    model_class = AutoModelForCausalLM
+    if training_args.pipeline_parallel_degree > 1:
+        model_class = AutoModelForCausalLMPipe
+
+    actor_model = model_class.from_pretrained(
+        model_args.model_name_or_path,
+        config=model_config,
+    )
+
+    if True:  # test export_evaluate_model
+        # 随机初始化
+        config = copy.deepcopy(model_config)
+        if training_args.pipeline_parallel_degree <= 1:
+            config.tensor_parallel_degree = -1
+            config.tensor_parallel_rank = 0
+
+        actor_eval_model = AutoModelForCausalLM.from_config(config)
+        # ground truth模型
+        actor_gt_model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, config=config)
+
+        trainer = Trainer(
+            model=actor_model,
+            args=training_args,
+        )
+        trainer.export_evaluate_model(actor_model, actor_eval_model)
+
+        gp_state = actor_gt_model.state_dict()
+        export_state = actor_eval_model.state_dict()
+
+        for k, v in gp_state.items():
+            assert (
+                v._md5sum() == export_state[k]._md5sum()
+            ), f"{k} groud_truth: {v.shape}, export: {export_state[k].shape}"
+
+        split_group = tp_group
+        if training_args.pipeline_parallel_degree > 1:
+            split_group = pp_group
+
+        input_ids = paddle.randint(low=1, high=50, shape=[8, 64])
+        paddle.distributed.broadcast(input_ids, src=0)
+
+        split_input_ids = data_group_split(input_ids, group=split_group)
+        ret = actor_eval_model(input_ids=split_input_ids, return_dict=True)
+        eval_loggits = data_group_merge(ret.logits, group=split_group)
+
+        gt_ret = actor_gt_model(input_ids=input_ids, return_dict=True)
+        gt_loggits = gt_ret.logits
+        numpy.testing.assert_almost_equal(eval_loggits.numpy(), gt_loggits.numpy(), decimal=5)
+
+
+if __name__ == "__main__":
+    main()

llm/Alignment/PPO/tests/test_export.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2024 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.
+
+import os
+
+from tests.parallel_launch import TestMultipleGpus
+
+environment_variables = {
+    "NCCL_ALGO": "Tree",
+    "NVIDIA_TF32_OVERRIDE": "0",
+    "NCCL_IB_TIMEOUT": "22",
+    "NCCL_DEBUG": "INFO",
+    "FLAGS_embedding_deterministic": "1",
+    "FLAGS_cudnn_deterministic": "1",
+    "Flags_mp_aysnc_allreduce": "1",
+    "Flags_skip_mp_c_identity": "1",
+    "FLAGS_shard_norm_align_dp": "0",
+    "FLAGS_shard_use_reduce": "1",
+    "test_ci_no_save_model": "1",
+}
+
+
+class TestExportEvalModel(TestMultipleGpus):
+    def setUp(self):
+        os.environ.update(environment_variables)
+        super().setUp()
+
+    def test_pptp_to_tp(self):
+        config = {
+            "output_dir": "./tmp",
+            "model_name_or_path": "__internal_testing__/tiny-random-llama",
+            "tensor_parallel_degree": 2,
+            "pipeline_parallel_degree": 2,
+        }
+        scripts = "tests/run_model.py"
+        self.run_4gpu(scripts, **config)
+
+    def test_tp_to_single(self):
+        config = {
+            "output_dir": "./tmp",
+            "model_name_or_path": "__internal_testing__/tiny-random-llama",
+            "tensor_parallel_degree": 2,
+            "pipeline_parallel_degree": 1,
+        }
+        scripts = "tests/run_model.py"
+        self.run_2gpu(scripts, **config)
+
+    def test_group_rank_guard(self):
+        config = {
+            "output_dir": "./tmp",
+            "model_name_or_path": "__internal_testing__/tiny-random-llama",
+            "tensor_parallel_degree": 2,
+            "pipeline_parallel_degree": 1,
+            "test_mode": "rank_guard",
+        }
+        scripts = "tests/run_model.py"
+        self.run_2gpu(scripts, **config)

llm/Alignment/PPO/trainer_utils.py

@@ -0,0 +1,662 @@
+# 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 inspect
+import os
+import time
+from contextlib import contextmanager
+from typing import Dict
+
+import numpy as np
+import paddle
+import tqdm
+from paddle.distributed import fleet
+from paddle.io import DataLoader
+
+from paddlenlp.generation.utils import GenerationMixin
+from paddlenlp.trainer.trainer import (
+    TRAINER_STATE_NAME,
+    HybridParallelOptimizer,
+    NlpDistributedBatchSampler,
+    ShardingOption,
+    Trainer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+    _obtain_optimizer_parameters_list,
+    distributed_file,
+    distributed_isfile,
+    fused_allreduce_gradients,
+    logger,
+    reshard_util,
+    split_inputs_sequence_dim,
+)
+from paddlenlp.transformers import BatchEncoding, PretrainedModel, PretrainedTokenizer
+from paddlenlp.transformers.configuration_utils import PretrainedConfig
+from paddlenlp.transformers.model_outputs import ModelOutput
+from paddlenlp.transformers.tokenizer_utils_base import (
+    PaddingStrategy,
+    TruncationStrategy,
+)
+
+
+# ########## patches for Trianer ##########
+def init_train_model_opt(
+    self: Trainer, max_steps: int, resume_from_checkpoint: bool = False, clear_master_weight: bool = False
+) -> PretrainedModel:
+    # Copy of model/optimizer init and resuming related code in `Trainer.train`.
+    # NOTE: this `_load_from_checkpoint` is indeed to load model states in the
+    # following elif-else branches, though they are apart away in `Trainer.train`.
+    if not self.args.should_load_sharding_stage1_model:
+        self._load_from_checkpoint(resume_from_checkpoint)
+
+    # delay_optimizer_creation = (
+    #     self.sharding is not None
+    #     and ShardingOption.SHARD_OP in self.args.sharding
+    # )
+    delay_optimizer_creation = False
+
+    if not delay_optimizer_creation:
+        self.create_optimizer_and_scheduler(num_training_steps=max_steps)
+
+    if self.args.should_load_sharding_stage1_model:
+        model = self._wrap_model_and_load_sharded_checkpoint(resume_from_checkpoint)
+    elif self.args.should_save_sharding_stage1_model:
+        # In the non-sharded mode, should invoke _load_from_checkpoint before _wrap_model.
+        # In this mode, the rank0 load all params and the _wrap_model implicitly broadcast params from rank0 to the other ranks.
+        model = self._wrap_model(self.model_wrapped)
+        if self.sharding_io is not None:
+            assert delay_optimizer_creation is False, "delay_optimizer_creation should be False"
+            # the self.optimizer should be wrapped and it is done in _wrap_model
+            self.sharding_io.set_optimizer(self.optimizer)
+        # for the rest of this function `model` is the outside model, whether it was wrapped or not
+        if model is not self.model:
+            self.model_wrapped = model
+        if delay_optimizer_creation:
+            self.create_optimizer_and_scheduler(num_training_steps=max_steps)
+        self._load_optimizer_and_scheduler(resume_from_checkpoint)
+    else:
+        model = self._wrap_model(self.model_wrapped)
+        # for the rest of this function `model` is the outside model, whether it was wrapped or not
+        if model is not self.model:
+            self.model_wrapped = model
+        if delay_optimizer_creation:
+            self.create_optimizer_and_scheduler(num_training_steps=max_steps)
+        self._load_optimizer_and_scheduler(resume_from_checkpoint)
+
+    if ShardingOption.FULL_SHARD in self.args.sharding and clear_master_weight:
+        # for inference model to use Trainer sharding stage3, clear master_weight
+        # which is created in GroupShardedStage3.__init__
+        self.optimizer._master_weights = None
+
+    if self.args.device == "npu" and self.args.flatten_param_grads:
+        from .plugins.npu_plugin import npu_accelerate_plugin
+
+        npu_accelerate_plugin(self.optimizer)
+
+    return model
+
+
+def init_train_state(
+    self: Trainer,
+    resume_from_checkpoint: bool,
+    train_dataloader: DataLoader,
+    max_steps: int,
+    num_train_epochs: int,
+    num_update_steps_per_epoch: int,
+):
+    args = self.args
+
+    self.state = TrainerState()
+    self.state.epoch = 0
+    epochs_trained = 0
+    steps_trained_in_current_epoch = 0
+    steps_trained_progress_bar = None
+
+    # Check if continuing training from a checkpoint
+    if resume_from_checkpoint is not None and distributed_isfile(
+        os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
+    ):
+        self.state = TrainerState.load_from_json(
+            distributed_file(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
+        )
+        epochs_trained = self.state.global_step // num_update_steps_per_epoch
+        if not args.ignore_data_skip:
+            steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
+            steps_trained_in_current_epoch *= args.gradient_accumulation_steps
+        else:
+            steps_trained_in_current_epoch = 0
+
+        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+        logger.info(f"  Continuing training from epoch {epochs_trained}")
+        logger.info(f"  Continuing training from global step {self.state.global_step}")
+        if not args.ignore_data_skip:
+            logger.info(
+                f"  Will skip the first {epochs_trained} epochs then the first {steps_trained_in_current_epoch} "
+                "batches in the first epoch. If this takes a lot of time, you can add the `--ignore_data_skip` "
+                "flag to your launch command, but you will resume the training on data already seen by your model."
+            )
+            if self.is_local_process_zero() and not args.disable_tqdm:
+                steps_trained_progress_bar = tqdm(total=steps_trained_in_current_epoch)
+                steps_trained_progress_bar.set_description("Skipping the first batches")
+        if not args.ignore_data_skip:
+            if isinstance(train_dataloader, paddle.io.DataLoader) and isinstance(
+                train_dataloader.batch_sampler, NlpDistributedBatchSampler
+            ):
+                consumed_samples = (
+                    self.state.global_step
+                    * args.train_batch_size
+                    * args.gradient_accumulation_steps
+                    * args.dataset_world_size
+                )
+                train_dataloader.batch_sampler.set_epoch(consumed_samples=consumed_samples)
+                logger.info(f"Set DistributedBatchSampler consumed_samples to {consumed_samples}")
+
+    self.state.max_steps = int(max_steps)
+    self.state.num_train_epochs = num_train_epochs
+    self.state.is_local_process_zero = self.is_local_process_zero()
+    self.state.is_world_process_zero = self.is_world_process_zero()
+
+    return epochs_trained, steps_trained_in_current_epoch, steps_trained_progress_bar
+
+
+def init_train_log(
+    self: Trainer,
+    num_examples: int,
+    num_train_epochs: int,
+    total_train_batch_size: int,
+    max_steps: int,
+    num_train_samples: int,
+    model: PretrainedModel,
+):
+    args = self.args
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_examples:,}")
+    logger.info(f"  Num Epochs = {num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {max_steps:,}")
+    logger.info(f"  Total num train samples = {num_train_samples:,}")
+    # per_device_trainable_numel = sum(p.numel().item() for p in model.parameters() if not p.stop_gradient)
+    # TODO: Temporary fix since Tensor.numel() not supported in distributed mode
+    per_device_trainable_numel = sum(np.prod(p.shape) for p in model.parameters() if not p.stop_gradient)
+    logger.info(f"  Number of trainable parameters = {per_device_trainable_numel:,} (per device)")
+    if self.args.use_hybrid_parallel:
+        # todo fix for pipeline_parallel_degree
+        parts_num = max(self.args.tensor_parallel_degree, 1) * max(self.args.pipeline_parallel_degree, 1)
+        if parts_num > 1:
+            all_reduce_dtype = "int64"
+            if paddle.get_device().split(":")[0] in ["npu", "xpu"]:
+                # TODO(duanyanhui): fix when NPU all_reduce supports int64
+                all_reduce_dtype = "float32"
+            trainable_numel_tensor = paddle.to_tensor(per_device_trainable_numel, dtype=all_reduce_dtype)
+            paddle.distributed.all_reduce(trainable_numel_tensor)
+            trainable_numel = int(trainable_numel_tensor.item()) // self.args.dataset_world_size
+            # the numel is roughly, because the tensor parallel still hold own bias or layer_norm weight without splited
+            # so, the trainable numel is a little bigger than real.
+            logger.info(f"  Number of trainable parameters = {trainable_numel:,} (all devices, roughly)")
+
+
+def full_training_step(self: Trainer, inputs: Dict[str, paddle.Tensor], **kwargs):
+    """
+    Just a copy of single training step complete code in Trainer.train while loop
+    which including forward+backward+step, while wraps the inputs and outputs to
+    make the complicated copied code no need to change. Maybe a better way is to
+    add fine-grained methods including these steps to Trainer which is similar to
+    DeepSpeed engine.
+    """
+
+    # TODO(guosheng): step, steps_trained_in_current_epoch and steps_trained_progress_bar
+    # should use reference since they would be overwrite.
+    # for state update
+    epoch = kwargs.get("epoch", 0)
+    step = kwargs.get("step", 0)
+    steps_in_epoch = kwargs.get("steps_in_epoch", 0)
+    step_control = kwargs.get("step_control", 0)
+    # for step and progress update when resuming data
+    train_dataloader = kwargs.get("train_dataloader", None)
+    resume_from_checkpoint = kwargs.get("resume_from_checkpoint", None)
+    steps_trained_in_current_epoch = kwargs.get("steps_trained_in_current_epoch", 0)
+    steps_trained_progress_bar = kwargs.get("steps_trained_progress_bar", None)
+    # for eval output ignore to gather
+    ignore_keys_for_eval = kwargs.get("ignore_keys_for_eval", None)
+    tr_loss = kwargs.get("tr_loss", 0.0)
+    model = kwargs.get("model", self.model_wrapped)
+    # needed in _maybe_log_save_evaluate
+    self._globalstep_last_logged = getattr(self, "_globalstep_last_logged", 0)
+    self._globalstep_last_start_time = getattr(self, "_globalstep_last_start_time", time.time())
+
+    args = self.args
+
+    if self.args.use_hybrid_parallel and self.args.sep_parallel_degree > 1:
+        inputs = split_inputs_sequence_dim(inputs)
+    # self.timers and self.timers("read-data").stop()
+    os.environ["TRAINER_GLOBAL_STEP"] = str(self.state.global_step)
+    self.callback_handler.on_load_data_end(args, self.state, self.control, inputs=inputs)
+
+    # Skip past any already trained steps if resuming training
+    # for paddlenlp.utils.batch_sampler.DistributedBatchSampler
+    # We use consumed_samples to reset the status
+    if isinstance(train_dataloader, paddle.io.DataLoader) and isinstance(
+        train_dataloader.batch_sampler, NlpDistributedBatchSampler
+    ):
+        if step == 0:
+            if steps_trained_progress_bar is not None:
+                steps_trained_progress_bar.update(steps_trained_in_current_epoch)
+                steps_trained_progress_bar.close()
+                steps_trained_progress_bar = None
+            self._load_rng_state(resume_from_checkpoint)
+        step += steps_trained_in_current_epoch
+    elif steps_trained_in_current_epoch > 0:
+        steps_trained_in_current_epoch -= 1
+        if steps_trained_progress_bar is not None:
+            steps_trained_progress_bar.update(1)
+        if steps_trained_in_current_epoch == 0:
+            self._load_rng_state(resume_from_checkpoint)
+        # continue
+        final_local_vars = locals()
+        for k in kwargs.keys():
+            if k in final_local_vars:
+                kwargs[k] = final_local_vars[k]
+        return kwargs
+    elif steps_trained_progress_bar is not None:
+        steps_trained_progress_bar.close()
+        steps_trained_progress_bar = None
+
+    if step_control % args.gradient_accumulation_steps == 0:
+        self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
+        self.timers and self.timers("forward-backward").start()
+
+    dp_enabled = self.args.data_parallel_degree > 1 if self.args.use_hybrid_parallel else args.local_rank != -1
+    forbidden_no_sync = False
+    # stage2 and stage3 should not no_sync, because the is no DDP wrapper and no_sync API
+    # hybrid_parallel (tp or pp or sharding stage 1) should not no_sync
+    if self.args.use_hybrid_parallel:
+        forbidden_no_sync = True
+
+    availiable_no_sync = dp_enabled and not forbidden_no_sync
+
+    is_no_sync = (
+        ((step_control + 1) % args.gradient_accumulation_steps != 0)
+        and availiable_no_sync
+        and args._no_sync_in_gradient_accumulation
+    ) or (args.recompute and availiable_no_sync)
+    # sharding
+    # stage1. the same as ddp
+    # stage2. manualy collect gradient on dp group
+
+    dp_master_grad = self.args.world_size > 1 and self.args.amp_master_grad and not self.args.use_hybrid_parallel
+    if dp_master_grad:
+        is_no_sync = True
+
+    if is_no_sync:
+        # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
+        with model.no_sync():
+            tr_loss_step = self.training_step(model, inputs)
+    else:
+        tr_loss_step = self.training_step(model, inputs)
+
+    tr_loss += tr_loss_step
+
+    if (step_control + 1) % args.gradient_accumulation_steps == 0 or (
+        # last step in epoch but step is always smaller than gradient_accumulation_steps
+        steps_in_epoch <= args.gradient_accumulation_steps
+        and (step + 1) == steps_in_epoch
+    ):
+        if self.args.pipeline_parallel_degree <= 1 and self._enable_delay_scale_loss():
+            tr_loss /= self.args.gradient_accumulation_steps
+
+        self.timers and self.timers("forward-backward").stop()
+        # Maunally collect gradients
+        # Case 1: Use recompute and dp
+        # Case 2: Hack dp with master_grad
+        # Case 3: Pipeline or sharding overlap
+        # local_rank != -1 don't means dp in networks.
+        self.timers and self.timers("all-reduce").start()
+
+        # Case 1: Use recompute and dp / sharding stage1,
+        # manualy collect gradient for dp.
+        if args.recompute and availiable_no_sync:
+            fused_allreduce_gradients(list(model.parameters()), None)
+
+        # Case 2: hack dp with master_grad
+        if dp_master_grad and not (args.recompute and availiable_no_sync):
+            fused_allreduce_gradients(list(model.parameters()), None)
+
+        # Pipeline parallel mode,  handle gradient reduce here to overlap
+        pipeline_parallel_config = (
+            set(args.pipeline_parallel_config.split(" ")) if args.pipeline_parallel_degree > 1 else set()
+        )
+        enable_dp_comm_overlap = "enable_dp_comm_overlap" in pipeline_parallel_config
+        enable_release_grads = "enable_release_grads" in pipeline_parallel_config
+
+        # Case 3: Pipeline parallel mode, overlap with dp
+        if isinstance(self.optimizer, HybridParallelOptimizer) and not self.do_grad_scaling:
+            parameters_list = _obtain_optimizer_parameters_list(self.optimizer._inner_opt)
+
+            if not enable_dp_comm_overlap:
+                if self.optimizer._sharding_enable:
+                    assert reshard_util.is_sharding_opt(self.optimizer)
+                    self.optimizer._inner_opt.reduce_gradients(list(parameters_list), self.optimizer._hcg)
+
+                if self.optimizer._dp_enable or getattr(self.optimizer, "_sep_enable", False):
+                    fused_allreduce_gradients(list(parameters_list), self.optimizer._hcg)
+
+        self.timers and self.timers("all-reduce").stop()
+        self.timers and self.timers("optimizer-step").start()
+
+        if self.args.gradient_accumulation_steps > 1 and self._enable_delay_scale_loss():
+            for p in model._layers.parameters():
+                with paddle.no_grad():
+                    if hasattr(p, "main_grad") and p.main_grad is not None:
+                        assert p.grad is None
+                        p.main_grad.scale_(1.0 / self.args.gradient_accumulation_steps)
+                    elif p.grad is not None:
+                        p.grad.scale_(1.0 / self.args.gradient_accumulation_steps)
+
+        # Optimizer step
+        self.callback_handler.on_optimizer_begin(
+            args, self.state, self.control, scaler=self.scaler if self.do_grad_scaling else None
+        )
+        optimizer_was_run = True
+        if self.do_grad_scaling:
+            scale_before = paddle.assign(self.scaler._scale)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            scale_after = self.scaler._scale
+            # Compatible with paddlepaddle 2.6.0 using typo word.
+            if hasattr(self.scaler, "_cache_founf_inf"):
+                optimizer_was_run = not self.scaler._cache_founf_inf
+            else:
+                optimizer_was_run = not self.scaler._cache_found_inf
+            if not optimizer_was_run:
+                scale_before_value = scale_before.cpu().numpy()
+                scale_after_value = scale_after.cpu().numpy()
+                logger.warning(
+                    f"optimizer not run, scale_before: {scale_before_value[0]}, scale_after: {scale_after_value[0]}"
+                )
+        elif isinstance(self.optimizer, HybridParallelOptimizer):
+            self.optimizer._step(parameters_list)
+        else:
+            self.optimizer.step()
+
+        self.timers and self.timers("optimizer-step").stop()
+
+        if optimizer_was_run:
+            self.lr_scheduler.step()
+
+        if enable_release_grads and args.pipeline_parallel_degree > 1:
+            self.optimizer.clear_grad(set_to_zero=False)
+            for _, buffers in model._chunk_2_comm_buffers.items():
+                for buffer in buffers:
+                    buffer._clear_grad_storage()
+        else:
+            self.optimizer.clear_grad(set_to_zero=False)
+
+        self.callback_handler.on_optimizer_end(
+            args, self.state, self.control, scaler=self.scaler if self.do_grad_scaling else None
+        )
+
+        self.state.global_step += 1
+        self.state.epoch = epoch + (step + 1) / steps_in_epoch
+        self.control = self.callback_handler.on_step_end(args, self.state, self.control)
+        self._maybe_log_save_evaluate(tr_loss, model, epoch, ignore_keys_for_eval, inputs=inputs)
+        self._print_timer()
+        step_control = 0
+    else:
+        self.control = self.callback_handler.on_substep_end(args, self.state, self.control)
+        step_control += 1
+
+    if self.control.should_epoch_stop or self.control.should_training_stop:
+        # break
+        final_local_vars = locals()
+        for k in kwargs.keys():
+            if k in final_local_vars:
+                kwargs[k] = final_local_vars[k]
+        return kwargs
+    # self.timers and self.timers("read-data").start()
+
+    final_local_vars = locals()
+    for k in kwargs.keys():
+        if k in final_local_vars:
+            kwargs[k] = final_local_vars[k]
+    return kwargs
+
+
+Trainer.init_train_model_opt = init_train_model_opt
+Trainer.init_train_log = init_train_log
+Trainer.init_train_state = init_train_state
+Trainer.full_training_step = full_training_step
+# ########## patches for Trianer ##########
+
+
+class MuteDefaultFlowCallback(TrainerCallback):
+    """
+    Add this callback can cencel logging/evaluation/saving by DefaultFlowCallback.
+    Use this when having multi trainer.
+    """
+
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        control.should_save = False
+        control.should_evaluate = False
+        control.should_log = False
+        return control
+
+
+@contextmanager
+def guard_set_args(args, arg_name_values):
+    for k, v in arg_name_values.items():
+        old_value = getattr(args, k, None)
+        setattr(args, k, v)
+        arg_name_values[k] = old_value
+    yield
+    for k, v in arg_name_values.items():
+        old_value = getattr(args, k)
+        setattr(args, k, v)
+        arg_name_values[k] = old_value
+
+
+class PipeEvalModel(GenerationMixin):
+    """
+    Wrapper for PipelineParallel to do evaluate and generate. Currently only
+    support .
+    """
+
+    def __init__(self, trainer: Trainer):
+        eval_model = getattr(trainer, "_inner_eval_model", None)
+        self.model: fleet.model.PipelineParallel = trainer.model_wrapped if eval_model is None else eval_model
+        self.config: PretrainedConfig = trainer.model.config
+        self._is_gen = False
+        self.update_model_kwargs_for_generation = (
+            self.model._layers._non_pipe_model_class.update_model_kwargs_for_generation
+        )
+
+    @property
+    def pp_group(self):
+        return self.model.pp_group
+
+    def eval(self):
+        self.model.eval()
+
+    def train(self):
+        self.model.train()
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.model, name)
+
+    def _broadcast_outputs(self, outputs):
+        # outputs is PipelineParallel.eval_batch which is a list of batches.
+        out = []
+        outputs = (outputs,) if isinstance(outputs, paddle.Tensor) else outputs
+        for tensors in outputs:
+            if not self.model.is_pipeline_last_stage():
+                tensor = tensors if isinstance(tensors, paddle.Tensor) else tensors[0]
+                head_out_meta = (
+                    (self.model._layers.head_out_meta,)
+                    if isinstance(self.model._layers.head_out_meta, paddle.static.InputSpec)
+                    else self.model._layers.head_out_meta
+                )
+                tensors = tuple(
+                    paddle.empty(
+                        shape=[
+                            tensor.shape[i] if (meta.shape[i] is None or meta.shape[i] < 0) else meta.shape[i]
+                            for i in range(len(meta.shape))
+                        ],
+                        dtype=tensor.dtype if meta.dtype is None else meta.dtype,
+                    )
+                    for meta in head_out_meta
+                )
+            else:
+                # Currently use tuple instead of ModelOutput and require the
+                # caller use the return result as tuple.
+                tensors = (
+                    (tensors,)
+                    if isinstance(tensors, paddle.Tensor)
+                    else tensors.to_tuple()
+                    if isinstance(tensors, ModelOutput)
+                    else tensors
+                )
+
+            # use map_structure seems hung
+            for tensor in tensors:
+                paddle.distributed.broadcast(tensor, src=self.model.pp_group.ranks[-1], group=self.model.pp_group)
+            out.append(tensors[0] if len(tensors) == 1 else tensors)
+        return out[0] if len(out) == 1 else out
+
+    def __call__(self, *args, **kwargs):
+        model = self.model
+        assert self.model.training is False
+        if self._is_gen:
+            # inputs by `prepare_inputs_for_generation` is a dict with following keys:
+            # "input_ids", "position_ids", "past_key_values", "use_cache", "attention_mask"
+            # NOTE: 1. cache/past_key_values should be passed across decoding steps
+            # by using as model attr rather than input args to reduce comm overhead.
+            # Also, pipe model defined for training not support this cache input.
+            # 2. ignore use_cache since _check_data_vaild requires tensor if not None.
+            # 3. attention_mask can reuse _prepare_decoder_attention_mask in LlamaEmbeddingPipe.
+            # 4. position_ids pass through _prepare_pipeline_inputs_func and PipeLayer.
+            inputs, labels = model._prepare_pipeline_inputs_func(*args, **kwargs)
+            # currently, set accumulate_steps to 1 to avoid multi-batch eval/gen
+            with guard_set_args(model, {"_compute_loss": False, "accumulate_steps": 1}):
+                outputs = model.eval_batch([inputs, labels], compute_loss=False)
+            # TODO(guosheng): Broadcasted logits are used to get next_scores, remove
+            # it to reduce comm overhead. Also note that we still need broadcast
+            # next_tokens though logits are broadcasted since pp ranks' seeds differs.
+            # Currently, just slice the last token to reduce comm overhead.
+            outputs = [
+                micro_batch_output[:, -1, :].unsqueeze(1)
+                if isinstance(micro_batch_output, paddle.Tensor)
+                else micro_batch_output[0][:, -1, :].unsqueeze(1)
+                for micro_batch_output in outputs
+            ]
+            outputs = self._broadcast_outputs(outputs)
+        else:
+            # use _prepare_pipeline_inputs_func to convert pipeline inputs
+            inputs, labels = model._prepare_pipeline_inputs_func(*args, **kwargs)
+            # NOTE(guosheng): bug seems exist. pp.eval_batch(compute_loss=False)
+            # will set pp._compute_loss to False and would not set it back. Thus
+            # hack here to set it back.
+            with guard_set_args(model, {"_compute_loss": False, "accumulate_steps": 1}):
+                outputs = model.eval_batch([inputs, labels], compute_loss=False)
+            outputs = self._broadcast_outputs(outputs)
+        return outputs
+
+    def generate(self, *args, **kwargs):
+        self._is_gen = True
+        # patch DecoderLayerPipe to use cache, DecoderLayerPipe is subclass of
+        # DecoderLayer, and would call super().forward
+        ori_decoder_layer_forward = self.model._layers._non_pipe_decoder_layer_class.forward
+
+        def decoder_layer_forward(layer_self, *args, **kwargs):
+            kwargs.update({"use_cache": True, "past_key_value": getattr(layer_self, "_cache", None)})
+            outputs = ori_decoder_layer_forward(layer_self, *args, **kwargs)
+            output = outputs[0]
+            layer_self._cache = outputs[1]
+            self._has_cache = True
+            return output
+
+        with guard_set_args(self.model._layers._non_pipe_decoder_layer_class, {"forward": decoder_layer_forward}):
+            outputs = super().generate(*args, **kwargs)
+        self._is_gen = False
+        # clear cache of decoder layers, sublayers is incursive thus suitable
+        # to both 1F1B and interleave
+        for layer in self.model._layers.sublayers():
+            if isinstance(layer, self.model._layers._non_pipe_decoder_layer_class):
+                layer._cache = None
+        self._has_cache = False
+        return outputs
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        arg_bind = inspect.signature(self.model._layers._non_pipe_model_class.prepare_inputs_for_generation).bind(
+            *((self,) + args), **kwargs
+        )
+        arg_bind.apply_defaults()
+        arg_dict = arg_bind.arguments
+        last_arg_name, last_arg_value = arg_dict.popitem()
+        if arg_bind.signature.parameters[last_arg_name].kind == inspect.Parameter.VAR_KEYWORD:
+            arg_dict.update(last_arg_value)
+        else:
+            arg_dict[last_arg_name] = last_arg_value
+        arg_dict.pop("self")
+        past_key_values = arg_dict.get("past_key_values", None)
+        # prepare_inputs_for_generation use past_key_values to discrimate prefill
+        # or decode and slice inputs accordingly.
+        if getattr(self, "_has_cache", False):
+            arg_dict.update({"past_key_values": True})
+        model_inputs = self.model._layers._non_pipe_model_class.prepare_inputs_for_generation(self, **arg_dict)
+        model_inputs.update({"past_key_values": past_key_values})
+        return model_inputs
+
+
+def is_same_tokenizer(
+    tokenizer: PretrainedTokenizer,
+    other_tokenizer: PretrainedTokenizer,
+) -> bool:
+    """Check if two tokenizers are the same."""
+    return tokenizer is other_tokenizer or (
+        tokenizer.__class__ == other_tokenizer.__class__ and tokenizer.get_vocab() == other_tokenizer.get_vocab()
+    )
+
+
+def batch_retokenize(
+    input_ids: paddle.Tensor,
+    src_tokenizer: PretrainedTokenizer,
+    dest_tokenizer: PretrainedTokenizer,
+    *,
+    padding: bool | str | PaddingStrategy = PaddingStrategy.LONGEST,
+    truncation: bool | str | TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
+    skip_special_tokens: bool = True,
+) -> BatchEncoding:
+    """Re-tokenize a batch of input ids from one tokenizer to another."""
+    output = dest_tokenizer(
+        [
+            text + dest_tokenizer.eos_token
+            for text in src_tokenizer.batch_decode(
+                input_ids,
+                skip_special_tokens=skip_special_tokens,
+            )
+        ],
+        padding=padding,
+        truncation=truncation,
+        return_tensors="pd",
+    )
+    return output

llm/Alignment/README.md

@@ -0,0 +1,250 @@
+# Alignment
+
+提供了基于强化学习 PPO 算法对 LLM 进行人类偏好对齐的代码及完整使用示例，支持**3D 分布式并行训练以及 rollout 阶段使用预测优化进行生成加速**。其中 PPO 代码实现细节参考了 [PKU-Alignment/safe-rlhf](https://github.com/PKU-Alignment/safe-rlhf)（PKU Beaver） 中的 PPO 实现，支持reward normalization、pretraining loss等常用的 PPO 稳定训练策略；示例使用 PKU-Alignment/safe-rlhf 提供的部分数据集和模型。后续将持续完善扩展，支持更好效果、更低成本、更高性能、更大规模的 RLHF 能力。
+
+## 快速开始
+
+项目整体组织结构如下：
+
+```
+.
+├── PPO                          # PPO 训练相关目录
+│   ├── comm_utils.py            # 通信相关工具py文件
+│   ├── data                     # 数据集相关目录
+│   │   ├── alpaca.py            # alpaca(raw)数据集py文件
+│   │   ├── base.py              # 数据集基类及工具py文件
+│   │   ├── __init__.py
+│   │   ├── preference.py        # 偏好数据集py文件
+│   │   ├── prompt_only.py       # prompt only数据集py文件
+│   │   ├── safe_rlhf.py         # safe_rlhf(raw)数据集py文件
+│   │   └── supervised.py        # supervised数据集py文件
+│   ├── infer_utils.py           # 生成加速工具py脚本
+│   ├── models                   # 模型相关目录
+│   │   ├── infer_model_utils.py # 预测加速模型补丁及工具py文件
+│   │   ├── __init__.py
+│   │   ├── model_pp.py          # 流水线并行模型py文件
+│   │   ├── pp_model_utils.py    # 流水线并行补丁及工具py文件
+│   │   ├── ppo_model.py         # PPO 相关模型实现
+│   │   ├── ppo_model_utils.py   # PPO loss等模型策略py文件
+│   │   ├── score_model.py       # score model模型定义py文件
+│   │   └── score_model_utils.py # score model基类及工具py文件
+│   ├── ppo_main.py              # RLHF训练脚本
+│   ├── ppo_trainer.py           # RLHF训练执行器py脚本
+│   ├── tests                    # 测试相关目录
+│   │   ├── run_model.py
+│   │   └── test_export.py
+│   └── trainer_utils.py         # Trainer补丁及工具py脚本
+├── README.md
+└── RM                           # Reward Model 训练相关目录
+    ├── models -> ../PPO/models
+    ├── reward_main.py            # reward model训练脚本
+    └── reward_trainer.py         # reward训练执行器py脚本
+```
+
+### 环境准备
+
+- Python >= 3.10
+- PaddlePaddle >= 2.6.0
+- PaddleNLP 最新版本
+
+如需使用生成加速功能，需要安装 [paddlenlp_ops](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/csrc) ，请使用 `git clone https://github.com/PaddlePaddle/PaddleNLP.git` 克隆 PaddleNLP 代码库并且将 PaddleNLP/llm 目录的路径加入 PYTHONPATH（后续将进行完善）。安装 paddlenlp_ops 后训练时将直接开启生成加速（开启流水线并行时不支持生成加速），否则使用原生动态图进行生成。
+
+### 数据准备
+
+PPO 训练包括 Supervised Fine-Tuning、Reward Model Fine-Tuning、RLHF 三个阶段（可见下文训练部分），会涉及到多个数据集，下面分别介绍并给出自定义数据的方法。
+
+#### Supervised Fine-Tuning 数据
+同[LLM 精调](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm#2-%E7%B2%BE%E8%B0%83)，可以直接参考对应内容进行数据准备。
+
+#### Reward Model Fine-Tuning 数据
+Reward Model Fine-Tuning 阶段需要使用人类偏好数据。示例使用 PKU-Alignment/safe-rlhf 提供的 [PKU-Alignment/PKU-SafeRLHF-30K](https://huggingface.co/datasets/PKU-Alignment/PKU-SafeRLHF-30K) 数据集，下面是其中一条样本，这里使用其中的`prompt、response_0、response_1、better_response_id`字段来组织偏好数据（safe字段，该数据集将helpful 和 harmless 分开标注，示例这里使用其 helpful 标注）。
+
+```json
+{
+    "prompt": "Is it possible to purchase a complete chess set at a pawn shop?",
+    "response_0": "Yes, it is possible to purchase a complete chess set at a pawn shop.",
+    "response_1": "Yes, it is possible to purchase a complete chess set at a pawn shop. Pawn shops may carry chess sets that have been either inherited or previously borrowed. There is also a possibility of finding individual chess pieces such as rooks, knights, bishops, queens, and kings. If you cannot find a complete set, you may be able to purchase the various chess pieces separately.",
+    "is_response_0_safe": true,
+    "is_response_1_safe": true,
+    "better_response_id": 1,
+    "safer_response_id": 1
+}
+```
+
+#### RLHF 数据
+RLHF 阶段使用 prompt only 数据，另外可以可选的提供额外的监督数据用于构建 LM 损失约束 RLHF 训练。示例使用 [PKU-Alignment/PKU-SafeRLHF](https://huggingface.co/datasets/PKU-Alignment/PKU-SafeRLHF) 数据集（同样是人类偏好数据集，这里只使用其 prompt 字段并对 prompt 去重）。此外还使用了 [tatsu-lab/alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca)的数据来构建额外的损失项。
+
+以上示例数据集在训练时将自动下载缓存使用。
+
+#### 自定义数据
+数据定义围绕 `RawSample` 和 `RawDataset` 两个预置的类进行；`RawSample` 提供了数据样本级别接入协议规范，`RawDataset` 提供了数据集级别接入协议规范；按照下面介绍的规范接入，即可通过预置的 `SupervisedDataset`、`PreferenceDataset`、`PromptOnlyDataset` 三类 RLHF 训练所需数据接口来使用自定义数据。
+
+自定义数据集需要：
+- 继承 `RawDataset` ，并定义类属性 `NAME` 用于注册数据集。
+- 实现 `__init__` 方法（加载数据），`__getitem__` 方法（根据 index 获取样本并转换为 `RawSample` 对象返回）、`__len__` 方法（数据集大小）。
+
+示例如下：
+
+```python
+from datasets import load_dataset
+from data import RawDataset, RawSample
+
+class MyRawDataset(RawDataset):
+    NAME = 'my-dataset-name'
+
+    def __init__(self, path=None) -> None:
+        # Load a dataset from Hugging Face or any other data source
+        # self.data = load_dataset(path or 'my-organization/my-dataset')['train']
+        self.data = [{
+            'col1': 'question',
+            'col2': 'answer1',
+            'col3': 'answer2',
+            'col4': 1,  # score of answer1
+            'col5': 2  # score of answer2
+        }] * 10  # dummy data for example
+
+    def __getitem__(self, index: int) -> RawSample:
+        data = self.data[index]
+        # Construct a `RawSample` dictionary from your custom dataset item
+        return RawSample(
+            input=data['col1'],
+            answer=data['col2'],
+            other_answer=data['col3'],
+            better=float(data['col4']) > float(data['col5']),
+        )
+
+    def __len__(self) -> int:
+        return len(self.data)  # dataset size
+```
+
+其中 `RawSample` 是整个 RLHF 训练过程用到的几种数据类型的超集，如下所示，其可以桥接各训练阶段所需样本类型。在自定义数据时，对于 SFT 数据使用 `RawSample` 的 `(input, answer)` 字段；对于人类偏好数据使用 `RawSample` 的 `(input, answer, other_answer, better)` 字段；对于 prompt only 数据，使用 `RawSample` 的 `(input)`字段。
+
+```python
+class RawSample(TypedDict, total=False):
+    """Raw sample type.
+
+    For SupervisedDataset, should provide (input, answer) or (dialogue).
+    For PreferenceDataset, should provide (input, answer, other_answer, better).
+    For SafetyPreferenceDataset, should provide (input, answer, other_answer, safer, is_safe, is_other_safe).
+    For PromptOnlyDataset, should provide (input).
+
+    When input is a list, it would be processed as a dialogue.
+    """
+
+    # Texts
+    input: NotRequired[str | list[str]]  # either `input` or `dialogue` should be provided
+    """User input text."""
+    answer: NotRequired[str]
+    """Assistant answer text."""
+    other_answer: NotRequired[str]
+    """Other assistant answer text via resampling."""
+    dialogue: NotRequired[list[str]]  # either `input` or `dialogue` should be provided
+    """Dialogue history."""
+
+    # Flags
+    better: NotRequired[bool]
+    """Whether ``answer`` is better than ``other_answer``."""
+    safer: NotRequired[bool]
+    """Whether ``answer`` is safer than ``other_answer``."""
+    is_safe: NotRequired[bool]
+    """Whether ``answer`` is safe."""
+    is_other_safe: NotRequired[bool]
+    """Whether ``other_answer`` is safe."""
+```
+
+如此定义的数据集将可以通过预置接口根据 `NAME` 来使用，当前内置支持`"PKU-SafeRLHF/train", "PKU-SafeRLHF/test", "PKU-SafeRLHF-30K/train", "PKU-SafeRLHF-30K/test", "PKU-SafeRLHF-10K/train", "alpaca"` 几个数据集。另外还支持使用多个数据集并指定数据比例，我们可以按照需要为每个阶段训练准备多份数据集。示例如下：
+
+```python
+from paddlenlp.transformers import AutoTokenizer
+from data import PreferenceDataset
+
+tokenizer = AutoTokenizer.from_pretrained('facebook/llama-7b')
+dataset = PreferenceDataset({
+    'alpaca': 0.75,
+    'my-dataset-name': 0.5
+}, tokenizer)
+```
+
+### 训练
+
+PPO 完整的训练过程包括以下 3 个阶段，如下图所示（来自[DeepSpeed-Chat](https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat)）：
+
+<p align="center">
+  <img src="https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/assets/image/ppo_trainer.png?raw=true" align="middle" width = "600" />
+</p>
+
+1. Supervised Fine-Tuning (SFT)
+
+同[LLM 精调](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm#2-%E7%B2%BE%E8%B0%83)，可以直接参考对应内容进行训练并使用其产出模型。
+
+2. Reward Model Fine-Tuning
+
+使用 `reward_main.py` 脚本根据 `rm.json` 训练奖励模型
+
+```
+cd RM
+python -u -m paddle.distributed.launch reward_main.py ../../config/llama/rm.json
+```
+
+`rm.json` 中的绝大部分参数释义同[LLM 精调](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm#2-%E7%B2%BE%E8%B0%83)，不再赘述；稍有区别的是 `train_datasets`/`eval_datasets` 分别使用数据集定义注册时的`NAME`属性给出训练和验证集。另外对于奖励模型训练有以下特殊参数配置及释义（使用 PKU-Alignment/PKU-SafeRLHF 中的默认值）：
+
+- `normalize_score_during_training`：是否在训练过程中对奖励进行 normalize，默认为 `False`。
+- `normalizer_type`：使用 normalizer 时计算 mean、var 的方式，可选`"RunningMeanStd", "ExponentialMovingAverage"`。
+- `normalizer_momentum`：使用 `ExponentialMovingAverage` normalizer 时指定的 momentum ，默认为 `0.9`。
+- `loss_type`：使用 token 级或是 sequence 级 loss 进行奖励模型训练，可选`"token-wise", "sequence-wise"`，默认为 `"sequence-wise"`。
+- `regularization`：奖励模型训练目标中对奖励的正则化系数，默认为 `0.001`。
+
+3. RLHF：
+
+RLHF 阶段需要 actor model、reference model、critic model、reward model 四个模型；actor-model/reference-model 使用 SFT 模型进行 initialize/frozen；critic-model/reward-model 使用 reward 模型进行 initialize/frozen (另外注意若 SFT 使用 LoRA 请先将 LoRA 权重合并）。这里使用 PKU-Alignment/PKU-SafeRLHF 提供的 SFT 模型（[PKU-Alignment/alpaca-7b-reproduced](https://huggingface.co/PKU-Alignment/alpaca-7b-reproduced)）和 reward 模型（[PKU-Alignment/beaver-7b-v1.0-reward](https://huggingface.co/PKU-Alignment/beaver-7b-v1.0-reward)，注意该模型只关注 helpful 未考量 harmless）作为示例，使用 `ppo_main.py` 脚本根据 `ppo.json` 进行 RLHF 训练。
+
+```
+# 类型提升 warning 暂时通过 loglevel 屏蔽，待后续修复
+cd PPO
+PYTHONPATH=../../ GLOG_minloglevel=2 python -u -m paddle.distributed.launch ppo_main.py ../../config/llama/ppo.json
+```
+
+`ppo.json` 中的绝大部分参数释义同[LLM 精调](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm#2-%E7%B2%BE%E8%B0%83)，不再赘述，重点给出以下参数配置及释义（使用 PKU-Alignment/PKU-SafeRLHF 中的默认值）：
+
+- `train_datasets`：使用数据集定义注册时的`NAME`属性给出训练集。
+- `eval_datasets`：使用数据集定义注册时的`NAME`属性给出验证集。
+- `ptx_datasets`：使用数据集定义注册时的`NAME`属性给出 ptx-loss 使用的数据集，未提供时将不使用 ptx-loss。
+- `actor_model_name_or_path`：actor-model/reference-model 用来 initialize/frozen 的模型名称或目录。
+- `reward_model_name_or_path`：reward-model 的模型名称或目录。
+- `reward_critic_model_name_or_path`：critic-model 的模型名称或目录，未提供时将使用`reward_model_name_or_path`进行 critic-model 的初始化。
+- `per_device_prompt_batch_size`：训练时 prompt only 数据集读取用于 rollout 生成的批次大小（每张卡）。
+- `per_device_train_batch_size`：根据 prompt 进行生成及训练使用的批次大小（每张卡）。
+- `num_return_sequences`：生成时每个 prompt 生成的回复个数，即 `GenerationConfig.num_return_sequences`，所有回复都将用来训练。
+- `temperature`：生成采样时使用的 `temperature` ，即 `GenerationConfig.temperature`。
+- `top_p`：生成采样时 top-p-filtering 阈值，即 `GenerationConfig.top_p`。
+- `repetition_penalty`：生成采样时长度惩罚系数，即 `GenerationConfig.repetition_penalty`。
+- `update_iters`：一次生成的数据被使用的次数。
+- `kl_coeff`：对 reward 进行 KL-Penalty 的系数。
+- `clip_range_score`：对 reward 进行裁剪的阈值。
+- `clip_range_value`：critic model（value function）对当前sequence的新值与Experience Buffer中旧值的差距超过该范围将进行裁剪。
+- `clip_range_ratio`：将当前sequence的新概率与Experience Buffer中旧概率比值裁剪到`(1-clip_range_ratio, 1+clip_range_ratio)`范围（PPO-Clip）。
+- `ptx_coeff`： 预训练损失项 ptx-loss 的系数。
+
+另外所有 [`TrainingArguments` 支持参数配置](https://paddlenlp.readthedocs.io/zh/latest/trainer.html#trainingarguments)将为 actor-model 和 critic-model 的训练复用（如`sharding_stage`），除单独提供了 `critic_learning_rate/critic_weight_decay/critic_lr_scheduler_type/critic_warmup_ratio/critic_recompute` 这些参数支持为 critic-model 训练单独指定相应配置。actor-model 和 critic-model 的 checkpoints 将分别保存在 `outpt_dir` 所指定目录的 policy 和 value 文件夹下。
+
+此外为了支持更高性、更大规模的 RLHF 训练提供了以下特殊参数配置，可以按需使用：
+- `use_fusemt`：安装 paddlenlp_ops 后将在 rollout 生成时开启生成加速（开启流水线并行时不支持生成加速），通过此设置可以禁用生成加速。
+- `eval_mode`：支持为空或者设置为 "single"、"tensor_parallel"；通常可以在使用流水线并行训练时设置为"tensor_parallel"，以此在 rollout 生成阶段使用非流水线并行模型并进行生成加速。
+- `offload_level`：支持设置为"freeze_model"、"optimizer"、"train_model"或者同时使用(空格分隔），分别指示 reward+reference 两个冻结模型、actor+critic 两个训练模型的优化器状态和模型参数的 offload/reload，用于在不同阶段 model/optimizer 使用结束后及时 offload 并在下次使用时 reload 相应参数权重以节省显存。
+
+另外注意，在使用流水线并行时（pipeline_parallel_degree大于1）建议将 `dataloader_drop_last` 设置为 true, 以此避免不同batch size带来的问题。
+
+
+
+
+### 推理
+
+训练完成后可以直接使用 `outpt_dir` 所指定目录中 policy 文件夹下的 checkpoints 按照[LLM 推理](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm#4-%E6%8E%A8%E7%90%86)部分的介绍来进行推理，请参考相应部分内容。
+
+## Acknowledge
+
+我们借鉴了[PKU-Alignment/safe-rlhf](https://github.com/PKU-Alignment/safe-rlhf)（PKU Beaver）的优秀设计实现，在此对其作者表示感谢。
+
+## 参考文献
+- Zheng R, Dou S, Gao S, et al. Secrets of rlhf in large language models part i: Ppo[J]. arXiv preprint arXiv:2307.04964, 2023.
+- Dai J, Pan X, Sun R, et al. Safe rlhf: Safe reinforcement learning from human feedback[J]. arXiv preprint arXiv:2310.12773, 2023.

llm/Alignment/RM/models

@@ -0,0 +1 @@
+../PPO/models

llm/Alignment/RM/reward_main.py

@@ -0,0 +1,202 @@
+# 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.
+
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Dict, Tuple
+
+import paddle
+
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+from data import PreferenceDataset, parse_dataset
+from models import AutoModelForScore
+from reward_trainer import RewardTrainer
+
+from paddlenlp.trainer import PdArgumentParser, TrainingArguments, get_last_checkpoint
+from paddlenlp.transformers import AutoConfig, AutoTokenizer, LlamaTokenizer
+from paddlenlp.utils.log import logger
+
+
+@dataclass
+class TrainingArguments(TrainingArguments):
+    loss_type: Literal["token-wise", "sequence-wise"] = field(
+        default="sequence-wise",
+        metadata={
+            "help": "Calculate ranking loss with all token-wise reward outputs in the sequence or the "
+            "sequence-wise reward output only (the reward of the last token in each sequence)."
+        },
+    )
+    # regularization
+    regularization: float = field(
+        default=0.0,
+        metadata={"help": "The regularization strength for the L2 regularization for score outputs."},
+    )
+
+
+@dataclass
+class ModelArgument:
+    model_name_or_path: str = field(
+        default=None, metadata={"help": "Build-in pretrained model name or the path to local model."}
+    )
+    normalize_score_during_training: bool = field(
+        default=False, metadata={"help": "Whether to normalize score during training."}
+    )
+    normalizer_type: Literal["RunningMeanStd", "ExponentialMovingAverage"] = field(
+        default=None, metadata={"help": "The type of the reward normalizer."}
+    )
+    normalizer_momentum: float = field(
+        default=None,
+        metadata={
+            "help": "The momentum use in ExponentialMovingAverage, EMA_{t+1} = momentum * x + (1 - momentum) * EMA_t."
+        },
+    )
+    use_flash_attention: bool = field(default=False, metadata={"help": "Whether to use flash attention"})
+
+    @property
+    def extra_model_kwargs(self):
+        """Extra keyword arguments for initializing the model."""
+        return {
+            "score_type": "reward",
+            "do_normalize": self.normalize_score_during_training,
+            "normalizer_type": self.normalizer_type,
+            "momentum": self.normalizer_momentum,
+        }
+
+
+@dataclass
+class DataArgument:
+    dataset_name_or_path: str = field(default=None, metadata={"help": "Name or path for dataset"})
+    task_name: str = field(default=None, metadata={"help": "Additional name to select a more specific task."})
+    train_datasets: str = field(default=None, metadata={"help": "Dataset name(s) registered in the raw dataset."})
+    eval_datasets: str = field(default=None, metadata={"help": "Dataset name(s) registered in the raw dataset."})
+    max_length: int = field(
+        default=2048,
+        metadata={"help": "The maximum length that model input tokens can ."},
+    )
+    eval_with_do_generation: bool = field(default=False, metadata={"help": "Whether to do generation for evaluation"})
+    save_generation_output: bool = field(
+        default=False,
+        metadata={"help": "Whether to save generated text to file when eval_with_do_generation set to True."},
+    )
+    lazy: bool = field(
+        default=False,
+        metadata={
+            "help": "Weather to return `MapDataset` or an `IterDataset`.True for `IterDataset`. False for `MapDataset`."
+        },
+    )
+
+    @property
+    def parsed_train_datasets(self) -> Tuple[str, Dict[str, Any]]:
+        """Parse dataset path and its proportion and optionally additional arguments from `train_datasets`."""
+        return [parse_dataset(string) for string in self.train_datasets.split(",")]
+
+    @property
+    def parsed_eval_datasets(self) -> Tuple[str, Dict[str, Any]]:
+        """Parse dataset path and its proportion and optionally additional arguments from `eval_datasets`."""
+        return [parse_dataset(string) for string in self.eval_datasets.split(",")]
+
+
+def main():
+    # Arguments
+    parser = PdArgumentParser((ModelArgument, DataArgument, TrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    training_args.print_config(model_args, "Model")
+    training_args.print_config(data_args, "Data")
+
+    # Setup GPU & distributed training
+    paddle.set_device(training_args.device)
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, world_size: {training_args.world_size}, "
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16 or training_args.bf16}"
+    )
+
+    # Detecting last checkpoint.
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 1:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Load model
+    if training_args.fp16_opt_level == "O2":
+        if training_args.fp16:
+            dtype = "float16"
+        elif training_args.bf16:
+            dtype = "bfloat16"
+        else:
+            raise ValueError("Please specific dtype: --fp16 or --bf16")
+    else:
+        dtype = "float32"
+
+    if training_args.pipeline_parallel_degree > 1:
+        raise ValueError("Not support pipeline parallel mode.")
+    else:
+        model_config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            tensor_parallel_output=False,
+            tensor_parallel_degree=training_args.tensor_parallel_degree,
+            tensor_parallel_rank=training_args.tensor_parallel_rank,
+            dtype=dtype,
+        )
+        if hasattr(model_config, "use_flash_attention"):
+            model_config.use_flash_attention = model_args.use_flash_attention
+        model = AutoModelForScore.from_pretrained(
+            model_args.model_name_or_path, config=model_config, **model_args.extra_model_kwargs
+        )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path, model_max_length=data_args.max_length, padding_side="right"
+    )
+    if isinstance(tokenizer, LlamaTokenizer) and tokenizer.pad_token_id is None:
+        # tokenizer.pad_token_id = tokenizer.eos_token_id
+        # to be consistent with PKU-Alignment/alpaca-7b-reproduced
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    train_ds = PreferenceDataset(data_args.parsed_train_datasets, tokenizer=tokenizer)
+
+    dev_ds = PreferenceDataset(data_args.parsed_eval_datasets, tokenizer=tokenizer)
+
+    trainer = RewardTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=dev_ds,
+        tokenizer=tokenizer,
+        data_collator=train_ds.get_collator(),
+    )
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+    trainer.train(resume_from_checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    main()

llm/Alignment/RM/reward_trainer.py

@@ -0,0 +1,219 @@
+# 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 typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from models import ScoreModelOutput
+from paddle.io import Dataset
+
+import paddlenlp.trainer.trainer as trainer
+from paddlenlp.data import DataCollator
+from paddlenlp.trainer import (
+    EvalPrediction,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from paddlenlp.trainer.utils import nested_detach
+from paddlenlp.transformers import PretrainedModel, PretrainedTokenizer
+
+_tr_acc = None
+
+speed_metrics = trainer.speed_metrics
+
+
+def patch_speed_metrics(split, start_time, num_samples=None, num_steps=None, seq_length=None):
+    # split: interval, train, eval, test
+    result = speed_metrics(split, start_time, num_samples, num_steps, seq_length)
+    if split not in ["train", "interval"]:
+        return result
+    # accuracy
+    global _tr_acc
+    tr_acc, total_acc_scalar, nested_gather = _tr_acc
+    tr_acc_scalar = nested_gather(tr_acc).mean().item()
+    total_acc_scalar += tr_acc_scalar
+    tr_acc.subtract_(tr_acc)
+    _tr_acc[1] = total_acc_scalar
+    result["accuracy"] = round(tr_acc_scalar / num_steps, 8)
+    if split == "train":
+        result["train_accuracy"] = round(total_acc_scalar / num_steps, 8)
+    return result
+
+
+trainer.speed_metrics = patch_speed_metrics
+
+
+def compute_accuracy(eval_pred) -> Dict[str, float]:
+    higher_end_rewards, lower_end_rewards = eval_pred
+    accuracy = (higher_end_rewards > lower_end_rewards).astype("float32").mean().item()
+    rewards = np.concatenate([higher_end_rewards, lower_end_rewards], axis=0)
+    reward_mean = rewards.mean().item()
+    reward_std = rewards.std().item()
+    return {
+        "accuracy": accuracy,
+        "rewards_mean": reward_mean,
+        "rewards_std": reward_std,
+    }
+
+
+class RewardTrainer(Trainer):
+    def __init__(
+        self,
+        model: Union[PretrainedModel, nn.Layer] = None,
+        criterion: nn.Layer = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Union[Dataset, Dict[str, Dataset]] = None,
+        tokenizer: Optional[PretrainedTokenizer] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
+        preprocess_logits_for_metrics: Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor] = None,
+    ):
+        if compute_metrics is None:
+            compute_metrics = compute_accuracy
+
+        super().__init__(
+            model,
+            criterion,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """
+        How the loss is computed by Trainer. By default, all models return the loss in the first element.
+        Subclass and override for custom behavior.
+        """
+        better_input_ids = inputs["better_input_ids"]
+        worse_input_ids = inputs["worse_input_ids"]
+        better_attention_mask = inputs["better_attention_mask"]
+        worse_attention_mask = inputs["worse_attention_mask"]
+
+        assert better_input_ids.shape[0] == worse_input_ids.shape[0], "batch size mismatch!"
+        batch_size = better_input_ids.shape[0]
+
+        output: ScoreModelOutput = model(
+            paddle.concat([better_input_ids, worse_input_ids], axis=0),
+            attention_mask=paddle.concat([better_attention_mask, worse_attention_mask], axis=0),
+        )
+        if isinstance(output, dict):
+            scores = output.scores  # size = (2 * B, L, 1)
+            end_scores = output.end_scores  # size = (2 * B, 1)
+        else:
+            scores, end_scores = output
+
+        # size = (B, L)
+        higher_rewards, lower_rewards = scores.squeeze(axis=-1).chunk(chunks=2, axis=0)
+        # size = (B,)
+        higher_end_rewards, lower_end_rewards = end_scores.squeeze(axis=-1).chunk(chunks=2, axis=0)
+
+        if self.args.loss_type == "token-wise":
+            losses = []
+            for i in range(batch_size):
+                assert not paddle.all(
+                    paddle.equal(better_input_ids[i], worse_input_ids[i]),
+                ).item(), "The better and worse answers are the same!"
+                higher_end_index = better_attention_mask[i].nonzero()[-1]
+                lower_end_index = worse_attention_mask[i].nonzero()[-1]
+                end_index = max(higher_end_index, lower_end_index)
+
+                diverge_index = (better_input_ids[i] != worse_input_ids[i]).nonzero()[0]
+                assert 0 <= diverge_index <= end_index, "diverge index is out of range!"
+
+                # size = (L,)
+                higher_truncated_rewards = higher_rewards[i, diverge_index : end_index + 1]
+                lower_truncated_rewards = lower_rewards[i, diverge_index : end_index + 1]
+
+                losses.append(
+                    -F.log_sigmoid(higher_truncated_rewards - lower_truncated_rewards).mean(),
+                )
+
+                if self.args.regularization > 0.0:
+                    losses[-1] = losses[-1] + self.args.regularization * (
+                        paddle.square(lower_truncated_rewards).mean() + paddle.square(higher_truncated_rewards).mean()
+                    )
+
+            loss = paddle.stack(losses).mean()  # size = ()
+        elif self.args.loss_type == "sequence-wise":
+            loss = -F.log_sigmoid(higher_end_rewards - lower_end_rewards).mean()
+
+            if self.args.regularization > 0.0:
+                loss = loss + self.args.regularization * (
+                    paddle.square(lower_end_rewards).mean() + paddle.square(higher_end_rewards).mean()
+                )
+        else:
+            raise ValueError(f"Unknown loss type: {self.args.loss_type}")
+
+        accuracy = (higher_end_rewards > lower_end_rewards).cast("float32").mean()  # size = ()
+
+        # TODO(guosheng): use a formal way to replace this hack for accuracy track
+        # in training
+        global _tr_acc
+        if _tr_acc is None:
+            _tr_acc = [paddle.to_tensor(0.0), 0.0, self._nested_gather]
+        _tr_acc[0] = _tr_acc[0] + accuracy.detach()
+
+        if return_outputs:
+            return loss, {
+                "higher_end_rewards": higher_end_rewards,
+                "lower_end_rewards": lower_end_rewards,
+                "accuracy": accuracy,
+            }
+        return loss
+
+    def prediction_step(
+        self,
+        model: nn.Layer,
+        inputs: Dict[str, Union[paddle.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ):
+        if self.args.pipeline_parallel_degree > 1:
+            # hack for pipeline mode
+            inputs = self._prepare_inputs(inputs)
+            return self.prediction_pipeline_step(model, inputs, prediction_loss_only, ignore_keys)
+        else:
+            inputs = self._prepare_inputs(inputs)
+
+        better_input_ids = inputs["better_input_ids"]
+        worse_input_ids = inputs["worse_input_ids"]
+        better_attention_mask = inputs["better_attention_mask"]
+        worse_attention_mask = inputs["worse_attention_mask"]
+
+        with paddle.no_grad():
+            with self.autocast_smart_context_manager():
+                higher_rewards = self.model(better_input_ids, better_attention_mask)
+                lower_rewards = self.model(worse_input_ids, worse_attention_mask)
+            if isinstance(higher_rewards, dict):
+                higher_end_rewards = higher_rewards.end_scores.squeeze(axis=-1)
+                lower_end_rewards = lower_rewards.end_scores.squeeze(axis=-1)
+            else:
+                higher_end_rewards = higher_rewards[-1].squeeze(axis=-1)
+                lower_end_rewards = lower_rewards[-1].squeeze(axis=-1)
+        higher_end_rewards = nested_detach(higher_end_rewards)
+        lower_end_rewards = nested_detach(lower_end_rewards)
+        return None, higher_end_rewards.cast("float32"), lower_end_rewards.cast("float32")

llm/config/llama/ppo.json

@@ -0,0 +1,57 @@
+{
+    "train_datasets": "PKU-SafeRLHF/train",
+    "eval_datasets": "PKU-SafeRLHF/test",
+    "ptx_datasets": "alpaca",
+    "actor_model_name_or_path": "PKU-Alignment/alpaca-7b-reproduced",
+    "reward_model_name_or_path": "PKU-Alignment/beaver-7b-v1.0-reward",
+    "output_dir": "checkpoints/llm_ppo",
+    "max_length": 512,
+    "top_p": 0.8,
+    "temperature": 1.0,
+    "num_return_sequences":1,
+    "repetition_penalty": 1.0,
+    "num_train_epochs": 1,
+    "update_iters": 1,
+    "per_device_prompt_batch_size": 16,
+    "per_device_train_batch_size": 16,
+    "gradient_accumulation_steps": 1,
+    "learning_rate": 1e-5,
+    "weight_decay": 0.01,
+    "lr_scheduler_type": "cosine",
+    "warmup_ratio": 0.03,
+    "recompute": true,
+    "critic_learning_rate": 5e-6,
+    "critic_weight_decay": 0.0,
+    "critic_lr_scheduler_type": "constant",
+    "critic_warmup_ratio": 0.03,
+    "critic_recompute": true,
+    "normalize_reward": false,
+    "kl_coeff": 0.02,
+    "clip_range_ratio": 0.2,
+    "clip_range_score": 50.0,
+    "clip_range_value": 5.0,
+    "ptx_coeff": 16.0,
+    "per_device_eval_batch_size": 16,
+    "logging_steps": 1,
+    "evaluation_strategy": "steps",
+    "eval_steps": 100,
+    "save_strategy": "epoch",
+    "save_steps": 100000,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "save_total_limit": 1,
+    "sharding_parallel_degree": 4,
+    "sharding": "stage1",
+    "tensor_parallel_degree": 2,
+    "pipeline_parallel_degree": 1,
+    "pipeline_parallel_config": "disable_p2p_cache_shape",
+    "max_grad_norm": 1.0, 
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.95,
+    "dataloader_drop_last": false,
+    "eval_mode": "",
+    "offload_level": "freeze_model"
+}

llm/config/llama/rm.json

@@ -0,0 +1,40 @@
+{
+    "model_name_or_path": "PKU-Alignment/alpaca-7b-reproduced",
+    "train_datasets": "PKU-SafeRLHF-30K/train",
+    "eval_datasets": "PKU-SafeRLHF-30K/test",
+    "output_dir": "checkpoints/llama_rm",
+    "per_device_train_batch_size": 16,
+    "gradient_accumulation_steps": 1,
+    "per_device_eval_batch_size": 16,
+    "eval_accumulation_steps": 1,
+    "num_train_epochs": 2,
+    "logging_steps": 1,
+    "evaluation_strategy": "epoch",
+    "eval_steps": 1000,
+    "save_strategy": "epoch",
+    "save_steps": 1000,
+    "max_length": 512,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "sharding_parallel_degree": 4,
+    "sharding": "stage2",
+    "weight_decay": 0.1,
+    "max_grad_norm": 0.0,
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.95,
+    "lr_scheduler_type": "cosine",
+    "learning_rate": 2e-5,
+    "warmup_ratio": 0.03,
+    "normalize_score_during_training": false,
+    "normalizer_type": "ExponentialMovingAverage",
+    "normalizer_momentum": 0.9,
+    "loss_type": "sequence-wise",
+    "regularization": 0.001
+}