lm_human_preferences/train_reward.py

#!/usr/bin/env python3

import json
import os
from dataclasses import dataclass, field
from functools import partial
from typing import Optional

import numpy as np
import tensorflow as tf
from mpi4py import MPI
from tensorflow.contrib import summary

from lm_human_preferences import label_types, lm_tasks, rewards
from lm_human_preferences.language import trained_models
from lm_human_preferences.policy import Policy
from lm_human_preferences.utils import core as utils
from lm_human_preferences.utils import gcs, hyperparams
from lm_human_preferences.utils.core import Schema


@dataclass
class LabelHParams(hyperparams.HParams):
    type: str = None
    num_train: int = None
    source: str = None


@dataclass
class RunHParams(hyperparams.HParams):
    seed: Optional[int] = None
    log_interval: int = 10
    save_interval: int = 50
    save_dir: Optional[str] = None

@dataclass
class HParams(hyperparams.HParams):
    run: RunHParams = field(default_factory=RunHParams)

    task: lm_tasks.TaskHParams = field(default_factory=lm_tasks.TaskHParams)
    labels: LabelHParams = field(default_factory=LabelHParams)

    batch_size: int = 40  # total across ranks
    lr: float = 5e-5

    rollout_batch_size: int = 64
    normalize_samples: int = 0  # Samples used to estimate reward mean and std
    debug_normalize: int = 0  # Samples used to check that normalization worked
    # Whether, before training, to normalize the rewards on the policy to the scales on the training buffer.
    # (For comparisons, just use mean 0, var 1.)
    normalize_before: bool = False
    # Whether, after training, to normalize the rewards on the ref policy to mean 0, var 1
    # (so the KL coefficient always has the same meaning).
    normalize_after: bool = False

    def validate(self, *, prefix=''):
        super().validate(prefix=prefix)
        utils.exact_div(self.labels.num_train, self.batch_size)

def round_down_to_multiple(n, divisor):
    return n - n % divisor


def download_labels(source, label_type, question_schemas, total_labels, comm):
    schemas = {**question_schemas, **label_type.label_schemas()}

    """
    if self.is_root:
        with tf.device('cpu:0'):
            self._enqueue_phs = {
                name: tf.placeholder(name=name, dtype=schema.dtype, shape=(None,) + schema.shape)
                for name, schema in self.schemas.items()
            }
            self._enqueue_answers = self.answer_queue.enqueue_many(self._enqueue_phs)
    else:
        self._enqueue_phs = None
        self._enqueue_answers = None
    """

    # TODO: download on just one rank?  then do: labels = utils.mpi_bcast_tensor_dict(labels, comm=comm)
    if source != 'test':
        with open(gcs.download_file_cached(source, comm=comm)) as f:
            results = json.load(f)
            print('Num labels found in source:', len(results))
    else:
        results = [
            {
                name: np.zeros(schema.shape, dtype=schema.dtype.as_numpy_dtype)
                for name, schema in schemas.items()
            }
            for _ in range(50)
        ]

    assert len(results) >= total_labels
    results = results[:total_labels]
    return {k: [a[k] for a in results] for k in schemas.keys()}


class RewardModelTrainer():
    def __init__(self, *, reward_model, policy, query_sampler, hparams, comm):
        self.reward_model = reward_model

        self.policy = policy
        self.hparams = hparams
        self.num_ranks = comm.Get_size()
        self.rank = comm.Get_rank()
        self.comm = comm

        self.label_type = label_types.get(hparams.labels.type)
        self.question_schemas = self.label_type.question_schemas(
            query_length=hparams.task.query_length,
            response_length=hparams.task.response_length,
        )

        data_schemas = {
            **self.question_schemas,
            **self.label_type.label_schemas(),
        }

        with tf.device(None), tf.device('/cpu:0'):
            with tf.variable_scope('label_buffer', use_resource=True, initializer=tf.zeros_initializer):
                self.train_buffer = utils.SampleBuffer(capacity=hparams.labels.num_train, schemas=data_schemas)

        with tf.name_scope('train_reward'):
            summary_writer = utils.get_summary_writer(self.hparams.run.save_dir, subdir='reward_model', comm=comm)

            @utils.graph_function(
                indices=Schema(tf.int32, (None,)),
                lr=Schema(tf.float32, ()))
            def train_batch(indices, lr):
                with tf.name_scope('minibatch'):
                    minibatch = self.train_buffer.read(indices)
                    stats = self.label_type.loss(reward_model=self.reward_model.get_rewards_op, labels=minibatch)

                    train_op = utils.minimize(
                        loss=stats['loss'], lr=lr, params=self.reward_model.get_params(), name='opt', comm=self.comm)

                    with tf.control_dependencies([train_op]):
                        step_var = tf.get_variable(name='train_step', dtype=tf.int64, shape=(), trainable=False, use_resource=True)
                        step = step_var.assign_add(1) - 1

                        stats = utils.FlatStats.from_dict(stats).map_flat(partial(utils.mpi_allreduce_mean, comm=comm)).as_dict()

                        train_stat_op = utils.record_stats(stats=stats, summary_writer=summary_writer, step=step, log_interval=hparams.run.log_interval, comm=comm)

                return train_stat_op
            self.train_batch = train_batch

        if self.hparams.normalize_before or self.hparams.normalize_after:
            @utils.graph_function()
            def target_mean_std():
                """Returns the means and variances to target for each reward model"""
                # Should be the same on all ranks because the train_buf should be the same
                scales = self.label_type.target_scales(self.train_buffer.data())
                if scales is None:
                    return tf.zeros([]), tf.ones([])
                else:
                    mean, var = tf.nn.moments(scales, axes=[0])
                    return mean, tf.sqrt(var)
            self.target_mean_std = target_mean_std

            def stats(query_responses):
                rewards = np.concatenate([self.reward_model.get_rewards(qs, rs) for qs, rs in query_responses], axis=0)
                assert len(rewards.shape) == 1, f'{rewards.shape}'
                sums = np.asarray([rewards.sum(axis=0), np.square(rewards).sum(axis=0)])
                means, sqr_means = self.comm.allreduce(sums, op=MPI.SUM) / (self.num_ranks * rewards.shape[0])
                stds = np.sqrt(sqr_means - means ** 2)
                return means, stds
            self.stats = stats

            def log_stats_after_normalize(stats):
                if comm.Get_rank() != 0:
                    return
                means, stds = stats
                print(f'after normalize: {means} +- {stds}')
            self.log_stats_after_normalize = log_stats_after_normalize

            def reset_reward_scales():
                self.reward_model.reset_reward_scale()
            self.reset_reward_scales = reset_reward_scales

            def set_reward_norms(mean, std, new_mean, new_std):
                print(f'targets: {new_mean} +- {new_std}')
                print(f'before normalize: {mean} +- {std}')
                assert np.isfinite((mean, std, new_mean, new_std)).all()
                self.reward_model.set_reward_norm(old_mean=mean, old_std=std, new_mean=new_mean, new_std=new_std)
            self.set_reward_norms = set_reward_norms

        if self.hparams.normalize_before or self.hparams.normalize_after:
            @utils.graph_function()
            def sample_policy_batch():
                queries = query_sampler('ref_queries')['tokens']
                responses = policy.respond_op(
                    queries=queries, length=hparams.task.response_length)['responses']
                return queries, responses

            def sample_policy_responses(n_samples):
                n_batches = utils.ceil_div(n_samples, hparams.rollout_batch_size)
                return [sample_policy_batch() for _ in range(n_batches)]
            self.sample_policy_responses = sample_policy_responses

        @utils.graph_function(labels=utils.add_batch_dim(data_schemas))
        def add_to_buffer(labels):
            return self.train_buffer.add(**labels)
        self.add_to_buffer = add_to_buffer

    def normalize(self, sample_fn, target_means, target_stds):
        if not self.hparams.normalize_samples:
            return

        self.reset_reward_scales()
        query_responses = sample_fn(self.hparams.normalize_samples)
        means, stds = self.stats(query_responses)

        self.set_reward_norms(means, stds, target_means, target_stds)
        if self.hparams.debug_normalize:
            query_responses = sample_fn(self.hparams.debug_normalize)
            stats = self.stats(query_responses)
            self.log_stats_after_normalize(stats)

    def train(self):
        labels = download_labels(
            self.hparams.labels.source,
            label_type=self.label_type,
            question_schemas=self.question_schemas,
            total_labels=self.hparams.labels.num_train,
            comm=self.comm
        )

        self.add_to_buffer(labels)

        if self.hparams.normalize_before:
            target_mean, target_std = self.target_mean_std()
            self.normalize(self.sample_policy_responses, target_mean, target_std)

        # Collect training data for reward model training.  train_indices will include the indices
        # trained on across all ranks, and its size must be a multiple of minibatch_size.
        per_rank_batch_size = utils.exact_div(self.hparams.batch_size, self.num_ranks)

        # Make sure each rank gets the same shuffle so we train on each point exactly once
        train_indices = self.comm.bcast(np.random.permutation(self.hparams.labels.num_train))

        # Train on train_indices
        print(self.rank, "training on", self.hparams.labels.num_train, "in batches of", per_rank_batch_size)
        for start_index in range(0, self.hparams.labels.num_train, self.hparams.batch_size):
            end_index = start_index + self.hparams.batch_size
            all_ranks_indices = train_indices[start_index:end_index]
            our_indices = all_ranks_indices[self.rank::self.num_ranks]
            lr = (1 - start_index / self.hparams.labels.num_train) * self.hparams.lr
            self.train_batch(our_indices, lr)

        if self.hparams.normalize_after:
            target_mean, target_std = np.zeros([]), np.ones([])
            self.normalize(self.sample_policy_responses, target_mean, target_std)



def train(hparams: HParams):
    with tf.Graph().as_default():
        hyperparams.dump(hparams)
        utils.set_mpi_seed(hparams.run.seed)

        m = trained_models.TrainedModel(hparams.task.policy.initial_model)
        encoder = m.encoding.get_encoder()
        hyperparams.dump(m.hparams(), name='model_hparams')

        comm = MPI.COMM_WORLD
        ref_policy = Policy(
            m, scope='ref_policy',
            is_root=comm.Get_rank() == 0,
            embed_queries=lm_tasks.query_formatter(hparams.task, encoder),
            temperature=hparams.task.policy.temperature,
            build_respond=False)

        reward_model = rewards.RewardModelTrainer(m, is_root=comm.Get_rank() == 0)

        query_sampler = lm_tasks.make_query_sampler(
            hparams=hparams.task, encoder=encoder, comm=comm,
            batch_size=utils.exact_div(hparams.rollout_batch_size, comm.Get_size())
        )

        tf.train.create_global_step()

        reward_trainer = RewardModelTrainer(
            reward_model=reward_model,
            policy=ref_policy,
            query_sampler=query_sampler,
            hparams=hparams,
            comm=comm,
        )

        save_dir = hparams.run.save_dir
        if comm.Get_rank() == 0 and save_dir:
            print(f"Will save to {save_dir}")
            saver = tf.train.Saver(max_to_keep=20, save_relative_paths=True)
            checkpoint_dir = os.path.join(save_dir, 'reward_model/checkpoints/model.ckpt')

            if not save_dir.startswith('gs://'):
                os.makedirs(os.path.join(save_dir, 'reward_model'), exist_ok=True)
            with tf.gfile.Open(os.path.join(save_dir, 'train_reward_hparams.json'), 'w') as f:
                json.dump(hparams.to_nested_dict(), f, indent=2)
            with tf.gfile.Open(os.path.join(save_dir, 'reward_model', 'hparams.json'), 'w') as f:
                json.dump(reward_model.hparams.to_nested_dict(), f, indent=2)
            with tf.gfile.Open(os.path.join(save_dir, 'reward_model', 'encoding'), 'w') as f:
                json.dump(reward_model.trained_model.encoding.name, f, indent=2)
        else:
            saver = None
            checkpoint_dir = None

        with utils.variables_on_gpu():
            init_ops = tf.group(
                tf.global_variables_initializer(),
                tf.local_variables_initializer(),
                summary.summary_writer_initializer_op())

            @utils.graph_function()
            def sync_models():
                return utils.variable_synchronizer(comm, vars=ref_policy.get_params() + reward_model.get_params())

        tf.get_default_graph().finalize()

        with utils.mpi_session() as sess:
            init_ops.run()
            sync_models()

            reward_trainer.train()

            if saver:
                saver.save(sess, checkpoint_dir)