Potential bug in accumulating KL

[pradeep-pyro]

Thanks for the paper and releasing this code base!
I had a question on whether this line of code: https://github.com/liziniu/ReMax/blob/master/step3_rlhf_finetuning/remax_trainer.py#L250 should be cumulative_kl += rather than cumulative_kl =. Seems like the KL accumulation is not really happening.

[liziniu]

You are right! Line 250 calculates the one-step KL divergence rather than the cumulative KL. To calculate the cumulative KL, it shoul bd cumulative_kl +=

[pradeep-pyro]

Thank you for the clarification.