Add gpt-neo-x support (#745)

* add gpt-neo-x configs

* fixes

* name fix

Author: sidthekidder

docs/source/exporters/onnx/package_reference/configuration.mdx

@@ -92,6 +92,7 @@ They specify which input generators should be used for the dummy inputs, but rem
 - GPT-2
 - GPT-J
 - GPT-Neo
+- GPT-NeoX
 - GroupVit
 - Hubert
 - IBert

optimum/exporters/onnx/model_configs.py

@@ -191,6 +191,11 @@ class GPTNeoOnnxConfig(TextDecoderOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads")
 
 
+class GPTNeoXOnnxConfig(TextDecoderOnnxConfig):
+    DEFAULT_ONNX_OPSET = 13
+    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
+
+
 class BloomDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
     def generate(self, input_name: str, framework: str = "pt"):
         past_key_shape = (

optimum/exporters/tasks.py

@@ -383,6 +383,13 @@ class TasksManager:
             "sequence-classification",
             onnx="GPTNeoOnnxConfig",
         ),
+        "gpt-neox": supported_tasks_mapping(
+            "default",
+            "default-with-past",
+            "causal-lm",
+            "causal-lm-with-past",
+            onnx="GPTNeoXOnnxConfig",
+        ),
         "groupvit": supported_tasks_mapping(
             "default",
             onnx="GroupViTOnnxConfig",

optimum/utils/normalized_config.py

@@ -193,6 +193,7 @@ class NormalizedConfigManager:
         "electra": NormalizedTextConfig,
         "gpt2": GPT2LikeNormalizedTextConfig,
         "gpt_neo": NormalizedTextConfig.with_args(num_attention_heads="num_heads"),
+        "gpt_neox": NormalizedTextConfig,
         "gptj": GPT2LikeNormalizedTextConfig,
         "longt5": T5LikeNormalizedTextConfig,
         "marian": BartLikeNormalizedTextConfig,

tests/exporters/exporters_utils.py

@@ -53,6 +53,7 @@
     "flaubert": "hf-internal-testing/tiny-random-flaubert",
     "gpt2": "hf-internal-testing/tiny-random-gpt2",
     "gpt-neo": "hf-internal-testing/tiny-random-GPTNeoModel",
+    "gpt-neox": "hf-internal-testing/tiny-random-GPTNeoXForCausalLM",
     "gptj": "hf-internal-testing/tiny-random-GPTJModel",
     "groupvit": "hf-internal-testing/tiny-random-groupvit",
     "ibert": "hf-internal-testing/tiny-random-IBertModel",
@@ -151,6 +152,7 @@
     "flaubert": "hf-internal-testing/tiny-random-flaubert",  # TODO
     "gpt2": "gpt2",
     "gpt-neo": "EleutherAI/gpt-neo-125M",
+    "gpt-neox": "EleutherAI/gpt-neox-20b",
     "gptj": "anton-l/gpt-j-tiny-random",  # TODO
     "groupvit": "nvidia/groupvit-gcc-yfcc",
     "ibert": "kssteven/ibert-roberta-base",

tests/onnxruntime/test_modeling.py

@@ -111,6 +111,7 @@
     "flaubert": "hf-internal-testing/tiny-random-flaubert",
     "gpt2": "hf-internal-testing/tiny-random-gpt2",
     "gpt_neo": "hf-internal-testing/tiny-random-GPTNeoModel",
+    "gpt_neox": "hf-internal-testing/tiny-random-GPTNeoXForCausalLM",
     "gptj": "hf-internal-testing/tiny-random-GPTJModel",
     "groupvit": "hf-internal-testing/tiny-random-groupvit",
     "ibert": "hf-internal-testing/tiny-random-IBertModel",
@@ -1731,6 +1732,7 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin):
         "codegen",
         "gpt2",
         "gpt_neo",
+        "gpt_neox",
         "gptj",
     ]