Ray CPU OOM with Megatron backend on long sequences and large node count

[jiaqiw09]

When computing log probabilities in the Megatron backend, values are cast to float32, which increases CPU usage in ray.get. Would it be possible to retain the original bfloat16 dtype, or is float32 required to avoid precision loss?

[jiaqiw09]

Here is a draft PR. If keeping bfloat16 has no precision issues, please let me know and I’ll mark it as ready for review.#3651

@vermouth1992