Add average_log_prob args for cpo

[Mecoli1219]

Summary

trl CPO implementation didn't average the log_probs, while the liger kernel averages it when computing the loss. This will cause a mismatch when integrating them.

Testing Done

• Updating unit test (still investigating why unit test fail locally)

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[kashif]

TRL is using the default as in the official repo for CPO: https://github.com/fe1ixxu/CPO_SIMPO/blob/main/scripts/cpo_trainer.py#L626

[austin362667]

What's the reasoning behind this adjustment?

[Mecoli1219]

@kashif and I find that after disabling average_log_prob for CPO, it will have a higher deviation from HF implementation when the model is large and the data type is bf16. Since the result is still close within both methods, we increase atol and rtol to make this test pass.

[kashif]

as bfloat16 is less accurate for larger numbers, this is needed to make the test pass and is the same as in the other bfloat16 tests

[austin362667]

Then adjusting tol makes sense. ❤️

[austin362667]

Thank you both for making this PR. Hopefully, it unblocks huggingface/trl#2506.

[kashif]

awesome thank you! we would still need a release of liger-kernel for the CI to pass but yes it will hopefully unblock!