Error when running evaluation and inference on QAModel
State: closed | Pull request: no
Labels: name: bug

**Bug Description**

Getting _Exception: cuda rng state model-parallel-rng is not added_ error when trying to evaluate/inference a pretrained QAModel with biomegatron lm after restoring the model from a .nemo file. However I am able to run the evaluate/inference right after training (i.e. trainer.fit).
Code :
`trainer.test(restored_model)`

### Traceback
<details>

``` python
Exception                                 Traceback (most recent call last)
<ipython-input-7-b3320e12a098> in <module>
      1 restored_model.setup_test_data(test_data_config=config.model.test_ds)
----> 2 trainer.test(restored_model)

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in test(self, model, test_dataloaders, ckpt_path, verbose, datamodule)
    753 
    754         if model is not None:
--> 755             results = self.__test_given_model(model, test_dataloaders)
    756         else:
    757             results = self.__test_using_best_weights(ckpt_path, test_dataloaders)

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in __test_given_model(self, model, test_dataloaders)
    818         self.testing = True
    819         self.model = model
--> 820         results = self.fit(model)
    821         self.testing = False
    822 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    471         self.call_hook('on_fit_start')
    472 
--> 473         results = self.accelerator_backend.train()
    474         self.accelerator_backend.teardown()
    475 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py in train(self)
     64 
     65         # train or test
---> 66         results = self.train_or_test()
     67         return results
     68 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py in train_or_test(self)
     65     def train_or_test(self):
     66         if self.trainer.testing:
---> 67             results = self.trainer.run_test()
     68         else:
     69             results = self.trainer.train()

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in run_test(self)
    660         # self.reset_test_dataloader(ref_model)
    661         with self.profiler.profile("run_test_evaluation"):
--> 662             eval_loop_results, _ = self.run_evaluation(test_mode=True)
    663 
    664         if len(eval_loop_results) == 0:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in run_evaluation(self, test_mode, max_batches)
    607                 # lightning module methods
    608                 with self.profiler.profile("evaluation_step_and_end"):
--> 609                     output = self.evaluation_loop.evaluation_step(test_mode, batch, batch_idx, dataloader_idx)
    610                     output = self.evaluation_loop.evaluation_step_end(output)
    611 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py in evaluation_step(self, test_mode, batch, batch_idx, dataloader_idx)
    173         if self.testing:
    174             model_ref._current_fx_name = "test_step"
--> 175             output = self.trainer.accelerator_backend.test_step(args)
    176         else:
    177             model_ref._current_fx_name = "validation_step"

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py in test_step(self, args)
     85 
     86     def test_step(self, args):
---> 87         return self._step(self.trainer.model.test_step, args)
     88 
     89     def to_device(self, batch):

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py in _step(self, model_step, args)
     72         if self.trainer.amp_backend == AMPType.NATIVE:
     73             with torch.cuda.amp.autocast():
---> 74                 output = model_step(*args)
     75         else:
     76             output = model_step(*args)

/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/question_answering/qa_model.py in test_step(self, batch, batch_idx)
    117 
    118     def test_step(self, batch, batch_idx):
--> 119         return self.validation_step(batch, batch_idx)
    120 
    121     def validation_epoch_end(self, outputs):

/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/question_answering/qa_model.py in validation_step(self, batch, batch_idx)
    103 
    104         input_ids, input_type_ids, input_mask, unique_ids, start_positions, end_positions = batch
--> 105         logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
    106         loss, start_logits, end_logits = self.loss(
    107             logits=logits, start_positions=start_positions, end_positions=end_positions

/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py in __call__(self, wrapped, instance, args, kwargs)
    533 
    534         # Call the method - this can be forward, or any other callable method
--> 535         outputs = wrapped(*args, **kwargs)
    536 
    537         instance._attach_and_validate_output_types(output_types=output_types, out_objects=outputs)

/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/question_answering/qa_model.py in forward(self, input_ids, token_type_ids, attention_mask)
     80     @typecheck()
     81     def forward(self, input_ids, token_type_ids, attention_mask):
---> 82         hidden_states = self.bert_model(
     83             input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
     84         )

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    742             result = self._slow_forward(*input, **kwargs)
    743         else:
--> 744             result = self.forward(*input, **kwargs)
    745         for hook in itertools.chain(
    746                 _global_forward_hooks.values(),

/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py in __call__(self, wrapped, instance, args, kwargs)
    533 
    534         # Call the method - this can be forward, or any other callable method
--> 535         outputs = wrapped(*args, **kwargs)
    536 
    537         instance._attach_and_validate_output_types(output_types=output_types, out_objects=outputs)

/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/megatron_bert.py in forward(self, input_ids, attention_mask, token_type_ids)
    119         position_ids = bert_position_ids(input_ids)
    120 
--> 121         sequence_output = self.language_model(
    122             input_ids=input_ids,
    123             position_ids=position_ids,

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    742             result = self._slow_forward(*input, **kwargs)
    743         else:
--> 744             result = self.forward(*input, **kwargs)
    745         for hook in itertools.chain(
    746                 _global_forward_hooks.values(),

/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py in forward(self, input_ids, position_ids, attention_mask, tokentype_ids, layer_past, get_key_value, pooling_sequence_index)
    304 
    305         # Transformer.
--> 306         transformer_output = self.transformer(embedding_output,
    307                                               attention_mask,
    308                                               layer_past=layer_past,

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    742             result = self._slow_forward(*input, **kwargs)
    743         else:
--> 744             result = self.forward(*input, **kwargs)
    745         for hook in itertools.chain(
    746                 _global_forward_hooks.values(),

/opt/conda/lib/python3.8/site-packages/megatron/model/transformer.py in forward(self, hidden_states, attention_mask, layer_past, get_key_value)
    579                 if layer_past is not None:
    580                     past = layer_past[index]
--> 581                 hidden_states = layer(hidden_states,
    582                                       attention_mask,
    583                                       layer_past=past,

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    742             result = self._slow_forward(*input, **kwargs)
    743         else:
--> 744             result = self.forward(*input, **kwargs)
    745         for hook in itertools.chain(
    746                 _global_forward_hooks.values(),

/opt/conda/lib/python3.8/site-packages/megatron/model/transformer.py in forward(self, hidden_states, attention_mask, layer_past, get_key_value)
    417         # Self attention.
    418         attention_output, attention_bias = \
--> 419             self.attention(layernorm_output,
    420                            attention_mask,
    421                            layer_past=layer_past,

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    742             result = self._slow_forward(*input, **kwargs)
    743         else:
--> 744             result = self.forward(*input, **kwargs)
    745         for hook in itertools.chain(
    746                 _global_forward_hooks.values(),

/opt/conda/lib/python3.8/site-packages/megatron/model/transformer.py in forward(self, hidden_states, attention_mask, layer_past, get_key_value)
    294         # This is actually dropping out entire tokens to attend to, which might
    295         # seem a bit unusual, but is taken from the original Transformer paper.
--> 296         with mpu.get_cuda_rng_tracker().fork():
    297             attention_probs = self.attention_dropout(attention_probs)
    298 

/opt/conda/lib/python3.8/contextlib.py in __enter__(self)
    111         del self.args, self.kwds, self.func
    112         try:
--> 113             return next(self.gen)
    114         except StopIteration:
    115             raise RuntimeError("generator didn't yield") from None

/opt/conda/lib/python3.8/site-packages/megatron/mpu/random.py in fork(self, name)
    181         # Check if we have added the state
    182         if name not in self.states_:
--> 183             raise Exception('cuda rng state {} is not added'.format(name))
    184         # Store current rng state.
    185         orig_cuda_rng_state = torch.cuda.get_rng_state()

Exception: cuda rng state model-parallel-rng is not added
```

</details>



**Environment overview**

 - Environment location: GCP
 - Method of NeMo install: python -m pip install git+https://github.com/NVIDIA/NeMo.git@main#egg=nemo_toolkit[all]
 - NVIDIA docker image: nvcr.io/nvidia/pytorch:20.12-py3



@ekmb @ericharper  this is independent from QAModel, since it works for BERT backbone with QAModel. Could you take a look when restoring .nemo model trained with Megatron? I checked it happens both with Megatron and BioMegatron

This should be fixed now, see #1724. Please let us know if you any more issues.

closing, since #1724 is merged. Please re-open if issue persists