check for first or last stage (#6708) (#6743)

* check for first or last stage



* remove redundant check



* fix typo



* add map_location



---------

Signed-off-by: ericharper <complex451@gmail.com>
Co-authored-by: Eric Harper <complex451@gmail.com>

examples/nlp/language_modeling/megatron_gpt_eval.py

@@ -203,6 +203,7 @@ def main(cfg) -> None:
             trainer=trainer,
             override_config_path=pretrained_cfg,
             save_restore_connector=save_restore_connector,
+            map_location=f'cuda:{trainer.local_rank}',  # map_location is needed for converted models
         )
     elif cfg.checkpoint_dir:
         app_state = AppState()

nemo/collections/nlp/modules/common/text_generation_utils.py

@@ -135,36 +135,41 @@ def megatron_gpt_generate(model, inputs, tokenizer, length_params, sampling_para
 
 
 def get_computeprob_response(tokenizer, response, inputs):
-    compute_prob_response = {}
-    new_token_ids = []
-    new_tokens = []
-    new_texts = []
-    log_probs = []
-    full_logprobs = []
-    offsets = []
-    for batch_id in range(len(response['tokens'])):
-        if isinstance(inputs, (list, tuple)):
-            if isinstance(inputs[0], str):
-                new_token_id = tokenizer.text_to_ids(inputs[batch_id])
-                new_text = inputs[batch_id]
-                token_len = len(new_token_id)
-            elif isinstance(inputs[0], torch.Tensor):
-                token_len = int(inputs[1][batch_id].item())
-                new_token_id = inputs[0][batch_id][:token_len].tolist()
-                new_text = tokenizer.ids_to_text(new_token_id)
-        new_token_ids.append(new_token_id)
-        new_tokens.append(response['tokens'][batch_id][:token_len])
-        new_texts.append(new_text)
-        log_probs.append(response['logprob'][batch_id][:token_len])
-        full_logprobs.append(response['full_logprob'][batch_id][:token_len])
-        offsets.append(response['offsets'][batch_id][:-1])
-    compute_prob_response['sentences'] = new_texts
-    compute_prob_response['tokens'] = new_tokens
-    compute_prob_response['token_ids'] = new_token_ids
-    compute_prob_response['logprob'] = log_probs
-    compute_prob_response['full_logprob'] = full_logprobs
-    compute_prob_response['offsets'] = offsets
-    return compute_prob_response
+    if parallel_state.is_pipeline_first_stage() or parallel_state.is_pipeline_last_stage():
+        # we only have a response on the first and last pipeline stages
+        compute_prob_response = {}
+        new_token_ids = []
+        new_tokens = []
+        new_texts = []
+        log_probs = []
+        full_logprobs = []
+        offsets = []
+        for batch_id in range(len(response['tokens'])):
+            if isinstance(inputs, (list, tuple)):
+                if isinstance(inputs[0], str):
+                    new_token_id = tokenizer.text_to_ids(inputs[batch_id])
+                    new_text = inputs[batch_id]
+                    token_len = len(new_token_id)
+                elif isinstance(inputs[0], torch.Tensor):
+                    token_len = int(inputs[1][batch_id].item())
+                    new_token_id = inputs[0][batch_id][:token_len].tolist()
+                    new_text = tokenizer.ids_to_text(new_token_id)
+            new_token_ids.append(new_token_id)
+            new_tokens.append(response['tokens'][batch_id][:token_len])
+            new_texts.append(new_text)
+            log_probs.append(response['logprob'][batch_id][:token_len])
+            full_logprobs.append(response['full_logprob'][batch_id][:token_len])
+            offsets.append(response['offsets'][batch_id][:-1])
+        compute_prob_response['sentences'] = new_texts
+        compute_prob_response['tokens'] = new_tokens
+        compute_prob_response['token_ids'] = new_token_ids
+        compute_prob_response['logprob'] = log_probs
+        compute_prob_response['full_logprob'] = full_logprobs
+        compute_prob_response['offsets'] = offsets
+        return compute_prob_response
+    else:
+        # intermediate stages
+        return None
 
 
 def get_batch(model, tokenizer, context_tokens):