Fix logits/chosen and logits/rejected metrics in kto_trainer.

[PhilipMay]

The calculation of the logits/chosen and logits/rejected metrics in kto_trainer seem to be wrong. A nansum() followed by nanmean() applied to the policy_rejected_logits is wrong.

Our fix is to apply nansum() followed by an other nansum() and then devide the result by count/chosen or count/rejected.

[PhilipMay]

Tagging @MAOJIASONG as the original author and @claralp and asking for a review. 🙏🏼

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Now I understand:

Let's take the following example:

              process 1   process 2
logits        [0, 1, 2]   [3, 4]
sum               3         7
gather              [3, 7]
mean                  5

Your proposed fix

              process 1   process 2
logits        [0, 1, 2]   [3, 4]
sum               3         7
gather              [3, 7]
sum                   10
/all_num_chosen        2

Would the following work?

              process 1   process 2
logits        [0, 1, 2]   [3, 4]
gather          [0, 1, 2, 3, 4]
sum                   10
/all_num_chosen        2

Probably, but too memory intensive, right?

[MAOJIASONG]

Now I understand:

Let's take the following example:

              process 1   process 2
logits        [0, 1, 2]   [3, 4]
sum               3         7
gather              [3, 7]
mean                  5

Your proposed fix

              process 1   process 2
logits        [0, 1, 2]   [3, 4]
sum               3         7
gather              [3, 7]
sum                   10
/all_num_chosen        2

Would the following work?

              process 1   process 2
logits        [0, 1, 2]   [3, 4]
gather          [0, 1, 2, 3, 4]
sum                   10
/all_num_chosen        2

Probably, but too memory intensive, right?

May I ask what's the difference between the first one, and the proposed second one? why is it better to use all_num_chosen for division? @qgallouedec

[qgallouedec]

Because otherwise the result is wrong (5)

[PhilipMay]

Probably, but too memory intensive, right?

@qgallouedec I am not sure about this to be honest.

[MAOJIASONG]

Because otherwise the result is wrong (5)

Thx, I misinterpreted the example. Thanks for pointing it out.