Add ONNX export support for OLMo and OLMo2 (#2121)

Add support for OLMo and OLMo2 Co-authored-by: Ella Charlaix <[email protected]>
huggingface · Dec 13, 2024 · 3f00766 · 3f00766
1 parent 22d93e7
commit 3f00766
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/docs/source/exporters/onnx/overview.mdx b/docs/source/exporters/onnx/overview.mdx
@@ -75,6 +75,8 @@ Supported architectures from [🤗 Transformers](https://huggingface.co/docs/tra
 - MT5
 - Musicgen (text-conditional only)
 - Nystromformer
+- OLMo
+- OLMo2
 - OWL-ViT
 - Pegasus
 - Perceiver

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -325,6 +325,15 @@ class LlamaOnnxConfig(TextDecoderWithPositionIdsOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
 
 
+class OlmoOnnxConfig(LlamaOnnxConfig):
+    ATOL_FOR_VALIDATION = 1e-4
+    MIN_TRANSFORMERS_VERSION = version.parse("4.40.0")
+
+
+class Olmo2OnnxConfig(OlmoOnnxConfig):
+    MIN_TRANSFORMERS_VERSION = version.parse("4.47.0")
+
+
 class Qwen2OnnxConfig(LlamaOnnxConfig):
     MIN_TRANSFORMERS_VERSION = version.parse("4.37.0")
 

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -954,6 +954,20 @@ class TasksManager:
             "text-generation-with-past",
             onnx="GraniteOnnxConfig",
         ),
+        "olmo": supported_tasks_mapping(
+            "feature-extraction",
+            "feature-extraction-with-past",
+            "text-generation",
+            "text-generation-with-past",
+            onnx="OlmoOnnxConfig",
+        ),
+        "olmo2": supported_tasks_mapping(
+            "feature-extraction",
+            "feature-extraction-with-past",
+            "text-generation",
+            "text-generation-with-past",
+            onnx="Olmo2OnnxConfig",
+        ),
         "pegasus": supported_tasks_mapping(
             "feature-extraction",
             "feature-extraction-with-past",

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -127,6 +127,8 @@
     "mt5": "lewtun/tiny-random-mt5",
     "musicgen": "hf-internal-testing/tiny-random-MusicgenForConditionalGeneration",
     "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
+    "olmo": "hf-internal-testing/tiny-random-OlmoForCausalLM",
+    "olmo2": "hf-internal-testing/tiny-random-Olmo2ForCausalLM",
     "opt": "hf-internal-testing/tiny-random-OPTModel",
     "owlv2": "hf-internal-testing/tiny-random-Owlv2Model",
     "owlvit": "hf-tiny-model-private/tiny-random-OwlViTModel",
-Original file line number
+Diff line change
@@ Expand Up @@
     - MT5
     - Musicgen (text-conditional only)
     - Nystromformer
+    - OLMo
+    - OLMo2
     - OWL-ViT
     - Pegasus
     - Perceiver
@@ Expand Down @@