opt ppo model

src/rlhf/ppo_old.py

-import math
-from pathlib import Path
-import copy
-from tqdm import tqdm
-from functools import partial
-from collections import deque, namedtuple
-from random import randrange
-
-from beartype import beartype
-from beartype.typing import List, Optional, Callable, Deque
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from torch.optim import Adam
-from torch.utils.data import Dataset, DataLoader
-from torch.nn.utils.rnn import pad_sequence
-
-from pytorch_lightning.utilities import rank_zero_info
-
-from einops import rearrange, repeat
-from einops.layers.torch import Rearrange
-
-from src.model import RWKV
-from src.rlhf.reward import RewardModel
-from src.rlhf.optimizer import get_optimizer
-from src.rlhf.utils import masked_mean, eval_decorator
-
-from accelerate import Accelerator
-
-# actor critic - PaLM with lora
-
-PPOActionCriticReturn = namedtuple('PPOActionCriticReturn', [
-    'actions',
-    'sequence',
-    'mask',
-    'prompt_mask',
-    'action_logits',
-    'values'
-])
-
-@beartype
-class ActorCritic(nn.Module):
-    def __init__(
-        self,
-        rwkv: RWKV,
-        critic_palm: Optional[RWKV] = None,
-        pooled_values = False,
-        actor_lora = True,
-        critic_lora = True,
-        actor_lora_r = 8,
-        critic_lora_r = 8,
-        actor_lora_scope = 'actor',
-        critic_lora_scope = 'critic',
-        actor_dropout = 0.,
-        critic_dropout = 0.
-    ):
-        super().__init__()
-        self.actor_palm = rwkv
-
-        self.critic_palm = critic_palm
-
-        if not exists(self.critic_palm):
-            self.critic_palm = copy.deepcopy(rwkv)
-
-        self.actor_palm.set_dropout(actor_dropout)
-        self.critic_palm.set_dropout(critic_dropout)
-
-        self.actor_lora = actor_lora
-        self.critic_lora = critic_lora
-
-        self.actor_lora_scope = actor_lora_scope if actor_lora else None
-        self.critic_lora_scope = critic_lora_scope if critic_lora else None
-
-        if self.actor_lora:
-            self.actor_palm.add_finetune_params(actor_lora_scope, lora_r = actor_lora_r)
-
-        if self.critic_lora:
-            self.critic_palm.add_finetune_params(critic_lora_scope, lora_r = critic_lora_r)
-
-        self.pooled_values = pooled_values
-        self.value_head = nn.Sequential(
-            nn.Linear(rwkv.dim, 1),
-            Rearrange('... 1 -> ...')
-        )
-
-        nn.init.zeros_(self.value_head[0].bias)
-        nn.init.orthogonal_(self.value_head[0].weight, gain = math.sqrt(2))
-
-    def actor_parameters(self):
-        if not self.actor_lora:
-            return self.actor_palm.parameters()
-
-        return [
-            *self.actor_palm.finetune_parameters(self.actor_lora_scope)
-        ]
-
-    def critic_parameters(self):
-        if not self.actor_lora:
-            return [*self.critic_palm.parameters(), *self.value_head.parameters()]
-
-        return [
-            *self.critic_palm.finetune_parameters(self.critic_lora_scope),
-            *self.value_head.parameters()
-        ]
-
-    @torch.no_grad()
-    @eval_decorator
-    def generate(
-        self,
-        state,
-        max_seq_len,
-        eos_token = None,
-        return_values = False,
-        **kwargs
-    ):
-        actions = self.actor_palm.generate(
-            max_seq_len,
-            prompt = state,       
-            eos_token = eos_token,     
-            finetune_scope = self.actor_lora_scope,
-            use_tqdm = True,
-            **kwargs
-        )
-
-        sequence = torch.cat((state, actions), dim = -1)
-        action_len = actions.shape[-1]
-        state_len = state.shape[-1]
-
-        prompt_mask = torch.arange(sequence.shape[-1], device = state.device) < state_len
-        prompt_mask = repeat(prompt_mask, 'n -> b n', b = sequence.shape[0])
-
-        action_mask = ~prompt_mask
-
-        mask = None
-        if exists(eos_token):
-            mask = ((sequence == eos_token).cumsum(dim = -1) == 0)
-            mask = F.pad(mask, (1, -1), value = True) # include eos token
-            action_mask &= mask
-
-        action_logits, value = self.forward(
-            sequence,
-            mask = action_mask,
-            return_values = return_values
-        )        
-
-        return PPOActionCriticReturn(
-            actions,
-            sequence,
-            mask,
-            prompt_mask,
-            action_logits,
-            value
-        )
-
-    def forward(
-        self,
-        x,
-        mask = None,
-        return_values = True
-    ):
-        action_logits = self.actor_palm(
-            x,
-            finetune_scope = self.actor_lora_scope
-        )
-
-        if not return_values:
-            return action_logits, None
-
-        critic_embeds = self.critic_palm(
-            x,
-            return_only_embedding = True,
-            finetune_scope = self.critic_lora_scope
-        )
-
-        if self.pooled_values:
-            critic_embeds = shift(critic_embeds, shift = 1, dim = -2)
-            critic_embeds = masked_mean(critic_embeds, mask, dim = 1)
-
-        values = self.value_head(critic_embeds)
-
-        return action_logits, values
-
-# data
-
-Memory = namedtuple('Memory', [
-    'sequence',
-    'prompt_mask',
-    'mask',
-    'action_prob',
-    'action_log_prob',
-    'reward',
-    'value'
-])
-
-@beartype
-class ExperienceDataset(Dataset):
-    def __init__(
-        self,
-        data: List[torch.Tensor],
-        device = None
-    ):
-        super().__init__()
-        self.data = data
-        self.device = device
-
-    def __len__(self):
-        return self.data[0].shape[0]
-
-    def __getitem__(self, ind):
-        return tuple(map(lambda t: t[ind].to(self.device), self.data))
-
-def create_dataloader(data, batch_size, shuffle = True, device = None, **kwargs):
-    ds = ExperienceDataset(data, device = device)
-    return DataLoader(ds, batch_size = batch_size, shuffle = shuffle, **kwargs)
-
-# helper functions
-
-def exists(val):
-    return val is not None
-
-def default(val, d):
-    return val if exists(val) else d
-
-def masked_normalize(t, eps = 1e-5, mask = None, dim = None):
-    dim = default(dim, tuple(range(t.ndim)))
-    kwargs = dict(dim = dim, keepdim = True)
-
-    mean = masked_mean(t, mask = mask, **kwargs)
-    mean_centered = t - mean
-    var = masked_mean(mean_centered ** 2, mask = mask, **kwargs)
-
-    return mean_centered * var.clamp(min = eps).rsqrt()
-
-def pad_sequence_fixed(sequences, *args, **kwargs):
-    first_el = sequences[0]
-    has_no_dimension = first_el.ndim == 0
-
-    # if no dimensions, add a single dimension
-    if has_no_dimension:
-        sequences = tuple(map(lambda t: t[None], sequences))
-
-    out = pad_sequence(sequences, *args, **kwargs)
-
-    if has_no_dimension:
-        out = rearrange(out, '... 1 -> ...')
-
-    return out
-
-def log(t, eps = 1e-20):
-    return torch.log(t.clamp(min = eps))
-
-def log_prob(prob, indices):
-    assert prob.shape[:2] == indices.shape, f'preceding shapes of prob {prob.shape[:2]} and indices {indices.shape} must match'
-    return log(prob.gather(-1, indices[..., None])).squeeze(-1)
-
-def shift(t, value = 0, shift = 1, dim = -1):
-    zeros = (0, 0) * (-dim - 1)
-    return F.pad(t, (*zeros, shift, -shift), value = value)
-
-def masked_entropy(prob, dim = -1, mask = None):
-    entropies = (prob * log(prob)).sum(dim = -1)
-    return masked_mean(entropies, mask = mask).mean()
-
-def masked_kl_div(prob1, prob2, mask = None):
-    """
-    need to account for variable sequence lengths, therefore not using the built-in functional version
-    """
-    kl_divs = (prob1 * (log(prob2) - log(prob1))).sum(dim = -1)
-
-    if not exists(mask):
-        return kl_divs.mean()
-
-    return masked_mean(kl_divs, mask).mean()
-
-def clipped_value_loss(values, rewards, old_values, clip):
-    value_clipped = old_values + (values - old_values).clamp(-clip, clip)
-    value_loss_1 = (value_clipped.flatten() - rewards) ** 2
-    value_loss_2 = (values.flatten() - rewards) ** 2
-    return torch.mean(torch.max(value_loss_1, value_loss_2))
-
-# rlhf trainer
-
-@beartype
-class RLHFTrainer(nn.Module):
-    def __init__(
-        self,
-        args,
-        accelerate_kwargs: dict = {}
-    ):
-        super().__init__()
-
-        self.args = args
-
-        self.accelerate = Accelerator(**accelerate_kwargs)
-
-        # 加载 RWKV 模型
-        rwkv = RWKV(args)
-
-        if len(args.load_sft_model) == 0:
-            rank_zero_info(f"SFT must load model, please input ")
-            exit(1)
-
-        rank_zero_info(f"########## Loading {args.load_sft_model}... ##########")
-        try:
-            load_dict = torch.load(args.load_sft_model, map_location="cpu")
-        except:
-            rank_zero_info(f"Bad checkpoint {args.load_sft_model}")
-            exit(1)
-
-        if args.load_partial == 1:
-            load_keys = load_dict.keys()
-            for k in rwkv.state_dict():
-                if k not in load_keys:
-                    load_dict[k] = rwkv.state_dict()[k]
-        rwkv.load_state_dict(load_dict)
-
-        self.rwkv = rwkv
-
-        # 使用 RWKV 初始化 actor_critic
-        actor_critic = ActorCritic(
-            rwkv = self.rwkv,
-            actor_lora = args.actor_lora,
-            critic_lora = args.critic_lora,
-            actor_lora_r = args.actor_lora_r,
-            critic_lora_r = args.critic_lora_r,
-            pooled_values = args.critic_pooled_values,
-            actor_dropout = args.actor_dropout,
-            critic_dropout = args.critic_dropout
-        ).to(self.rwkv.device)
-
-        self.actor_critic = actor_critic
-
-        # 加载 reward_model，并将 reward_model 设置为 evaluation 模式 
-        reward_model = RewardModel(args)
-        reward_model.load(args.load_rm_model)
-        self.reward_model = reward_model.eval()
-
-        # optimizers
-        self.actor_optim = get_optimizer(actor_critic.actor_parameters(), lr = self.args.actor_lr, wd = self.args.actor_wd, betas = self.args.betas, eps = self.args.actor_adam_eps, use_lion = self.args.use_lion)
-        self.critic_optim = get_optimizer(actor_critic.critic_parameters(), lr = self.args.critic_lr, wd = self.args.critic_wd, betas = self.args.betas, eps = self.args.critic_adam_eps, use_lion = self.args.use_lion)
-
-        # prepare with accelerator
-
-        (
-            self.actor_critic,
-            self.reward_model,
-            self.actor_optim,
-            self.critic_optim
-        ) = self.accelerate.prepare(
-            self.actor_critic,
-            self.reward_model,
-            self.actor_optim,
-            self.critic_optim
-        )
-
-
-    def print(self, msg):
-        return self.accelerate.print(msg)
-
-    def save(self, filepath = './checkpoint.pt'):
-        torch.save(self.actor_critic.state_dict(), filepath)
-
-    def load(self, filepath = './checkpoint.pt'):
-        state_dict = torch.load(filepath)
-        self.actor_critic.load_state_dict(state_dict)
-
-    @property
-    def device(self):
-        return self.accelerate.device
-
-    @torch.no_grad()
-    def generate(
-        self,
-        max_seq_len,
-        *args,
-        prompt,
-        num_samples = 4,  # sample 4 per prompt and select the one with highest reward
-        **kwargs
-    ):
-        assert prompt.ndim == 1, 'only one prompt allowed at a time for now'
-        prompt = repeat(prompt, 'n -> b n', b = num_samples)
-
-        actor_critic = self.accelerate.unwrap_model(self.actor_critic)
-        reward_model = self.accelerate.unwrap_model(self.reward_model)
-
-        actor_critic.eval()
-
-        (
-            actions,
-            sequences,
-            mask,
-            prompt_mask,
-            action_logits,
-            _
-        ) = actor_critic.generate(
-            prompt,
-            *args,
-            max_seq_len = max_seq_len,
-            return_values = False,
-            **kwargs
-        )
-
-        rewards = reward_model(
-            sequences,
-            prompt_mask = prompt_mask,
-            mask = mask,
-            sample = True
-        )
-
-        best_sequence_index = rewards.topk(1, dim = -1).indices
-
-        best_sequence = sequences[best_sequence_index]
-        best_sequence = rearrange(best_sequence, '1 ... -> ...')
-
-        return best_sequence
-
-    def learn(
-        self,
-        memories: Deque[Memory]
-    ):
-        # stack all data stored in the memories
-
-        all_memories_stacked_and_padded = list(map(partial(pad_sequence_fixed, batch_first = True), zip(*memories)))
-
-        # prepare dataloader for policy phase training
-
-        dl = create_dataloader(all_memories_stacked_and_padded, self.minibatch_size, device = self.device)
-
-        self.actor_critic.train()
-
-        # PPO training
-
-        for _ in range(self.epochs):
-            for (
-                sequences,
-                prompt_masks,
-                masks,
-                old_action_probs,
-                old_log_probs,
-                rewards,
-                old_values
-            ) in dl:
-                action_masks = ~prompt_masks & masks
-
-                action_logits, values = self.actor_critic(
-                    sequences,
-                    mask = action_masks
-                )
-
-                action_logits = shift(action_logits, shift = 1, dim = -2) # need to shift along sequence dimension by 1, since actions start from the last prompt (state) token
-                action_len = old_log_probs.shape[-1]
-
-                action_probs = action_logits.softmax(dim = -1)
-                action_log_probs = log_prob(action_probs, sequences)
-                action_log_probs = action_log_probs[:, -action_len:]
-
-                # calculate entropies, taking into account which part of the sequence is actually an action
-
-                entropies = masked_entropy(action_probs, mask = action_masks)
-
-                # calculate kl div between old action probs and new ones, taking into account which part of the sequence is action or not
-
-                kl_div_loss = 0.
-
-                if self.args.kl_div_loss_weight > 0:
-                    kl_div_loss = masked_kl_div(action_probs, old_action_probs, mask = action_masks) * self.args.kl_div_loss_weight
-
-                # handle non-pooled values
-
-                normalize_kwargs = dict()
-
-                if old_values.ndim == 2:
-                    old_values, values = map(lambda t: shift(t, shift = 1, dim = -2), (old_values, values))
-
-                    old_values = old_values[:, -action_len:]
-                    values = values[:, -action_len:]
-                    rewards = rearrange(rewards, 'b -> b 1')
-                    normalize_kwargs = dict(dim = -1, mask = action_masks[:, -action_len:])
-
-                if values.ndim < rewards.ndim:
-                    values = rearrange(values, '... -> ... 1')
-
-                # calculate clipped surrogate objective, classic PPO loss
-
-                ratios = (action_log_probs - old_log_probs).exp()
-                advantages = masked_normalize(rewards - old_values, **normalize_kwargs)
-
-                if advantages.ndim == 1:
-                    advantages = rearrange(advantages, 'b -> b 1')
-
-                surr1 = ratios * advantages
-                surr2 = ratios.clamp(1 - self.args.eps_clip, 1 + self.args.eps_clip) * advantages
-                policy_loss = - torch.min(surr1, surr2) - self.args.beta_s * entropies
-
-                # combine losses
-
-                loss = policy_loss.mean() + kl_div_loss
-
-                # update actor
-
-                self.accelerate.backward(loss)
-
-                self.print(f'policy_loss: {loss.item():.3f}')
-
-                if exists(self.args.max_norm):
-                    self.accelerator.clip_grad_norm_(self.actor_critic.actor_parameters(), self.args.max_norm)
-
-                self.actor_optim.step()
-                self.actor_optim.zero_grad()
-
-                # calculate value loss and update value network separate from policy network
-
-                value_loss = clipped_value_loss(values, rewards, old_values, self.args.value_clip)
-                value_loss = value_loss.mean()
-
-                self.print(f'critic_loss: {value_loss.item():.3f}')
-
-                self.accelerate.backward(value_loss)
-
-                if exists(self.args.max_norm):
-                    self.accelerator.clip_grad_norm_(self.actor_critic.critic_parameters(), self.args.max_norm)
-
-                self.critic_optim.step()
-                self.critic_optim.zero_grad()
-
-    def train(
-        self,
-        num_episodes = 50000,
-        max_timesteps = 500,
-        update_timesteps = 5000,
-        max_batch_size = 16,
-        eos_token = None,
-        temperature = 1.
-    ):
-        device = self.device
-
-        time = 0
-        memories = deque([])
-
-        for eps in tqdm(range(num_episodes), desc = 'episodes'):
-            for timestep in range(max_timesteps):
-                time += 1
-
-                # select a bunch of random states (prompts)
-                # and get the action (sampled sequence from palm as well as the action probs)
-                # also calculate the reward using reward model and store
-
-                rand_prompt_index = randrange(0, self.num_prompts)
-
-                state = self.prompt_token_ids[rand_prompt_index]
-
-                # remove padding from state
-
-                state_mask = state != self.args.pad_value
-                state = state[state_mask]
-
-                # get predicted sequence
-
-                (
-                    actions,
-                    sequence,
-                    mask,
-                    prompt_mask,
-                    action_logits,
-                    value
-                ) = self.actor_critic.generate(
-                    rearrange(state, 'n -> 1 n'),
-                    max_seq_len = self.args.ctx_len,
-                    eos_token = eos_token,
-                    temperature = temperature,
-                    return_values = True
-                )
-                action_logits = shift(action_logits, shift = 1, dim = -2) # need to shift along sequence dimension by 1, since actions start from the last prompt (state) token
-
-                action_prob = action_logits.softmax(dim = -1)
-
-                action_len = actions.shape[-1]
-                action_log_prob = log_prob(action_prob, sequence)
-                action_log_prob = action_log_prob[:, -action_len:]
-
-                actions = rearrange(actions, '1 ... -> ...')
-
-                # get reward as given by supervised trained reward model
-
-                sequence = torch.cat((state, actions), dim = 0)
-
-                prompt_length = len(state)
-                prompt_mask = torch.arange(sequence.shape[-1], device = device) < prompt_length
-
-                sequence = rearrange(sequence, 'n -> 1 n')
-                prompt_mask = rearrange(prompt_mask, 'n -> 1 n')
-                mask = rearrange(mask, 'n -> 1 n') if exists(mask) else torch.ones(sequence.shape, dtype = torch.bool, device = device)
-
-                reward = self.reward_model(
-                    sequence,
-                    prompt_mask = prompt_mask,
-                    mask = mask,
-                    sample = True
-                )
-
-                detach_to_cpu_ = lambda t: rearrange(t.detach().cpu(), '1 ... -> ...')
-
-                # store memory for learning
-
-                memories.append(Memory(*map(detach_to_cpu_, (
-                    sequence,
-                    prompt_mask,
-                    mask,
-                    action_prob,
-                    action_log_prob,
-                    reward,
-                    value
-                ))))
-
-                # learn from the stored memories
-
-                if time % update_timesteps == 0:
-                    self.learn(memories)
-                    memories.clear()
-
-        print('rlhf training complete')