Decoupling generation and loss batch sizes

[sidnarayanan]

This introduces a per_device_loss_batch_size to define microbatches to be used when computing the loss. Ideally, I would have liked to compute the loss in chunks of per_device_loss_batch_size and accumulate gradients. However, to compute the advantage, we need all per_device_train_batch_size * num_generations samples.

So instead, we compute the three tensors needed for the loss (reward, logp, KL) in chunks of per_device_loss_batch_size, concatenate the chunks, and compute the full loss all at once. I think this should result in a similar memory reduction, but it remains to be tested.

I also think this code is pretty compilation-unfriendly, since I'm slicing tensors dynamically. Oh well.

[jamesbraza]

LGTM, bonus points for a simple unit test