Converting biomegatron to onnx
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

Hi,
I am trying to convert Biomegatron Model finetuned for token classification tasks using nemo to onnx format. I am getting the following error:

ARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
[NeMo E 2022-08-08 14:55:29 export:164] Export failed. Please make sure your NeMo model class (<class 'nemo.collections.nlp.models.token_classification.token_classification_model.TokenClassificationModel'>) has working export() and that you have the latest NeMo package installed with [all] dependencies.
Traceback (most recent call last):
  File "export.py", line 173, in <module>
    nemo_export(sys.argv[1:])
  File "export.py", line 169, in nemo_export
    raise e
  File "export.py", line 155, in nemo_export
    _, descriptions = model.export(
  File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/exportable.py", line 68, in export
    out, descr, out_example = model._export(
  File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/exportable.py", line 164, in _export
    torch.onnx.export(
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/__init__.py", line 332, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 125, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 824, in _export
    proto, export_map, val_use_external_data_format, node_names = graph._export_onnx(
RuntimeError: ONNX export failed: Couldn't export Python operator ScaledMaskedSoftmax


**Steps/Code to reproduce bug**
I'm using the script export.py used in the below tutorial:
[(https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/MegatronBert_export.ipynb)]

Please, download export.py from here:
[(https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/export.py)]

Also, tried model.export('model.onnx') separately and getting the same error.


**Environment overview (please complete the following information)**

 - Environment location: Docker, GCP
 - docker pull nvcr.io/nvidia/nemo:22.05

**Environment details**

NeMo container:  nvcr.io/nvidia/nemo:22.05
NeMo version: 1.10.0
Branch: 'Main'
Biomegatron Model: "biomegatron345m_biovocab_50k_cased"