in lm_human_preferences/train_reward.py [0:0]
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