Repository: EleutherAI/gpt-neox

Issue:
Restore ability to get logits from generation?
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 7aed13305c48b0048ec8b9efdba01e40fc69f123. 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


Discussion:
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):

https://github.com/EleutherAI/gpt-neox/compare/main...moyix:logit_try2

Happy to make this into a PR if desired!
@moyix Have you looked at [this code](https://github.com/EleutherAI/gpt-neox/blob/8229d921d329266323706c01dd6778fa71649ac7/eval_tasks/eval_adapter.py#L133)? 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?
I hadn't seen that! I will try to take a look and see if I can expose it in `generate.py`.
@moyix Any updates on this?
I can take this on if no one else is working on this. 
@Kyle1668 please do :)
Great! Feel free to assign me the ticket. 

Resolving patch:
diff --git a/megatron/neox_arguments/neox_args.py b/megatron/neox_arguments/neox_args.py
--- a/megatron/neox_arguments/neox_args.py
+++ b/megatron/neox_arguments/neox_args.py
@@ -939,6 +939,11 @@ class NeoXArgsTextgen(NeoXArgsTemplate):
     integer between 0 and the models vocab size. Filters out any logits with a probability less than that of the top_kth token.
     """
 
+    return_logits: bool = False
+    """
+    Boolean for whether to return the logits for generated tokens
+    """
+
     maximum_tokens: int = 64
     """
     maximum number of tokens to be generated
diff --git a/megatron/text_generation_utils.py b/megatron/text_generation_utils.py
--- a/megatron/text_generation_utils.py
+++ b/megatron/text_generation_utils.py
@@ -278,6 +278,7 @@ def stream_tokens(
         # initialize generation variables
         state_is_done = torch.zeros([batch_size]).byte().cuda()
         token_generation_end_index = torch.ones([batch_size]).long().cuda() * (-1)
+        generation_logits = torch.empty(maximum_tokens, neox_args.padded_vocab_size).float().cuda()
 
         while token_index_to_generate <= last_token_index_to_generate:
             if recompute:  # recompute all tokens
@@ -333,6 +334,9 @@ def stream_tokens(
                         next_token_log_probs, num_samples=1
                     ).view(-1)
 
+                if neox_args.return_logits:
+                    generation_logits[token_index_to_generate - 1] = generated_token_logits[0]
+
             if neox_args.is_pipe_parallel:
                 # broadcast generated tokens to pipe parallel group
                 src_rank = model.grid.stage_to_global(model.num_stages - 1)
@@ -378,7 +382,7 @@ def stream_tokens(
 
             token_index_to_generate += 1
 
-            yield context_tokens, token_generation_start_index, token_generation_end_index, state_is_done.bool()
+            yield context_tokens, token_generation_start_index, token_generation_end_index, generation_logits, state_is_done.bool()
             if torch.all(state_is_done):
                 break
 
@@ -473,6 +477,7 @@ def generate_samples_from_prompt(
             batch_context_tokens,
             batch_token_generation_start_index,
             batch_token_generation_end_index,
+            batch_generated_token_logits,
             is_done,
         ) in stream_tokens(
             neox_args=neox_args,
@@ -526,6 +531,10 @@ def generate_samples_from_prompt(
                     "message": message,
                     "duration_seconds": float(time.time() - start_time),
                 }
+
+                if neox_args.return_logits:
+                    data["logits"] = batch_generated_token_logits.cpu().numpy().tolist()
+
                 generated_texts.append(data)
 
     return generated_texts


Test patch: