in lm_human_preferences/train_policy.py [0:0]
def __init__(self, *, policy, ref_policy, query_sampler, score_fn, hparams, comm):
self.comm = comm
self.policy = policy
self.ref_policy = ref_policy
self.score_fn = score_fn
self.hparams = hparams
if hparams.rewards.adaptive_kl is None:
self.kl_ctl = FixedKLController(hparams.rewards.kl_coef)
else:
self.kl_ctl = AdaptiveKLController(hparams.rewards.kl_coef, hparams=hparams.rewards.adaptive_kl)
response_length = hparams.task.response_length
query_length = hparams.task.query_length
@utils.graph_function()
def sample_queries():
return query_sampler()['tokens']
self.sample_queries = sample_queries
def compute_rewards(scores, logprobs, ref_logprobs):
kl = logprobs - ref_logprobs
non_score_reward = -self.kl_ctl.value * kl
rewards = non_score_reward.copy()
rewards[:, -1] += scores
return rewards, non_score_reward, self.kl_ctl.value
self.compute_rewards = compute_rewards
# per rank sizes
per_rank_rollout_batch_size = utils.exact_div(hparams.ppo.batch_size, comm.Get_size())
per_rank_minibatch_size = utils.exact_div(per_rank_rollout_batch_size, hparams.ppo.nminibatches)
@utils.graph_function(
rollouts=dict(
queries=Schema(tf.int32, (per_rank_minibatch_size, query_length)),
responses=Schema(tf.int32, (per_rank_minibatch_size, response_length)),
values=Schema(tf.float32, (per_rank_minibatch_size, response_length)),
logprobs=Schema(tf.float32, (per_rank_minibatch_size, response_length)),
rewards=Schema(tf.float32, (per_rank_minibatch_size, response_length)),
))
def train_minibatch(rollouts):
"""One step of PPO training."""
left = 1 - policy_frac(hparams)
lrnow = hparams.ppo.lr * left
ppo_loss, stats = self.loss(rollouts)
ppo_train_op = utils.minimize(
loss=ppo_loss, lr=lrnow, params=policy.get_params(), name='ppo_opt', comm=self.comm)
return ppo_train_op, stats
def train(rollouts):
stat_list = []
# Do multiple epochs of PPO training, with a fresh random shuffle in each epoch
for ppo_epoch_idx in range(hparams.ppo.noptepochs):
order = np.random.permutation(per_rank_rollout_batch_size)
for mb_start in range(0, per_rank_rollout_batch_size, per_rank_minibatch_size):
mb_data = {k: v[order[mb_start:mb_start+per_rank_minibatch_size]]
for k, v in rollouts.items()}
step = tf.train.get_global_step().eval()
_, stats = train_minibatch(mb_data)
stat_list.append(stats)
# Collect the stats. (They will be averaged later.)
return {k: [s[k] for s in stat_list] for k in stat_list[0].keys()}
self.train = train
# NOTE: must line up with stats created in self.loss (TODO: better solution?)
scalar_batch = Schema(tf.float32, (None,))
ppo_stat_schemas = utils.flatten_dict(dict(
loss=dict(policy=scalar_batch, value=scalar_batch, total=scalar_batch),
policy=dict(entropy=scalar_batch, approxkl=scalar_batch, clipfrac=scalar_batch),
returns=dict(mean=scalar_batch, var=scalar_batch),
val=dict(vpred=scalar_batch, error=scalar_batch, clipfrac=scalar_batch, mean=scalar_batch, var=scalar_batch),
), sep='/')
stat_data_schemas = dict(
logprobs=Schema(tf.float32, (None, hparams.task.response_length)),
ref_logprobs=Schema(tf.float32, (None, hparams.task.response_length)),
scores=scalar_batch,
non_score_reward=Schema(tf.float32, (None, hparams.task.response_length)),
score_stats=score_fn.stat_schemas,
train_stats=ppo_stat_schemas,
)
@utils.graph_function(
**stat_data_schemas, kl_coef=Schema(tf.float32, ()))
def record_step_stats(*, kl_coef, **data):
ppo_summary_writer = utils.get_summary_writer(self.hparams.run.save_dir, subdir='ppo', comm=self.comm)
kl = data['logprobs'] - data['ref_logprobs']
mean_kl = tf.reduce_mean(tf.reduce_sum(kl, axis=1))
mean_entropy = tf.reduce_mean(tf.reduce_sum(-data['logprobs'], axis=1))
mean_non_score_reward = tf.reduce_mean(tf.reduce_sum(data['non_score_reward'], axis=1))
stats = {
'objective/kl': mean_kl,
'objective/kl_coef': kl_coef,
'objective/entropy': mean_entropy,
}
for k, v in data['train_stats'].items():
stats[f'ppo/{k}'] = tf.reduce_mean(v, axis=0)
for k, v in data['score_stats'].items():
mean = tf.reduce_mean(v, axis=0)
stats[f'objective/{k}'] = mean
stats[f'objective/{k}_total'] = mean + mean_non_score_reward
stats = utils.FlatStats.from_dict(stats).map_flat(
partial(utils.mpi_allreduce_mean, comm=self.comm)).as_dict()
# Add more statistics
step = tf.train.get_global_step().read_value()
stats['ppo/val/var_explained'] = 1 - stats['ppo/val/error'] / stats['ppo/returns/var']
steps = step + 1
stats.update({
'elapsed/updates': steps,
'elapsed/steps/serial': steps * hparams.task.response_length,
'elapsed/steps/total': steps * hparams.ppo.batch_size * hparams.task.response_length,
'elapsed/episodes': steps * hparams.ppo.batch_size,
})
# Time statistics
total, delta = tf_times()
stats.update({
'elapsed/fps': tf.cast(hparams.ppo.batch_size * hparams.task.response_length / delta, tf.int32),
'elapsed/time': total,
})
if ppo_summary_writer:
record_op = utils.record_stats(
stats=stats, summary_writer=ppo_summary_writer, step=step, log_interval=hparams.run.log_interval, name='ppo_stats', comm=self.comm)
else:
record_op = tf.no_op()
return record_op, stats
self.record_step_stats = record_step_stats