Suggest to use larger gradient accumulation steps instead of multi GPUs

[hkunzhe]

In the case of the same batch size, it is recommended to use a larger number of gradient accumulation steps in a single GPU instead of multi-GPUs considering huggingface/diffusers#4046. It may lead to fluctuations in the reward.

[kvablack]

Thanks so much for pointing this out! What a terrible bug. I've been able to fix it so that gradients are synchronized properly across GPUs, but it uses more memory for some reason (up to 16GB from 10GB before the change).

[bhattg]

Hi! Does this bug effect any of the findings in the paper?

[kvablack]

@bhattg no, fortunately the results in the paper all used the original Jax codebase.