Reduce memory consumption in batched_forward_pass

[ohashi56225]

This PR reduces memory consumption in batched_forward_pass of PPOTrainer, by avoiding the storage of logits when they are not necessary.

Before this PR, batched_forward_pass stored all of the model's logits all the time like other tensors such as values and logprobs. Here, logits tensors have a much larger size (batch_size * tokens * vocabulary_size) compared to logprobs and values tensors (batch_size × tokens), consuming a significant amount of cuda memory.

I have modified batched_forward_pass to avoid unnecessary storage of logits, which is only required when calculating entropy in the loss method.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[younesbelkada]

Thanks a lot for fixing and for taking care of the memory consumption
This looks very good to me!
Would love to hear @lvwerra 's thoughts here

[lvwerra]

Looks great, thanks!