Is reward_fn equal to log_softmax

[EganGu]

I noticed that the scores in reward_fn is actually equal to logits_i - logsumexp(logits).
I think this expression can be calculated directly by log_softmax. Why not use log_softmax?

LMOps/minillm/minillm/reward.py

Line 33 in 5fbf5bc

def reward_fn(self, input_ids, gen_ids, inf_mask=None, output_pos=True):

[t1101675]

This is because for models with tensor parallelism, log_softmax should be computed as logits_i - logsumexp(logits). To maintain consistency and compare the results, we also use logits_i - logsumexp(logits) in normal scenarios.

[EganGu]

This is because for models with tensor parallelism, log_softmax should be computed as logits_i - logsumexp(logits). To maintain consistency and compare the results, we also use logits_i - logsumexp(logits) in normal scenarios.

Understood. Thanks for your reply.