Unable to export PyTorch .pt file from biomegatron .nemo model
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

When trying to do a `model.export()` of `qa_squadv2.0_megatron_uncased` nemo model I get `RuntimeError: Could not export Python function call '_ReduceFromModelParallelRegion'` 

**Steps/Code to reproduce bug**

```
wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv2_0_megatron_uncased/versions/1.0.0rc1/zip -O qa_squadv2_0_megatron_uncased_1.0.0rc1.zip

import nemo
import nemo.collections.nlp as nemo_nlp
model = nemo_nlp.models.question_answering.qa_model.QAModel.from_pretrained(model_name="qa_squadv2.0_megatron_uncased")
model.export('/QA.pt')
```

**Expected behavior**

The nemo model should be exported as a PyTorch .pt model.

**Environment overview (please complete the following information)**

 - Environment location: AWS
Running in docker container `nvcr.io/nvidia/pytorch:21.04-py3` with nemo installed via  `python -m pip install git+https://github.com/NVIDIA/NeMo.git@main#egg=nemo_toolkit[all]`

The last section of output is:

```
[NeMo W 2021-05-10 23:12:08 megatron_bert:185] Megatron-lm checkpoint version not found. Setting checkpoint_version to 0.
[NeMo I 2021-05-10 23:12:08 megatron_bert:193] Checkpoint loaded from from /root/.cache/torch/megatron/megatron-bert-345m-uncased
[NeMo I 2021-05-10 23:12:11 modelPT:434] Model QAModel was successfully restored from /root/.cache/torch/NeMo/NeMo_1.0.0/qa_squadv2.0_megatron_uncased/747e4a01f09deb5e4269b1b36cfa4d8e/qa_squadv2.0_megatron_uncased.nemo.
>>> model.export('/foo.pt')
[NeMo W 2021-05-10 23:12:48 export_utils:198] Swapped 49 modules
> initializing torch distributed ...
> initializing tensor model parallel with size 1
> initializing pipeline model parallel with size 1
> setting random seeds to 1234 ...
> initializing model parallel cuda seeds on global rank 0, model parallel rank 0, and data parallel rank 0 with model parallel seed: 3952 and data parallel seed: 1234
[NeMo W 2021-05-10 23:12:49 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/jit/_trace.py:720: UserWarning: `optimize` is deprecated and has no effect. Use `with torch.jit.optimized_execution() instead
      warnings.warn(

[NeMo W 2021-05-10 23:12:53 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/megatron/mpu/utils.py:22: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
      assert numerator % denominator == 0, '{} is not divisible by {}'.format(

[NeMo W 2021-05-10 23:12:53 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/_tensor.py:565: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.
    To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  ../aten/src/ATen/native/BinaryOps.cpp:335.)
      return torch.floor_divide(self, other)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/exportable.py", line 194, in export
    jitted_model.save(output)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py", line 495, in save
    return self._c.save(str(f), **kwargs)
RuntimeError:
Could not export Python function call '_ReduceFromModelParallelRegion'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
/opt/conda/lib/python3.8/site-packages/megatron/mpu/mappings.py(149): reduce_from_tensor_model_parallel_region
/opt/conda/lib/python3.8/site-packages/megatron/mpu/layers.py(191): forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1003): _slow_forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1015): _call_impl
/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py(183): forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1003): _slow_forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1015): _call_impl
/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py(321): forward
/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py(414): forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1003): _slow_forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1015): _call_impl
/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/megatron_bert.py(163): forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1003): _slow_forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1015): _call_impl
/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/question_answering/qa_model.py(81): forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1003): _slow_forward
/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py(1015): _call_impl
/opt/conda/lib/python3.8/site-packages/torch/jit/_trace.py(950): trace_module
/opt/conda/lib/python3.8/site-packages/torch/jit/_trace.py(733): trace
/opt/conda/lib/python3.8/site-packages/nemo/core/classes/exportable.py(184): export
<stdin>(1): <module>

```

What are you trying to accomplish? a) just to extract .pt from .nemo or b) export the model to Torchscript ? 
If a): then simply use 'tar xvf model.nemo'
If b): Megatron model currently is not exportable to Torchscript. You can export it to ONNX though by supplying name with .onnx extension to export()


I was trying to do B - export to torchscript.  I did end up going the ONNX route.