Adding Onnx Support For Megatron-Bert (#2169)

* Adding Onnx Support For Megatron-Bert * fix style * fix style * fix style --------- Co-authored-by: pragyan <[email protected]> Co-authored-by: Ella Charlaix <[email protected]>
huggingface · Jan 28, 2025 · 50531a4 · 50531a4
1 parent 4fdeea7
commit 50531a4
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/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
 

diff --git a/optimum/exporters/tasks.py b/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",

diff --git a/tests/exporters/exporters_utils.py b/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",