Adding Onnx Support For Megatron-Bert (#2169)

* Adding Onnx Support For Megatron-Bert

* fix style

* fix style

* fix style

---------

Co-authored-by: pragyan <pragyan.singh@tarkas.ai>
Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>

Author: 3 people

optimum/exporters/onnx/model_configs.py

@@ -171,6 +171,10 @@ class RemBertOnnxConfig(BertOnnxConfig):
     DEFAULT_ONNX_OPSET = 11
 
 
+class MegatronBertOnnxConfig(BertOnnxConfig):
+    DEFAULT_ONNX_OPSET = 11
+
+
 class DistilBertOnnxConfig(BertOnnxConfig):
     DEFAULT_ONNX_OPSET = 14  # now uses F.scaled_dot_product_attention by default for transformers>=4.46.0
 

optimum/exporters/tasks.py

@@ -877,6 +877,15 @@ class TasksManager:
             onnx="MobileBertOnnxConfig",
             tflite="MobileBertTFLiteConfig",
         ),
+        "megatron-bert": supported_tasks_mapping(
+            "feature-extraction",
+            "fill-mask",
+            "text-classification",
+            "multiple-choice",
+            "token-classification",
+            "question-answering",
+            onnx="MegatronBertOnnxConfig",
+        ),
         "mobilevit": supported_tasks_mapping(
             "feature-extraction",
             "image-classification",

tests/exporters/exporters_utils.py

@@ -197,6 +197,7 @@
         "hf-internal-testing/tiny-random-UniSpeechSatForXVector": ["audio-xvector"],
     },
     "audio-spectrogram-transformer": "Ericwang/tiny-random-ast",
+    "megatron-bert": "hf-internal-testing/tiny-random-MegatronBertModel",
     # Disabled for now because some operator seems to not be supported by ONNX.
     # "mctct": "hf-internal-testing/tiny-random-MCTCTModel",
     "speech-to-text": "hf-internal-testing/tiny-random-Speech2TextModel",