https://zhuanlan.zhihu.com/p/591474085
https://zhuanlan.zhihu.com/p/613315873?utm_id=0

https://karpathy.ai/stateofgpt.pdf
https://zhuanlan.zhihu.com/p/616708590
https://openreview.net/forum?id=10uNUgI5Kl
https://huggingface.co/blog/zh/rlhf
https://huggingface.co/datasets/CarperAI/openai_summarize_comparisons/viewer/CarperAI--openai_summarize_comparisons/train?row=0
https://zhuanlan.zhihu.com/p/450690041 

(
    'The USA entered World War I because Germany attempted to enlist Mexico as an ally, and for what other reason?',
    [
        "The United States entered World War I because of Germany's use of submarine warfare against ships in the Atlantic Ocean, which was hurting American exports to Europe. Additionally, Germany tried to enlist Mexico as an ally against the United States, an event which convinced American businessmen and industrialists that the United States should enter the war.",
        'The USA entered World War I because Germany attempted to enlist Mexico as an ally and for the Zimmerman Telegram.'
    ]
)

class WebGPT:
    name = "openai/webgpt_comparisons"

    def __init__(self, split: str = "train"):
        super().__init__()
        self.split = split
        dataset = load_dataset(self.name, split=self.split)
        self.dataset_dict = defaultdict(dict)
        for item in dataset:
            post_id = item["question"]["id"]
            if post_id not in self.dataset_dict.keys():
                self.dataset_dict[post_id] = {
                    "full_text": item["question"]["full_text"],
                    "answers": [],
                }
                if item["score_0"] > 0:
                    answers = [item["answer_0"], item["answer_1"]]
                elif item["score_0"] < 0:
                    answers = [item["answer_1"], item["answer_0"]]
                else:
                    answers = []
                answers = [re.sub(r"\[\d+\]", "", answer) for answer in answers]
                answers = [
                    ".".join([sent.strip() for sent in answer.split(".")])
                    for answer in answers
                ]
                if answers:
                    self.dataset_dict[post_id]["answers"].extend(answers)
                else:
                    _ = self.dataset_dict.pop(post_id)

        self.post_ids = list(self.dataset_dict.keys())

    def __len__(self):
        return len(self.post_ids)

    def __getitem__(self, idx):
        question, answers = self.dataset_dict[self.post_ids[idx]].values()
        return question, answers

@dataclass
class RMDataCollator:
    tokenizer: PreTrainedTokenizer
    max_length: int = 512

    def format_example(self, example, eos, prompt=False):
        sp_token = SPECIAL_TOKENS["prompter"] if prompt else SPECIAL_TOKENS["assistant"]
        return "{}{}{}".format(sp_token, example, eos)

    def process_example(self, example):
        trunc_len = 0
        eos = self.tokenizer.eos_token
        prefix, outputs = example
        prefix = self.format_example(example, eos, prompt=True)
        outputs = [self.format_example(output, eos) for output in outputs]

        prefix_tokens = self.tokenizer.encode(prefix)
        input_ids, attention_masks = [], []
        for output in outputs:
            out_tokens = self.tokenizer.encode(
                output,
            )
            if len(prefix_tokens) + len(out_tokens) > self.max_length:
                trunc_len = max(
                    0, len(prefix_tokens) + len(out_tokens) - self.max_length
                )
            prefix_tokens = prefix_tokens[trunc_len:]
            out_tokens = prefix_tokens + out_tokens
            out_tokens = out_tokens[: self.max_length]
            pad_len = self.max_length - len(out_tokens)
            attn_masks = [1] * len(out_tokens) + [0] * pad_len
            out_tokens += [self.tokenizer.pad_token_id] * pad_len
            input_ids.append(out_tokens)
            attention_masks.append(attn_masks)
        return input_ids, attention_masks

    def __call__(self, examples):
        batch_k_lens = [0]
        input_ids, attention_masks = [], []
        for i, example in enumerate(examples):
            inp_ids, attn_masks = self.process_example(example)
            input_ids.extend(inp_ids)
            attention_masks.extend(attn_masks)
            batch_k_lens.append(batch_k_lens[i] + len(inp_ids))

        return {
            "input_ids": torch.tensor(input_ids),
            "attention_mask": torch.tensor(attention_masks),
            "k_lens": batch_k_lens,
        }

@dataclass
class GPTNeoxRMOuptput(ModelOutput):
    """
    Reward Model Output
    """

    logits: torch.FloatTensor = None


class GPTNeoXRM(GPTNeoXPreTrainedModel):
    """ """

    def __init__(
        self,
        config,
    ):
        super().__init__(config)
        self.gpt_neox = GPTNeoXModel(config)
        self.out_layer = nn.Linear(config.hidden_size, 1)

    def forward(
        self,
        input_ids,
        attention_mask,
        **kwargs,
    ):
        return_dict = (
            kwargs.get("return_dict")
            if kwargs.get("return_dict") is not None
            else self.config.use_return_dict
        )
        outputs = self.gpt_neox(
            input_ids,
            attention_mask,
            return_dict=return_dict,
            **kwargs,
        )
        hidden_states = outputs[0]
        if attention_mask is None:
            hidden_states = hidden_states.mean(dim=1)
        else:
            hidden_states = (hidden_states * attention_mask.unsqueeze(-1)).sum(
                dim=1
            ) / attention_mask.sum(dim=1).unsqueeze(-1)
        lm_logits = self.out_layer(hidden_states)

        if not return_dict:
            return (lm_logits,) + outputs[1:]

        return GPTNeoxRMOuptput(logits=lm_logits)

class RMLoss(nn.Module):
    """ """

    def __init__(
        self,
        reduction=None,
        beta=0.001,
    ):
        super().__init__()
        self.reduction = reduction
        self.beta = beta

    def forward(
        self,
        logits,
        k_lens=None,
    ):
        total_loss = []
        indices = list(zip(k_lens[:-1], k_lens[1:]))
        for start, end in indices:
            combinations = torch.combinations(
                torch.arange(start, end, device=logits.device), 2
            )
            positive = logits[combinations[:, 0]]
            negative = logits[combinations[:, 1]]
            l2 = 0.5 * (positive**2 + negative**2)
            loss = (
                -1 * nn.functional.logsigmoid(positive - negative) + self.beta * l2
            ).mean()
            total_loss.append(loss)

        total_loss = torch.stack(total_loss)
        if self.reduction == "mean":
            total_loss = total_loss.mean()
        return total_loss
view raw

https://explodinggradients.com/reward-modeling-for-large-language-models-with-code
https://huggingface.co/datasets/openai/summarize_from_feedback/viewer/axis/test?row=0 

# 单GPU
cd sft/ && CUDA_VISIBLE_DEVICES=0 python3 train_gptj_summarize.py
# 多GPU
cd sft/ && deepspeed train_gptj_summarize.py

def create_comparison_dataset(path="CarperAI/openai_summarize_comparisons", split="train"):
    dataset = load_dataset(path, split=split)
    pairs = []
    for sample in tqdm(dataset):
        pair = {}
        prompt = sample["prompt"]
        chosen_summary = sample["chosen"]
        rejected_summary = sample["rejected"]
        if chosen_summary == rejected_summary:
            continue
        if len(chosen_summary.split()) < 5 or len(rejected_summary.split()) < 5:
            continue
        pair["chosen"] = prompt + "\n" + chosen_summary
        pair["rejected"] = prompt + "\n" + rejected_summary
        pairs.append(pair)
    return pairs 

class PairwiseDataset(Dataset):
    def __init__(self, pairs, tokenizer, max_length):
        self.chosen_input_ids = []
        self.chosen_attn_masks = []
        self.rejected_input_ids = []
        self.rejected_attn_masks = []
        for pair in tqdm(pairs):
            chosen, rejected = pair["chosen"], pair["rejected"]
            chosen_encodings_dict = tokenizer(
                "<|startoftext|>" + chosen + "<|endoftext|>",
                truncation=True,
                max_length=max_length,
                padding="max_length",
                return_tensors="pt",
            )
            rejected_encodings_dict = tokenizer(
                "<|startoftext|>" + rejected + "<|endoftext|>",
                truncation=True,
                max_length=max_length,
                padding="max_length",
                return_tensors="pt",
            )
            self.chosen_input_ids.append(chosen_encodings_dict["input_ids"])
            self.chosen_attn_masks.append(chosen_encodings_dict["attention_mask"])
            self.rejected_input_ids.append(rejected_encodings_dict["input_ids"])
            self.rejected_attn_masks.append(rejected_encodings_dict["attention_mask"])

    def __len__(self):
        return len(self.chosen_input_ids)

    def __getitem__(self, idx):
        return (
            self.chosen_input_ids[idx],
            self.chosen_attn_masks[idx],
            self.rejected_input_ids[idx],
            self.rejected_attn_masks[idx],
        )

class DataCollatorReward:
    def __call__(self, data):
        batch = {}
        batch["input_ids"] = torch.cat([f[0] for f in data] + [f[2] for f in data])
        batch["attention_mask"] = torch.cat([f[1] for f in data] + [f[3] for f in data])
        batch["labels"] = torch.tensor([0] * len(data) + [1] * len(data))
        return batch

class GPTRewardModel(nn.Module):
    def __init__(self, model_path):
        super().__init__()
        model = AutoModelForCausalLM.from_pretrained(model_path)
        self.config = model.config
        # `gpt-neo(x)` models use `hidden_size` attribute names instead of `n_embd``
        self.config.n_embd = self.config.hidden_size if hasattr(self.config, "hidden_size") else self.config.n_embd
        self.transformer = model.transformer
        self.v_head = nn.Linear(self.config.n_embd, 1, bias=False)
        self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.PAD_ID = self.tokenizer(self.tokenizer.pad_token)["input_ids"][0]

    def forward(
        self,
        input_ids=None,
        past_key_values=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        mc_token_ids=None,
        labels=None,
        return_dict=False,
        output_attentions=False,
        output_hidden_states=False,
    ):
        loss = None
        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )

        hidden_states = transformer_outputs[0]

        rewards = self.v_head(hidden_states).squeeze(-1)
        chosen_end_scores = []
        rejected_end_scores = []

        # Split the inputs and rewards into two parts, chosen and rejected
        assert len(input_ids.shape) == 2
        bs = input_ids.shape[0] // 2
        chosen = input_ids[:bs]
        rejected = input_ids[bs:]
        chosen_rewards = rewards[:bs]
        rejected_rewards = rewards[bs:]

        loss = 0
        inference = False
        for i in range(bs):
            if torch.all(torch.eq(chosen[i], rejected[i])).item():
                c_inds = (chosen[i] == self.PAD_ID).nonzero()
                c_ind = c_inds[0].item() if len(c_inds) > 0 else chosen.shape[1]
                chosen_end_scores.append(chosen_rewards[i, c_ind - 1])
                inference = True
                continue

            # Check if there is any padding otherwise take length of sequence
            c_inds = (chosen[i] == self.PAD_ID).nonzero()
            c_ind = c_inds[0].item() if len(c_inds) > 0 else chosen.shape[1]
            r_inds = (rejected[i] == self.PAD_ID).nonzero()
            r_ind = r_inds[0].item() if len(r_inds) > 0 else rejected.shape[1]
            end_ind = max(c_ind, r_ind)

            # Retrieve first index where trajectories diverge
            divergence_ind = (chosen[i] != rejected[i]).nonzero()[0]
            assert divergence_ind > 0

            # Index into the correct rewards
            c_truncated_reward = chosen_rewards[i][divergence_ind:end_ind]
            r_truncated_reward = rejected_rewards[i][divergence_ind:end_ind]

            # Append the last rewards to the list of end scores
            chosen_end_scores.append(c_truncated_reward[-1])
            rejected_end_scores.append(r_truncated_reward[-1])

            # Compute loss based on truncated rewards (ignore padding)
            loss += -torch.log(torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()
        loss = loss / bs

        if not inference:
            chosen_end_scores = torch.stack(chosen_end_scores)
            rejected_end_scores = torch.stack(rejected_end_scores)

        if inference:
            chosen_end_scores = torch.stack(chosen_end_scores)
            return {"chosen_end_scores": chosen_end_scores}

        return {
            "loss": loss,
            "chosen_end_scores": chosen_end_scores,
            "rejected_end_scores": rejected_end_scores,
        }

    # Initialize the reward model from the (supervised) fine-tuned GPT-J
    model = GPTRewardModel("CarperAI/openai_summarize_tldr_sft")

    # Freeze the first 70% of the hidden layers of the reward model backbone
    layers = model.transformer.h
    num_layers = len(layers)
    num_unfrozen = int(0.3 * num_layers)
    for layer in layers[:-num_unfrozen]:
        layer.requires_grad_(False)

    # Create the comparisons datasets
    data_path = "CarperAI/openai_summarize_comparisons"
    train_pairs = create_comparison_dataset(data_path, "train")
    val_pairs = create_comparison_dataset(data_path, "test")

    # Make pairwise datasets for training
    max_length = 550
    train_dataset = PairwiseDataset(train_pairs, tokenizer, max_length=max_length)
    val_dataset = PairwiseDataset(val_pairs, tokenizer, max_length=max_length)

    # Create the collator to gather batches of pairwise comparisons
    data_collator = DataCollatorReward()

    Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        compute_metrics=compute_metrics,
        eval_dataset=val_dataset,
        data_collator=data_collator,
    ).train()

cd reward_model/ && deepspeed train_reward_model_gptj.py

mkdir reward_model/rm_checkpoint
wget https://huggingface.co/CarperAI/openai_summarize_tldr_rm_checkpoint/resolve/main/pytorch_model.bin -O reward_model/rm_checkpoint/pytorch_model.bin

https://huggingface.co/datasets/CarperAI/openai_summarize_comparisons/viewer/CarperAI--openai_summarize_comparisons/train?row=0
https://link.zhihu.com/?target=https%3A//github.com/CarperAI/trlx/tree/main/examples/summarize_rlhf
https://github.com/CarperAI/trlx 
https://github.com/CarperAI/trlx/tree/main/examples/summarize_rlhf

https://github.com/allenai/RL4LMs