Add DeepseekV3 Export (#47)

Author: Abdennacer-Badaoui

docs/source/onnx/overview.mdx

@@ -41,6 +41,7 @@ Supported architectures from [🤗 Transformers](https://huggingface.co/docs/tra
 - Deberta
 - Deberta-v2
 - Decision Transformer
+- DeepSeek-V3
 - Deit
 - Detr
 - DINOv2

optimum/exporters/onnx/model_configs.py

@@ -55,6 +55,7 @@
     ASTDummyAudioInputGenerator,
     BartDummyTextInputGenerator,
     BloomDummyPastKeyValuesGenerator,
+    DeepSeekV3DummyPastKeyValuesGenerator,
     Dinov2DummyInputGenerator,
     DummyCodegenDecoderTextInputGenerator,
     DummyDecisionTransformerInputGenerator,
@@ -441,6 +442,13 @@ class ArceeOnnxConfig(LlamaOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA
 
 
+@register_tasks_manager_onnx("deepseek_v3", *COMMON_TEXT_GENERATION_TASKS)
+class DeepSeekV3OnnxConfig(LlamaOnnxConfig):
+    MIN_TRANSFORMERS_VERSION = version.parse("4.51.0")
+    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DeepSeekV3DummyPastKeyValuesGenerator)
+    DUMMY_PKV_GENERATOR_CLASS = DeepSeekV3DummyPastKeyValuesGenerator
+
+
 @register_tasks_manager_onnx("cohere", *COMMON_TEXT_GENERATION_TASKS)
 class CohereOnnxConfig(LlamaOnnxConfig):
     MIN_TRANSFORMERS_VERSION = version.parse("4.38.0")

optimum/exporters/onnx/utils.py

@@ -71,6 +71,7 @@
 MODEL_TYPES_REQUIRING_POSITION_IDS = {
     "arcee",
     "codegen",
+    "deepseek_v3",
     "cohere",
     "falcon",
     "gemma",

optimum/onnxruntime/modeling_decoder.py

@@ -208,11 +208,16 @@ def __init__(
             self.embed_size_per_head = self.config.head_dim
         elif self.config.model_type == "gpt_bigcode":
             self.embed_size_per_head = self.config.hidden_size // self.config.num_attention_heads * 2
+        elif self.config.model_type == "deepseek_v3":
+            # For deepseek_v3, keys and values have different head dimensions
+            self.qk_head_dim = self.config.qk_rope_head_dim + self.config.qk_nope_head_dim
+            self.v_head_dim = self.config.v_head_dim
         else:
             self.embed_size_per_head = self.config.hidden_size // self.config.num_attention_heads
 
         if self.config.model_type in {
             "arcee",
+            "deepseek_v3",
             "cohere",
             "gemma",
             "helium",
@@ -345,6 +350,10 @@ def forward(
                 v_shape = (batch_size * self.num_key_value_heads, 0, self.embed_size_per_head)
             elif self.config.model_type == "gpt_bigcode" and self.config.multi_query:
                 k_shape = v_shape = (batch_size, 0, self.embed_size_per_head)
+            elif self.config.model_type == "deepseek_v3":
+                # For deepseek_v3, keys and values have different head dimensions
+                k_shape = (batch_size, self.num_key_value_heads, 0, self.qk_head_dim)
+                v_shape = (batch_size, self.num_key_value_heads, 0, self.v_head_dim)
             else:
                 k_shape = v_shape = (batch_size, self.num_key_value_heads, 0, self.embed_size_per_head)
             k_tensor = torch.zeros(k_shape, dtype=self.dtype, device=self.device)
@@ -375,6 +384,10 @@ def forward(
             elif self.config.model_type == "gpt_bigcode" and self.config.multi_query:
                 embed_size_per_head = past_key_values[0].shape[-1]
                 k_shape = v_shape = (batch_size, pkv_seq_len + seq_len, embed_size_per_head)
+            elif self.config.model_type == "deepseek_v3":
+                # For deepseek_v3, keys and values have different head dimensions
+                k_shape = (batch_size, self.num_key_value_heads, pkv_seq_len + seq_len, self.qk_head_dim)
+                v_shape = (batch_size, self.num_key_value_heads, pkv_seq_len + seq_len, self.v_head_dim)
             else:
                 embed_size_per_head = past_key_values[0].shape[-1]
                 k_shape = v_shape = (batch_size, self.num_key_value_heads, pkv_seq_len + seq_len, embed_size_per_head)

pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
-    "optimum @ git+https://github.com/huggingface/optimum",
+    "optimum @ git+https://github.com/huggingface/optimum@add-deepseekv3-dummypastkeyvaluesgenerator",
     "transformers>=4.36,<4.54.0",
     "onnx",
 ]

tests/exporters/onnx/utils_tests.py

@@ -76,6 +76,7 @@
     "deberta-v2": "hf-internal-testing/tiny-random-DebertaV2Model",
     "decision_transformer": "edbeeching/decision-transformer-gym-hopper-medium",
     "deit": "hf-internal-testing/tiny-random-DeiTModel",
+    "deepseek_v3": "hf-internal-testing/tiny-random-DeepseekV3ForCausalLM",
     "dinov2": "hf-internal-testing/tiny-random-Dinov2Model",
     "donut": "fxmarty/tiny-doc-qa-vision-encoder-decoder",
     "donut-swin": "hf-internal-testing/tiny-random-DonutSwinModel",

tests/onnxruntime/test_decoder.py

@@ -31,6 +31,7 @@
     ArceeOnnxConfig,
     BloomOnnxConfig,
     CohereOnnxConfig,
+    DeepSeekV3OnnxConfig,
     GemmaOnnxConfig,
     GraniteOnnxConfig,
     HeliumOnnxConfig,
@@ -122,6 +123,8 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin):
         SUPPORTED_ARCHITECTURES.append("internlm2")
     if is_transformers_version(">=", str(SmolLM3OnnxConfig.MIN_TRANSFORMERS_VERSION)):
         SUPPORTED_ARCHITECTURES.append("smollm3")
+    if is_transformers_version(">=", str(DeepSeekV3OnnxConfig.MIN_TRANSFORMERS_VERSION)):
+        SUPPORTED_ARCHITECTURES.append("deepseek_v3")
     if is_transformers_version(">=", str(StableLMOnnxConfig.MIN_TRANSFORMERS_VERSION)):
         SUPPORTED_ARCHITECTURES.append("stablelm")
 

tests/onnxruntime/testing_utils.py

@@ -50,6 +50,7 @@
     "data2vec-audio": "hf-internal-testing/tiny-random-Data2VecAudioModel",
     "deberta": "hf-internal-testing/tiny-random-DebertaModel",
     "deberta-v2": "hf-internal-testing/tiny-random-DebertaV2Model",
+    "deepseek_v3": "hf-internal-testing/tiny-random-DeepseekV3ForCausalLM",
     "deit": "hf-internal-testing/tiny-random-DeiTModel",
     "donut": "fxmarty/tiny-doc-qa-vision-encoder-decoder",
     "detr": "hf-internal-testing/tiny-random-detr",