Add support for OLMo and OLMo2

huggingface · Dec 7, 2024 · bbc91c0 · bbc91c0
1 parent 4a7cb29
commit bbc91c0
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
@@ -74,6 +74,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
@@ -324,6 +324,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
@@ -949,6 +949,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
@@ -126,6 +126,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 @@