Restore ability to get logits from generation?

[moyix]

It's often useful when generating samples to be able to get the logits / probabilities of the generated tokens (e.g. for ranking suggestions). It looks like this used to be available but was removed by 7aed133. It would be great if this functionality could be ported to the most recent version of the code.

I attempted to hack it back in, but must have messed something up because it made evaluation ~3.5x slower; I assume I accidentally introduced some GPU-CPU copies but unfortunately I don't really understand the Megatron/Deepspeed code well enough to see how to fix it. That branch is here: https://github.com/moyix/gpt-neox/tree/return_logits

Edit: fixed the link to my attempt at implementation, not sure how that happened

[moyix]

Okay, I think I managed to get it working (at least for my use case). I think the problem previously was that I was trying to save all 52K logits for every token. This version just saves the logit for the token that was actually chosen (actually it saves the normalized log-probability of the token, which is what I really wanted — but it's easy enough to undo that if raw logits are desired):

main...moyix:logit_try2

Happy to make this into a PR if desired!

[StellaAthena]

@moyix Have you looked at this code? It’s not currently exposed by generate.py but looks like it has the desired functionality. Maybe we can clean up the generate.py interface and introduce this as an option?

[moyix]

I hadn't seen that! I will try to take a look and see if I can expose it in generate.py.

[StellaAthena]

@moyix Any updates on this?

[Kyle1668]

I can take this on if no one else is working on this.

[StellaAthena]

@Kyle1668 please do :)

[Kyle1668]

Great! Feel free to assign me the ticket.